pipeline.py

"""Model load, inference, and overlay rendering for the Gradio Space."""
from __future__ import annotations

import logging
from functools import lru_cache
from pathlib import Path
from typing import Any

import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont
from safetensors.torch import load_file

from regnetx_infer.engine.infer import infer_image
from regnetx_infer.modeling.model import RegNetXPolylineModel
from regnetx_infer.serialization.reading_order import apply_reading_order
from regnetx_infer.serialization.serialize import serialize_pagexml

logger = logging.getLogger("regnetx-space")

DEFAULT_REPO = "johnlockejrr/regnetx-8gf-polyline-baseline-stage0"
DEFAULT_WEIGHTS = "best_cbad_f1.safetensors"

# Manuscript palette — ink / ochre / wax seal (same as sibling Spaces)
COLOR_BASELINE = (36, 72, 92)
COLOR_POLYGON = (176, 92, 48)
COLOR_SCORE = (28, 28, 28)


def resolve_weights(
    *,
    repo_id: str = DEFAULT_REPO,
    filename: str = DEFAULT_WEIGHTS,
    local_path: str | Path | None = None,
) -> Path:
    if local_path:
        p = Path(local_path)
        if not p.is_file():
            raise FileNotFoundError(f"Local weights not found: {p}")
        return p
    from huggingface_hub import hf_hub_download

    return Path(hf_hub_download(repo_id=repo_id, filename=filename))


@lru_cache(maxsize=4)
def load_model(weights_path: str, image_size: int = 1280) -> RegNetXPolylineModel:
    """Load HybridEncoder weights for a given canvas size.

    Checkpoint buffers decoder.anchors / decoder.valid_mask are tied to the
    train-time canvas. Skip them so freshly generated buffers for image_size
    are kept when they differ.

    pretrained_backbone=False — full weights come from the safetensors file;
    do not download ImageNet timm weights on Space startup.
    """
    logger.info("Loading RegNetXPolylineModel from %s (canvas=%d)", weights_path, image_size)
    model = RegNetXPolylineModel(
        backbone_name="regnetx_080.tv2_in1k",
        pretrained_backbone=False,
        num_queries=300,
        num_ctrl=8,
        eval_spatial_size=(image_size, image_size),
    )
    state = load_file(weights_path)
    cleaned: dict[str, torch.Tensor] = {}
    skipped_spatial = 0
    for k, v in state.items():
        nk = k.replace("model.", "").replace("ema_model.", "")
        if nk in {"decoder.anchors", "decoder.valid_mask"}:
            skipped_spatial += 1
            continue
        cleaned[nk] = v
    missing, unexpected = model.load_state_dict(cleaned, strict=False)
    logger.info(
        "Loaded weights (missing=%d unexpected=%d skipped_spatial=%d)",
        len(missing),
        len(unexpected),
        skipped_spatial,
    )
    model.eval()
    return model


def run_detection(
    image: Image.Image,
    model: RegNetXPolylineModel,
    *,
    conf: float = 0.4,
    image_size: int = 1280,
    reading_order: str = "rtl",
    polygonize: bool = True,
    device: str = "cpu",
) -> list[dict[str, Any]]:
    lines = infer_image(
        model,
        image,
        conf_thresh=float(conf),
        image_size=int(image_size),
        letterbox=True,
        device=device,
        polygonize=bool(polygonize),
        fix_intersections=True,
    )
    return apply_reading_order(lines, reading_order=reading_order.lower())


def _font(size: int = 14) -> ImageFont.ImageFont:
    for name in (
        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
        "/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf",
    ):
        try:
            return ImageFont.truetype(name, size=size)
        except OSError:
            continue
    return ImageFont.load_default()


def render_overlay(
    image: Image.Image,
    lines: list[dict[str, Any]],
    *,
    show_baselines: bool = True,
    show_polygons: bool = True,
    show_scores: bool = False,
    polygon_opacity: float = 0.28,
    baseline_width: int = 3,
) -> Image.Image:
    """Composite polygon fills + baseline strokes onto a copy of the page."""
    base = image.convert("RGBA")
    poly_layer = Image.new("RGBA", base.size, (0, 0, 0, 0))
    line_layer = Image.new("RGBA", base.size, (0, 0, 0, 0))
    draw_poly = ImageDraw.Draw(poly_layer, "RGBA")
    draw_line = ImageDraw.Draw(line_layer, "RGBA")
    font = _font(13)
    alpha = int(max(0.0, min(1.0, polygon_opacity)) * 255)
    fill = (*COLOR_POLYGON, alpha)
    outline = (*COLOR_POLYGON, min(255, alpha + 90))
    stroke = (*COLOR_BASELINE, 230)

    for item in lines:
        boundary = item.get("boundary") or []
        baseline = item.get("baseline") or []
        score = float(item.get("score", 0.0))

        if show_polygons and len(boundary) >= 3:
            pts = [(int(round(x)), int(round(y))) for x, y in boundary]
            draw_poly.polygon(pts, fill=fill, outline=outline)

        if show_baselines and len(baseline) >= 2:
            pts = [(int(round(x)), int(round(y))) for x, y in baseline]
            draw_line.line(pts, fill=stroke, width=max(1, baseline_width))

        if show_scores and baseline:
            x0, y0 = baseline[0]
            label = f"{score:.2f}"
            tx, ty = int(round(x0)) + 2, int(round(y0)) - 16
            draw_line.text((tx, ty), label, fill=(*COLOR_SCORE, 240), font=font)

    out = Image.alpha_composite(base, poly_layer)
    out = Image.alpha_composite(out, line_layer)
    return out.convert("RGB")


def lines_to_pagexml(
    image: Image.Image, lines: list[dict[str, Any]], filename: str = "page.jpg"
) -> str:
    return serialize_pagexml(
        image_filename=filename,
        image_size=image.size,
        lines=lines,
    )


def summarize_lines(lines: list[dict[str, Any]]) -> str:
    if not lines:
        return "No lines detected — try a lower confidence threshold."
    scores = [float(x.get("score", 0.0)) for x in lines]
    return (
        f"**{len(lines)}** lines · score min `{min(scores):.2f}` · "
        f"mean `{float(np.mean(scores)):.2f}` · max `{max(scores):.2f}`"
    )