211 lines
6.7 KiB
Python
211 lines
6.7 KiB
Python
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"
|