91 lines
2.6 KiB
Python
91 lines
2.6 KiB
Python
import os
|
|
import json
|
|
import pymupdf
|
|
from pathlib import Path
|
|
from typing import Dict, Any, List
|
|
|
|
def extract_text_from_pdf(pdf_path: str) -> Dict[str, Any]:
|
|
"""Extract text and metadata from PDF"""
|
|
try:
|
|
doc = pymupdf.open(pdf_path)
|
|
|
|
text = ""
|
|
for page_num in range(len(doc)):
|
|
page = doc[page_num]
|
|
text += page.get_text()
|
|
|
|
metadata = doc.metadata
|
|
doc.close()
|
|
|
|
# Basic cleaning
|
|
text = clean_text(text)
|
|
|
|
return {
|
|
"title": metadata.get('title', Path(pdf_path).stem),
|
|
"text": text,
|
|
"author": metadata.get('author', ''),
|
|
"subject": metadata.get('subject', ''),
|
|
"num_pages": len(doc)
|
|
}
|
|
except Exception as e:
|
|
print(f"Error processing {pdf_path}: {e}")
|
|
return None
|
|
|
|
def process_pdf_directory(input_dir: str, output_dir: str) -> int:
|
|
"""Process all PDFs in a directory"""
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
count = 0
|
|
for pdf_file in Path(input_dir).glob("*.pdf"):
|
|
result = extract_text_from_pdf(str(pdf_file))
|
|
|
|
if result and result.get('text'):
|
|
output_file = os.path.join(output_dir, f"{pdf_file.stem}.json")
|
|
|
|
doc = {
|
|
"id": f"pdf_{pdf_file.stem}",
|
|
"title": result['title'],
|
|
"text": result['text'],
|
|
"source": "pdf",
|
|
"metadata": {
|
|
"author": result.get('author'),
|
|
"subject": result.get('subject'),
|
|
"pages": result.get('num_pages')
|
|
}
|
|
}
|
|
|
|
with open(output_file, 'w', encoding='utf-8') as f:
|
|
json.dump(doc, f, ensure_ascii=False, indent=2)
|
|
|
|
count += 1
|
|
if count % 100 == 0:
|
|
print(f"Processed {count} PDFs")
|
|
|
|
return count
|
|
|
|
def clean_text(text: str) -> str:
|
|
"""Clean extracted PDF text"""
|
|
import re
|
|
|
|
# Remove multiple spaces
|
|
text = re.sub(r' +', ' ', text)
|
|
# Remove multiple newlines
|
|
text = re.sub(r'\n\n+', '\n', text)
|
|
# Remove special characters
|
|
text = re.sub(r'[\x00-\x08\x0b-\x0c\x0e-\x1f]', '', text)
|
|
|
|
return text.strip()
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
if len(sys.argv) < 3:
|
|
print("Usage: python parse_pdf.py <input_dir> <output_dir>")
|
|
sys.exit(1)
|
|
|
|
input_dir = sys.argv[1]
|
|
output_dir = sys.argv[2]
|
|
|
|
count = process_pdf_directory(input_dir, output_dir)
|
|
print(f"Successfully processed {count} PDFs")
|