BytecodeApps/docverse-api

1import io2import fitz # PyMuPDF for PDF handling3import easyocr4import numpy as np5from PIL import Image6 7class OCREngine:8    def __init__(self, languages=['en']):9        # Initialize EasyOCR reader (this will download models on first run)10        # In production, use gpu=True if running on a CUDA-enabled machine11        self.reader = easyocr.Reader(languages, gpu=False)12 13    def process_document(self, file_bytes: bytes, filename: str) -> str:14        """15        Processes a PDF or Image file and returns extracted text.16        """17        if filename.lower().endswith('.pdf'):18            return self._extract_from_pdf(file_bytes)19        elif filename.lower().endswith(('.png', '.jpg', '.jpeg')):20            return self._extract_from_image(file_bytes)21        else:22            raise ValueError("Unsupported file format for OCR.")23 24    def _extract_from_image(self, image_bytes: bytes) -> str:25        """ Extract text from a single image byte array """26        try:27            # EasyOCR can read bytes directly, but converting to PIL then numpy array is safe28            image = Image.open(io.BytesIO(image_bytes))29            # Convert to RGB (EasyOCR expects numpy arrays)30            img_np = np.array(image.convert('RGB'))31            32            # Extract text33            results = self.reader.readtext(img_np, detail=0) # detail=0 returns just text34            return "\n".join(results)35        except Exception as e:36            # Log error in production37            print(f"OCR Image Error: {e}")38            return ""39 40    def _extract_from_pdf(self, pdf_bytes: bytes) -> str:41        """ Extract text from a multi-page PDF document by rendering pages as images """42        try:43            doc = fitz.open(stream=pdf_bytes, filetype="pdf")44            extracted_text = []45 46            for page_num in range(len(doc)):47                page = doc.load_page(page_num)48                # Render page to an image (pixmap)49                pix = page.get_pixmap(matrix=fitz.Matrix(2, 2)) # 2x zoom for better OCR50                img_bytes = pix.tobytes("png")51                52                # Run OCR on the rendered page53                page_text = self._extract_from_image(img_bytes)54                extracted_text.append(page_text)55            56            return "\n\n".join(extracted_text)57        except Exception as e:58            print(f"OCR PDF Error: {e}")59            return ""60 61# -- Usage Example --62# engine = OCREngine()63# with open("sample_rx.jpg", "rb") as f:64#     text = engine.process_document(f.read(), "sample_rx.jpg")65