timing/backend/ocr.py

"""
Startnummer-deteksjon og OCR.
Bruker EasyOCR for å lese sifre fra bib-nummerlapp.
Returnerer en OcrResult med tall, konfidens og partial-flagg.
"""

import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import cv2
import numpy as np

# EasyOCR lastes lazy for å unngå lang oppstartstid
_reader = None


def _get_reader():
    global _reader
    if _reader is None:
        import easyocr
        _reader = easyocr.Reader(["en"], gpu=False, verbose=False)
    return _reader


@dataclass
class OcrResult:
    digits: Optional[str]       # Gjenkjente sifre, f.eks. "42", None hvis ingen
    confidence: float           # 0.0–1.0
    partial: bool               # True hvis nummeret trolig er delvis skjult
    proximity_score: float = 0.0  # Areal av detektert bib-boks (px²) — større = nærmere kamera
    raw_texts: list[str] = field(default_factory=list)  # Alle OCR-treff for debug


def _preprocess(image_path: Path) -> np.ndarray:
    """Forbered bilde for OCR: gråskala, kontrast, skarphet."""
    img = cv2.imread(str(image_path))
    if img is None:
        raise ValueError(f"Kunne ikke lese bilde: {image_path}")

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # CLAHE for bedre kontrast i varierende lys
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    enhanced = clahe.apply(gray)

    # Skaler opp hvis bildet er lite (OCR er bedre på større tekst)
    h, w = enhanced.shape
    if max(h, w) < 800:
        scale = 800 / max(h, w)
        enhanced = cv2.resize(
            enhanced, (int(w * scale), int(h * scale)),
            interpolation=cv2.INTER_CUBIC,
        )

    return enhanced


def _bbox_area(bbox) -> float:
    """Beregn areal av EasyOCR bounding box [[x1,y1],[x2,y2],[x3,y3],[x4,y4]]."""
    xs = [p[0] for p in bbox]
    ys = [p[1] for p in bbox]
    return (max(xs) - min(xs)) * (max(ys) - min(ys))


def read_all_bibs(image_path: Path) -> list["OcrResult"]:
    """
    Les alle startnumre fra bildet.
    Returnerer én OcrResult per unikt startnummer funnet (minst 2 sifre).
    Tom liste hvis ingen funnet eller ved feil.
    """
    try:
        processed = _preprocess(image_path)
        reader = _get_reader()
        results = reader.readtext(processed, detail=1, paragraph=False)

        raw_texts = [text for (_, text, _) in results]

        # Best konfidens + areal per unikt siffersekvens (minst 2 sifre)
        best: dict[str, tuple[float, float]] = {}  # digits -> (conf, area)
        for (bbox, text, conf) in results:
            digits = re.sub(r"[^0-9]", "", text)
            if len(digits) < 2:
                continue
            area = _bbox_area(bbox)
            if digits not in best or float(conf) > best[digits][0]:
                best[digits] = (float(conf), area)

        return [
            OcrResult(
                digits=digits,
                confidence=conf,
                partial=False,
                proximity_score=area,
                raw_texts=raw_texts,
            )
            for digits, (conf, area) in best.items()
        ]
    except Exception as e:
        return [OcrResult(digits=None, confidence=0.0, partial=False, raw_texts=[f"ERROR: {e}"])]


def read_bib(image_path: Path) -> OcrResult:
    """
    Les beste startnummer fra bildet (bakoverkompatibel).
    Returnerer OcrResult. Aldri exception — fallback til konfidens 0 ved feil.
    """
    bibs = read_all_bibs(image_path)
    if not bibs or bibs[0].digits is None:
        # Sjekk også enkle sifre (partial) for bakoverkompatibilitet
        try:
            processed = _preprocess(image_path)
            reader = _get_reader()
            results = reader.readtext(processed, detail=1, paragraph=False)
            raw_texts = [text for (_, text, _) in results]
            candidates = []
            for (bbox, text, conf) in results:
                digits = re.sub(r"[^0-9]", "", text)
                if digits:
                    candidates.append((digits, float(conf), _bbox_area(bbox)))
            if candidates:
                best_digits, best_conf, best_area = max(candidates, key=lambda x: x[1])
                return OcrResult(
                    digits=best_digits,
                    confidence=best_conf,
                    partial=len(best_digits) < 2,
                    proximity_score=best_area,
                    raw_texts=raw_texts,
                )
        except Exception:
            pass
        return OcrResult(digits=None, confidence=0.0, partial=False)
    return max(bibs, key=lambda r: r.confidence)