Changes committed

docengine/app/services/ocr_service.py

@@ -0,0 +1,210 @@
+from __future__ import annotations
+
+import uuid
+from pathlib import Path
+
+import cv2
+import numpy as np
+from paddleocr import PaddleOCR
+from sqlalchemy.orm import Session
+
+from app.core.config import settings
+from app.core.logging_config import get_logger
+from app.models.document import Document, DocumentPage
+from app.repositories.document_repository import (
+    DocumentPageRepository,
+    DocumentRepository,
+    DocumentTextBlockRepository,
+)
+from app.storage.provider import get_storage_provider
+
+logger = get_logger(__name__)
+
+_ocr_instance: PaddleOCR | None = None
+
+
+def get_ocr_engine() -> PaddleOCR:
+    """Get or create singleton PaddleOCR instance."""
+    global _ocr_instance
+    if _ocr_instance is None:
+        _ocr_instance = PaddleOCR(
+            use_angle_cls=True,
+            lang=settings.ocr_language,
+            use_gpu=settings.ocr_use_gpu,
+            show_log=False,
+            det_db_thresh=0.3,
+            det_db_box_thresh=0.5,
+            rec_batch_num=6,
+        )
+    return _ocr_instance
+
+
+class OCRService:
+    """OCR processing service using PaddleOCR for scanned documents."""
+
+    def __init__(self, db: Session) -> None:
+        self.db = db
+        self.storage = get_storage_provider()
+        self.doc_repo = DocumentRepository(db)
+        self.page_repo = DocumentPageRepository(db)
+        self.text_block_repo = DocumentTextBlockRepository(db)
+
+    def process_image(self, document: Document) -> Document:
+        """Process a scanned image document with OCR."""
+        file_path = self.storage.get_absolute_path(document.storage_path)
+        image = cv2.imread(file_path)
+        if image is None:
+            raise ValueError(f"Failed to read image: {file_path}")
+
+        height, width = image.shape[:2]
+        document.page_count = 1
+        document.is_scanned = True
+
+        # Save page image
+        image_filename = f"{document.id}_page_1.png"
+        image_bytes = cv2.imencode(".png", image)[1].tobytes()
+        image_path = self.storage.save_file(image_bytes, "images", image_filename)
+
+        # Create page record
+        doc_page = self.page_repo.create_page(
+            document_id=document.id,
+            page_number=1,
+            width=float(width),
+            height=float(height),
+            image_path=image_path,
+        )
+
+        # Run OCR
+        self._run_ocr_on_page(doc_page, file_path)
+
+        return document
+
+    def process_scanned_pdf(self, document: Document) -> Document:
+        """Process a scanned PDF document - convert pages to images and OCR each."""
+        file_path = self.storage.get_absolute_path(document.storage_path)
+
+        try:
+            from pdf2image import convert_from_path
+            images = convert_from_path(file_path, dpi=300)
+        except Exception as e:
+            logger.error("pdf_to_image_failed", document_id=str(document.id), error=str(e))
+            raise
+
+        document.page_count = len(images)
+        document.is_scanned = True
+
+        for page_num, pil_image in enumerate(images, start=1):
+            # Convert PIL to OpenCV format
+            np_image = np.array(pil_image)
+            cv_image = cv2.cvtColor(np_image, cv2.COLOR_RGB2BGR)
+            height, width = cv_image.shape[:2]
+
+            # Save page image
+            image_filename = f"{document.id}_page_{page_num}.png"
+            image_bytes = cv2.imencode(".png", cv_image)[1].tobytes()
+            image_path = self.storage.save_file(image_bytes, "images", image_filename)
+
+            # Create page record
+            doc_page = self.page_repo.create_page(
+                document_id=document.id,
+                page_number=page_num,
+                width=float(width),
+                height=float(height),
+                image_path=image_path,
+            )
+
+            # Run OCR on saved image
+            temp_path = self.storage.get_absolute_path(image_path)
+            self._run_ocr_on_page(doc_page, temp_path)
+
+        return document
+
+    def _run_ocr_on_page(self, doc_page: DocumentPage, image_path: str) -> None:
+        """Run PaddleOCR on a single page image and store results."""
+        ocr = get_ocr_engine()
+
+        try:
+            results = ocr.ocr(image_path, cls=True)
+        except Exception as e:
+            logger.error("ocr_failed", page_id=str(doc_page.id), error=str(e))
+            return
+
+        if not results or not results[0]:
+            logger.info("ocr_no_results", page_id=str(doc_page.id))
+            return
+
+        full_text_parts = []
+        sequence = 0
+
+        for line in results[0]:
+            if not line or len(line) < 2:
+                continue
+
+            bbox_points = line[0]  # List of 4 corner points
+            text_info = line[1]    # (text, confidence)
+
+            text = text_info[0] if isinstance(text_info, (list, tuple)) else str(text_info)
+            confidence = float(text_info[1]) if isinstance(text_info, (list, tuple)) and len(text_info) > 1 else 0.0
+
+            if not text.strip():
+                continue
+
+            # Convert bbox points to x, y, width, height
+            xs = [p[0] for p in bbox_points]
+            ys = [p[1] for p in bbox_points]
+            x = min(xs)
+            y = min(ys)
+            width = max(xs) - x
+            height = max(ys) - y
+
+            # Determine block type based on position
+            block_type = self._classify_text_block(
+                y, height, doc_page.height, text
+            )
+
+            self.text_block_repo.create_text_block(
+                page_id=doc_page.id,
+                text=text,
+                x=x,
+                y=y,
+                width=width,
+                height=height,
+                confidence=confidence,
+                block_type=block_type,
+                sequence=sequence,
+            )
+            full_text_parts.append(text)
+            sequence += 1
+
+        # Update page text content
+        doc_page.text_content = "\n".join(full_text_parts)
+
+    def _classify_text_block(
+        self,
+        y: float,
+        height: float,
+        page_height: float,
+        text: str,
+    ) -> str:
+        """Classify a text block as header, footer, watermark, or regular text."""
+        if page_height <= 0:
+            return "text"
+
+        relative_y = y / page_height
+
+        # Header: top 10%
+        if relative_y < 0.10:
+            return "header"
+
+        # Footer: bottom 10%
+        if relative_y > 0.90:
+            return "footer"
+
+        # Watermark detection heuristic: large text in center
+        if 0.3 < relative_y < 0.7 and height > page_height * 0.05:
+            # Check for common watermark words
+            watermark_keywords = {"confidential", "draft", "copy", "sample", "watermark", "void"}
+            if text.strip().lower() in watermark_keywords:
+                return "watermark"
+
+        return "text"