95 lines
3.1 KiB
Python
95 lines
3.1 KiB
Python
import pytesseract
|
|
from pdf2image import convert_from_path
|
|
from PIL import Image
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
import tempfile
|
|
|
|
from config import settings
|
|
from logger import get_logger
|
|
|
|
log = get_logger("ocr")
|
|
|
|
|
|
class OCRExtractor:
|
|
def __init__(self, lang: str = None):
|
|
self.lang = lang or settings.TESSERACT_LANG
|
|
log.debug(f"OCRExtractor initialized: lang={self.lang}")
|
|
|
|
def extract_from_image(self, image_path: Path | str) -> str:
|
|
log.debug(f"Extracting from image: {image_path}")
|
|
image = Image.open(image_path)
|
|
log.debug(f"Image size: {image.size}")
|
|
|
|
text = pytesseract.image_to_string(image, lang=self.lang)
|
|
log.info(f"Image OCR: {len(text)} chars extracted")
|
|
|
|
return text.strip()
|
|
|
|
def extract_from_pdf(self, pdf_path: Path | str) -> list[dict]:
|
|
log.debug(f"Converting PDF to images: {pdf_path}")
|
|
|
|
pages = convert_from_path(pdf_path)
|
|
log.debug(f"PDF has {len(pages)} pages")
|
|
|
|
results = []
|
|
|
|
for i, page in enumerate(pages, 1):
|
|
log.debug(f"Processing page {i}/{len(pages)}")
|
|
text = pytesseract.image_to_string(page, lang=self.lang)
|
|
if text.strip():
|
|
results.append({
|
|
"page": i,
|
|
"text": text.strip()
|
|
})
|
|
log.debug(f"Page {i}: {len(text)} chars")
|
|
else:
|
|
log.debug(f"Page {i}: empty (skipped)")
|
|
|
|
log.info(f"PDF OCR complete: {len(results)} pages with text")
|
|
return results
|
|
|
|
def extract(self, file_path: Path | str) -> dict:
|
|
path = Path(file_path)
|
|
log.info(f"Extracting: {path.name}")
|
|
|
|
if not path.exists():
|
|
log.error(f"File not found: {path}")
|
|
raise FileNotFoundError(f"File not found: {path}")
|
|
|
|
suffix = path.suffix.lower()
|
|
log.debug(f"File type: {suffix}")
|
|
|
|
if suffix == ".pdf":
|
|
pages = self.extract_from_pdf(path)
|
|
full_text = "\n\n".join(p["text"] for p in pages)
|
|
result = {
|
|
"source": str(path),
|
|
"type": "pdf",
|
|
"pages": pages,
|
|
"full_text": full_text,
|
|
"total_pages": len(pages)
|
|
}
|
|
log.info(f"PDF extracted: {len(full_text)} chars, {len(pages)} pages")
|
|
return result
|
|
|
|
elif suffix in {".png", ".jpg", ".jpeg", ".tiff", ".bmp"}:
|
|
text = self.extract_from_image(path)
|
|
result = {
|
|
"source": str(path),
|
|
"type": "image",
|
|
"full_text": text,
|
|
"total_pages": 1
|
|
}
|
|
log.info(f"Image extracted: {len(text)} chars")
|
|
return result
|
|
|
|
else:
|
|
log.error(f"Unsupported file type: {suffix}")
|
|
raise ValueError(f"Unsupported file type: {suffix}")
|
|
|
|
|
|
def extract_text(file_path: Path | str, lang: str = None) -> dict:
|
|
extractor = OCRExtractor(lang=lang)
|
|
return extractor.extract(file_path)
|