64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
|
"""PDF extractor with optional OCR fallback."""
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import List
|
||
|
|
|
||
|
|
from src.ingest.models import DocumentChunk, NormalizedDocument
|
||
|
|
|
||
|
|
|
||
|
|
def _extract_with_pymupdf(path: Path) -> tuple[str, List[DocumentChunk]]:
|
||
|
|
import fitz # pymupdf
|
||
|
|
|
||
|
|
doc = fitz.open(str(path))
|
||
|
|
chunks: List[DocumentChunk] = []
|
||
|
|
parts: List[str] = []
|
||
|
|
for page_num, page in enumerate(doc, start=1):
|
||
|
|
text = page.get_text("text").strip()
|
||
|
|
if text:
|
||
|
|
parts.append(f"--- Страница {page_num} ---\n{text}")
|
||
|
|
chunks.append(DocumentChunk(text=text, source=path.name, page=page_num))
|
||
|
|
doc.close()
|
||
|
|
return "\n\n".join(parts).strip(), chunks
|
||
|
|
|
||
|
|
|
||
|
|
def _extract_with_ocr(path: Path) -> tuple[str, List[DocumentChunk]]:
|
||
|
|
try:
|
||
|
|
import fitz
|
||
|
|
import pytesseract
|
||
|
|
from PIL import Image
|
||
|
|
import io
|
||
|
|
except ImportError:
|
||
|
|
return "", []
|
||
|
|
|
||
|
|
doc = fitz.open(str(path))
|
||
|
|
chunks: List[DocumentChunk] = []
|
||
|
|
parts: List[str] = []
|
||
|
|
for page_num, page in enumerate(doc, start=1):
|
||
|
|
pix = page.get_pixmap(dpi=200)
|
||
|
|
img = Image.open(io.BytesIO(pix.tobytes("png")))
|
||
|
|
text = pytesseract.image_to_string(img, lang="rus+eng").strip()
|
||
|
|
if text:
|
||
|
|
parts.append(f"--- Страница {page_num} (OCR) ---\n{text}")
|
||
|
|
chunks.append(DocumentChunk(text=text, source=path.name, page=page_num))
|
||
|
|
doc.close()
|
||
|
|
return "\n\n".join(parts).strip(), chunks
|
||
|
|
|
||
|
|
|
||
|
|
def extract_pdf(path: Path, document_id: str, project: str, doc_type: str, use_ocr: bool = True) -> NormalizedDocument:
|
||
|
|
full_text, chunks = _extract_with_pymupdf(path)
|
||
|
|
ocr_used = False
|
||
|
|
if not full_text and use_ocr:
|
||
|
|
full_text, chunks = _extract_with_ocr(path)
|
||
|
|
ocr_used = bool(full_text)
|
||
|
|
|
||
|
|
return NormalizedDocument(
|
||
|
|
document_id=document_id,
|
||
|
|
filename=path.name,
|
||
|
|
doc_type=doc_type,
|
||
|
|
project=project,
|
||
|
|
full_text=full_text,
|
||
|
|
chunks=chunks or ([DocumentChunk(text=full_text, source=path.name)] if full_text else []),
|
||
|
|
metadata={"format": "pdf", "ocr_used": ocr_used, "pages": len(chunks)},
|
||
|
|
mime_hint="application/pdf",
|
||
|
|
)
|