ocr/base.py

import io
import json
import os
from PIL import Image
from google import genai
from google.genai import types
from dotenv import load_dotenv

load_dotenv()

# Models tried in order; first successful response wins.
MODELS = [
    "gemini-3-flash-preview",
    "gemini-2.5-flash",
    "gemini-2.0-flash",
]

_client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))

_JSON_CONFIG = types.GenerateContentConfig(
    response_mime_type="application/json",
    thinking_config=types.ThinkingConfig(thinking_budget=0),
)


def smart_resize(image_bytes: bytes, threshold: int = 1400) -> bytes:
    """Downscale images larger than `threshold` px on either side to 1024×1024."""
    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    if img.width > threshold or img.height > threshold:
        img.thumbnail((1024, 1024))
        buf = io.BytesIO()
        img.save(buf, format="JPEG", quality=85)
        return buf.getvalue()
    return image_bytes


class GeminiExtractor:
    """
    Generic Gemini JSON extractor.

    Sends an image + prompt to Gemini and returns parsed JSON.
    Falls back across `models` on failure, raising RuntimeError if all fail.

    Args:
        prompt:  Instruction prompt sent alongside the image.
        models:  Model list to try in order (defaults to module-level MODELS).
        resize:  Pre-shrink large images before sending (default True).
                 Set False when the returned coordinates must align with the
                 original image dimensions (e.g. bounding-box detection).
    """

    def __init__(
        self,
        prompt: str,
        models: list[str] | None = None,
        resize: bool = True,
    ) -> None:
        self._prompt = prompt
        self._models = models or MODELS
        self._resize = resize

    def extract(self, image_bytes: bytes) -> dict:
        """Return ``{"engine_used": str, "data": dict}``."""
        if self._resize:
            image_bytes = smart_resize(image_bytes)

        image_part = types.Part.from_bytes(data=image_bytes, mime_type="image/jpeg")
        last_error = None

        for model in self._models:
            try:
                response = _client.models.generate_content(
                    model=model,
                    contents=[image_part, self._prompt],
                    config=_JSON_CONFIG,
                )
                return {"engine_used": model, "data": json.loads(response.text)}
            except Exception as e:
                last_error = str(e)

        raise RuntimeError(f"All models failed. Last error: {last_error}")