src/vlps/ocr.py

"""
OCR utilities for Phase 3 — Track C (OCR-focused Question Answering).

Wraps EasyOCR for scene-text extraction, caches results to parquet (extraction is
slow), and provides CER/WER helpers (via jiwer) for OCR-quality evaluation.
"""
import os
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd
from PIL import Image


# A single global reader is reused across calls — model load is expensive.
_READER = None
_READER_LANGS = None


def get_reader(languages=("en",), gpu: Optional[bool] = None):
    """
    Return a cached EasyOCR Reader for the given languages.
    languages: tuple of EasyOCR language codes, e.g. ("en",) or ("pt", "en").
    """
    global _READER, _READER_LANGS
    import easyocr
    import torch

    langs = tuple(languages)
    if _READER is not None and _READER_LANGS == langs:
        return _READER

    if gpu is None:
        gpu = torch.cuda.is_available()
    _READER = easyocr.Reader(list(langs), gpu=gpu)
    _READER_LANGS = langs
    return _READER


def ocr_image(img: Image.Image, reader=None, languages=("en",),
              detail: bool = True) -> dict:
    """
    Run OCR on a single PIL image.
    Returns {'text': joined string, 'tokens': [...], 'confidences': [...]}.
    """
    reader = reader or get_reader(languages)
    arr = np.array(img.convert("RGB"))
    results = reader.readtext(arr, detail=1, paragraph=False)
    tokens = [r[1] for r in results]
    confs = [float(r[2]) for r in results]
    return {
        "text": " ".join(tokens),
        "tokens": tokens,
        "confidences": confs,
    }


def extract_ocr(image_lookup: dict, save_path, languages=("en",),
                gpu: Optional[bool] = None) -> pd.DataFrame:
    """
    Run OCR over an {image_id: PIL.Image} mapping and cache to parquet.
    Reloads from cache if save_path already exists.
    Columns: [image_id, ocr_text, ocr_tokens, ocr_confidences].
    """
    save_path = Path(save_path)
    if save_path.exists():
        print(f"OCR cache found at {save_path}, loading...")
        return pd.read_parquet(save_path)

    reader = get_reader(languages, gpu=gpu)
    rows = []
    n = len(image_lookup)
    for i, (image_id, img) in enumerate(image_lookup.items()):
        try:
            res = ocr_image(img, reader=reader)
        except Exception as e:  # noqa: BLE001 — keep going on a bad image
            print(f"  OCR failed for {image_id}: {e}")
            res = {"text": "", "tokens": [], "confidences": []}
        rows.append({
            "image_id": image_id,
            "ocr_text": res["text"],
            "ocr_tokens": res["tokens"],
            "ocr_confidences": res["confidences"],
        })
        if (i + 1) % 50 == 0:
            print(f"  OCR {i + 1}/{n} images...")

    df = pd.DataFrame(rows)
    save_path.parent.mkdir(parents=True, exist_ok=True)
    df.to_parquet(save_path)
    print(f"Saved OCR for {len(df)} images to {save_path}")
    return df


# ---------------------------------------------------------------------------
# OCR-quality metrics
# ---------------------------------------------------------------------------

def _normalise(text: str) -> str:
    return " ".join(str(text).lower().split())


def cer(reference: str, hypothesis: str) -> float:
    """Character Error Rate (lower is better)."""
    from jiwer import cer as _cer
    ref, hyp = _normalise(reference), _normalise(hypothesis)
    if not ref:
        return 0.0 if not hyp else 1.0
    return float(_cer(ref, hyp))


def wer(reference: str, hypothesis: str) -> float:
    """Word Error Rate (lower is better)."""
    from jiwer import wer as _wer
    ref, hyp = _normalise(reference), _normalise(hypothesis)
    if not ref:
        return 0.0 if not hyp else 1.0
    return float(_wer(ref, hyp))


def evaluate_ocr_quality(df: pd.DataFrame, ref_col: str, hyp_col: str = "ocr_text") -> dict:
    """
    Compute mean CER and WER over a DataFrame where each row has a reference OCR
    string (ref_col, e.g. joined TextVQA ocr_tokens) and our extracted text (hyp_col).
    """
    cers, wers = [], []
    for _, row in df.iterrows():
        ref, hyp = row[ref_col], row[hyp_col]
        cers.append(cer(ref, hyp))
        wers.append(wer(ref, hyp))
    return {"mean_cer": float(np.mean(cers)), "mean_wer": float(np.mean(wers)), "n": len(df)}