adapters/base.py

"""
Base Abstract Adapter for Multi-Model OCR Engines.
Standardizes inference execution, timing measurement, and region bounding box propagation.
Zero-fallback: Errors strictly marked as ERROR, without synthetic outputs.
"""

import time
import logging
from abc import ABC, abstractmethod
from typing import Dict, Any, Optional, List, Tuple
from PIL import Image

from core.models import Region

logger = logging.getLogger("BaseOCRAdapter")


class BaseOCRAdapter(ABC):
    """
    Abstract Base Class for all OCR model adapters.
    """

    def __init__(self, model_id: str, display_name: str, default_output_type: str = "markdown"):
        self.model_id = model_id
        self.display_name = display_name
        self.default_output_type = default_output_type
        self._is_loaded = False

    @property
    def is_loaded(self) -> bool:
        return self._is_loaded

    @abstractmethod
    def load_model(self) -> None:
        """Load model weights and processor into memory."""
        pass

    @abstractmethod
    def unload_model(self) -> None:
        """Unload model from memory/VRAM to free resources."""
        self._is_loaded = False

    @abstractmethod
    def run_inference(self, image: Image.Image, **kwargs) -> Dict[str, Any]:
        """
        Execute OCR inference on a PIL Image.
        Returns dictionary containing:
        - 'text': str
        - 'markdown': Optional[str]
        - 'json': Optional[Any]
        - 'regions': Optional[List[Region]]
        - 'output_type': str ('markdown', 'json', 'text')
        """
        pass

    def get_device_and_dtype(self) -> Tuple[str, Any]:
        """
        Determines the optimal device ('cuda' or 'cpu') and torch dtype.
        """
        try:
            import torch
            if torch.cuda.is_available():
                return "cuda", torch.float16
            return "cpu", torch.float32
        except ImportError:
            return "cpu", None

    def process(self, image: Image.Image, **kwargs) -> Dict[str, Any]:
        """
        Public execution wrapper.
        Measures exact runtime using time.perf_counter() around inference.
        Captures any runtime exceptions cleanly without crashing the server.
        """
        t_start = time.perf_counter()

        try:
            # Ensure model is loaded
            if not self._is_loaded:
                logger.info(f"Loading {self.display_name} model...")
                self.load_model()

            # Execute actual model inference
            raw_output = self.run_inference(image, **kwargs)
            elapsed = round(time.perf_counter() - t_start, 3)

            # Ensure regions format
            regions = raw_output.get("regions", [])

            return {
                "model_name": self.display_name,
                "model_id": self.model_id,
                "status": "SUCCESS",
                "inference_time_seconds": elapsed,
                "inference_time_str": f"{elapsed:.2f}s",
                "text": raw_output.get("text", ""),
                "markdown": raw_output.get("markdown", raw_output.get("text", "")),
                "json": raw_output.get("json"),
                "regions": regions,
                "output_type": raw_output.get("output_type", self.default_output_type),
                "error": None
            }

        except Exception as exc:
            elapsed_err = round(time.perf_counter() - t_start, 3)
            logger.error(f"Inference error in {self.display_name}: {exc}", exc_info=True)
            return {
                "model_name": self.display_name,
                "model_id": self.model_id,
                "status": "ERROR",
                "inference_time_seconds": None,
                "inference_time_str": "N/A",
                "text": None,
                "markdown": None,
                "json": None,
                "regions": [],
                "output_type": self.default_output_type,
                "error": f"{type(exc).__name__}: {str(exc)}"
            }