predict.py

"""TrOCR inference logic: load both models lazily and run beam-search decoding."""
import os, io, logging
from PIL import Image
from transformers import TrOCRProcessor, VisionEncoderDecoderModel

logger = logging.getLogger(__name__)
NUM_BEAMS = int(os.getenv("NUM_BEAMS", 4))
MAX_LENGTH = int(os.getenv("MAX_LENGTH", 160))

# Fine-tuned BiLinguaScript checkpoints (LoRA merged, so they load with a
# plain from_pretrained). English uses the Cameroonian-adapted model —
# 18.62% CER on held-out Cameroonian writers vs 26.47% for the IAM-only model.
EN_MODEL_ID = os.getenv("EN_MODEL_ID", "unixio/trocr-cameroon-english-lora")
FR_MODEL_ID = os.getenv("FR_MODEL_ID", "unixio/trocr-rimes-french-lora")

# Both checkpoints live in private HF repos — a read token is required.
HF_TOKEN = os.getenv("HF_TOKEN") or None

_models: dict = {}


def normalize_lang(language: str) -> str:
    """Map any accepted spelling ('english', 'EN', 'fr', 'french') to 'en'/'fr'.

    The Django backend sends the full names ('english'/'french'); without this
    normalisation the model lookup below would silently route English requests
    to the French model.
    """
    lang = (language or "").strip().lower()
    if lang in ("fr", "french", "français", "francais"):
        return "fr"
    return "en"


def _load(lang: str):
    lang = normalize_lang(lang)
    if lang not in _models:
        model_id = EN_MODEL_ID if lang == "en" else FR_MODEL_ID
        logger.info("Loading %s model: %s", lang, model_id)
        processor = TrOCRProcessor.from_pretrained(model_id, token=HF_TOKEN)
        model = VisionEncoderDecoderModel.from_pretrained(model_id, token=HF_TOKEN)
        model.eval()
        _models[lang] = (processor, model)
    return _models[lang]


def run_inference(image_bytes: bytes, lang: str) -> dict:
    """Run TrOCR inference on a PNG line image. Returns transcription + confidence."""
    import torch
    processor, model = _load(lang)
    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    pixel_values = processor(images=img, return_tensors="pt").pixel_values
    with torch.no_grad():
        outputs = model.generate(
            pixel_values,
            num_beams=NUM_BEAMS,
            max_length=MAX_LENGTH,
            output_scores=True,
            return_dict_in_generate=True,
        )
    transcription = processor.batch_decode(outputs.sequences, skip_special_tokens=True)[0]
    score = outputs.sequences_scores[0].exp().item() if hasattr(outputs, "sequences_scores") else None
    return {"transcription": transcription, "confidence": score}