src/document/ocr/base.py

"""
Base OCR Interface

Defines the abstract OCR engine interface and common data structures.
"""

from abc import ABC, abstractmethod
from typing import List, Optional, Dict, Any, Tuple
from dataclasses import dataclass, field
from enum import Enum
import numpy as np
from pydantic import BaseModel, Field

from ..schemas.core import BoundingBox, OCRRegion


class OCRLanguage(str, Enum):
    """Supported OCR languages."""
    ENGLISH = "en"
    CHINESE_SIMPLIFIED = "ch"
    CHINESE_TRADITIONAL = "chinese_cht"
    FRENCH = "fr"
    GERMAN = "german"
    SPANISH = "es"
    ITALIAN = "it"
    PORTUGUESE = "pt"
    RUSSIAN = "ru"
    JAPANESE = "japan"
    KOREAN = "korean"
    ARABIC = "ar"
    HINDI = "hi"
    LATIN = "latin"


class OCRConfig(BaseModel):
    """Configuration for OCR processing."""
    # Engine selection
    engine: str = Field(default="paddle", description="OCR engine: paddle or tesseract")

    # Language settings
    languages: List[str] = Field(
        default=["en"],
        description="Languages to detect (ISO codes)"
    )

    # Detection settings
    det_db_thresh: float = Field(
        default=0.3,
        ge=0.0,
        le=1.0,
        description="Detection threshold for text regions"
    )
    det_db_box_thresh: float = Field(
        default=0.5,
        ge=0.0,
        le=1.0,
        description="Box detection threshold"
    )

    # Recognition settings
    rec_batch_num: int = Field(
        default=6,
        ge=1,
        description="Recognition batch size"
    )
    min_confidence: float = Field(
        default=0.5,
        ge=0.0,
        le=1.0,
        description="Minimum confidence threshold"
    )

    # Performance settings
    use_gpu: bool = Field(default=True, description="Use GPU acceleration")
    gpu_id: int = Field(default=0, ge=0, description="GPU device ID")
    use_angle_cls: bool = Field(
        default=True,
        description="Use angle classification for rotated text"
    )
    use_dilation: bool = Field(
        default=False,
        description="Use dilation for detection"
    )

    # Output settings
    drop_score: float = Field(
        default=0.5,
        ge=0.0,
        le=1.0,
        description="Drop results below this score"
    )
    return_word_boxes: bool = Field(
        default=False,
        description="Return word-level boxes (vs line-level)"
    )

    # Preprocessing
    preprocess_resize: Optional[int] = Field(
        default=None,
        description="Resize image max dimension before OCR"
    )
    preprocess_denoise: bool = Field(
        default=False,
        description="Apply denoising before OCR"
    )


@dataclass
class OCRResult:
    """
    Result of OCR processing for a single image/page.
    """
    regions: List[OCRRegion] = field(default_factory=list)
    full_text: str = ""
    confidence_avg: float = 0.0
    processing_time_ms: float = 0.0
    engine: str = "unknown"
    language_detected: Optional[str] = None

    # Error handling
    success: bool = True
    error: Optional[str] = None

    def get_text_in_bbox(self, bbox: BoundingBox) -> str:
        """Get text within a bounding box."""
        texts = []
        for region in self.regions:
            if bbox.contains(region.bbox) or bbox.iou(region.bbox) > 0.5:
                texts.append(region.text)
        return " ".join(texts)

    def filter_by_confidence(self, min_confidence: float) -> "OCRResult":
        """Return new result with regions above confidence threshold."""
        filtered_regions = [r for r in self.regions if r.confidence >= min_confidence]
        return OCRResult(
            regions=filtered_regions,
            full_text=" ".join(r.text for r in filtered_regions),
            confidence_avg=sum(r.confidence for r in filtered_regions) / len(filtered_regions) if filtered_regions else 0,
            processing_time_ms=self.processing_time_ms,
            engine=self.engine,
            language_detected=self.language_detected,
            success=self.success,
            error=self.error,
        )


class OCREngine(ABC):
    """
    Abstract base class for OCR engines.
    Defines the interface that all OCR implementations must follow.
    """

    def __init__(self, config: Optional[OCRConfig] = None):
        """
        Initialize OCR engine.

        Args:
            config: OCR configuration
        """
        self.config = config or OCRConfig()
        self._initialized = False

    @abstractmethod
    def initialize(self):
        """Initialize the OCR engine (load models, etc.)."""
        pass

    @abstractmethod
    def recognize(
        self,
        image: np.ndarray,
        page_number: int = 0,
    ) -> OCRResult:
        """
        Perform OCR on an image.

        Args:
            image: Image as numpy array (RGB, HWC format)
            page_number: Page number for multi-page documents

        Returns:
            OCRResult with recognized text and regions
        """
        pass

    def recognize_batch(
        self,
        images: List[np.ndarray],
        page_numbers: Optional[List[int]] = None,
    ) -> List[OCRResult]:
        """
        Perform OCR on multiple images.

        Args:
            images: List of images
            page_numbers: Optional page numbers

        Returns:
            List of OCRResult
        """
        if page_numbers is None:
            page_numbers = list(range(len(images)))

        results = []
        for img, page_num in zip(images, page_numbers):
            results.append(self.recognize(img, page_num))
        return results

    @abstractmethod
    def get_supported_languages(self) -> List[str]:
        """Return list of supported language codes."""
        pass

    @property
    def name(self) -> str:
        """Return engine name."""
        return self.__class__.__name__

    @property
    def is_initialized(self) -> bool:
        """Check if engine is initialized."""
        return self._initialized

    def __enter__(self):
        """Context manager entry."""
        if not self._initialized:
            self.initialize()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        pass