File size: 963 Bytes
734b5b4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | 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
|