src/agents/ocr_agent.py

from __future__ import annotations
from pathlib import Path
import io
import fitz
from PIL import Image
import pytesseract
from src.services.tesseract_service import TesseractService

class OfflineOCRAgent:
    name = "ocr_agent"

    def __init__(self):
        self.tesseract = TesseractService()

    def run(self, input_pdf: Path, output_pdf: Path, language: str = "eng") -> Path:
        self.tesseract.ensure_available()
        src = fitz.open(input_pdf)
        out = fitz.open()
        for page in src:
            pix = page.get_pixmap(dpi=300, alpha=False)
            img = Image.open(io.BytesIO(pix.tobytes("png")))
            pdf_bytes = pytesseract.image_to_pdf_or_hocr(img, extension="pdf", lang=language)
            ocr_doc = fitz.open("pdf", pdf_bytes)
            out.insert_pdf(ocr_doc)
        output_pdf.parent.mkdir(parents=True, exist_ok=True)
        out.save(output_pdf, garbage=4, deflate=True, clean=True)
        return output_pdf