inference_classifier.py

#!/usr/bin/env python3
"""Standalone binary page-orientation inference (copied to Hub as ``inference_classifier.py``)."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch
import torch.nn as nn
from PIL import Image
from transformers import AutoImageProcessor, AutoModel

DINOV3_MODEL_ID = "facebook/dinov3-vits16-pretrain-lvd1689m"
DEFAULT_LABELS = ("non_flipped", "flipped")
DEFAULT_PREPROCESS = "resize_letterbox"
DEFAULT_PREPROCESS_SIZE = 448


class DINOv3Classifier(nn.Module):
    def __init__(self, model_id: str, num_classes: int, dropout: float = 0.1):
        super().__init__()
        self.backbone = AutoModel.from_pretrained(model_id)
        hidden = self.backbone.config.hidden_size
        self.head = nn.Sequential(
            nn.LayerNorm(hidden),
            nn.Dropout(dropout),
            nn.Linear(hidden, 128),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(128, num_classes),
        )

    def forward(self, pixel_values):
        out = self.backbone(pixel_values=pixel_values)
        cls = out.last_hidden_state[:, 0, :]
        return self.head(cls)


def _resize_short_edge(img: Image.Image, target: int) -> Image.Image:
    w, h = img.size
    if h <= w:
        new_h = target
        new_w = max(1, int(w * target / h))
    else:
        new_w = target
        new_h = max(1, int(h * target / w))
    return img.resize((new_w, new_h), Image.BICUBIC)


def _center_crop(img: Image.Image, size: int = 224) -> Image.Image:
    img = _resize_short_edge(img, size)
    w, h = img.size
    left = max(0, (w - size) // 2)
    top = max(0, (h - size) // 2)
    crop = img.crop((left, top, left + size, top + size))
    if crop.size != (size, size):
        padded = Image.new("RGB", (size, size), (255, 255, 255))
        padded.paste(crop, (0, 0))
        return padded
    return crop


def _letterbox_resize(img: Image.Image, size: int, fill: int = 255) -> Image.Image:
    w, h = img.size
    scale = size / max(w, h)
    nw, nh = round(w * scale), round(h * scale)
    img = img.resize((nw, nh), Image.BILINEAR)
    pad_l = (size - nw) // 2
    pad_t = (size - nh) // 2
    canvas = Image.new("RGB", (size, size), (fill, fill, fill))
    canvas.paste(img, (pad_l, pad_t))
    return canvas


def apply_preprocess(img: Image.Image, mode: str | None, *, size: int = 448) -> Image.Image:
    if not mode or mode == "none":
        return img
    if mode in ("center_crop", "center_crop_whole_page"):
        return _center_crop(img, size)
    if mode == "resize_letterbox":
        return _letterbox_resize(img, size)
    raise ValueError(f"Unknown preprocess mode: {mode!r}")


def processor_skip_resize(mode: str | None) -> bool:
    return mode in ("center_crop", "center_crop_whole_page", "resize_letterbox")


def label_order(ckpt: dict) -> list[str]:
    idx = ckpt.get("idx_to_label") or {}
    if idx:
        return [str(idx[k]) for k in sorted(idx.keys(), key=lambda x: int(x))]
    raw = ckpt.get("label_to_idx") or {}
    if raw:
        return sorted(raw.keys(), key=lambda k: raw[k])
    return list(DEFAULT_LABELS)


def load_model_card_defaults(checkpoint: Path) -> tuple[str, int]:
    card_path = checkpoint.parent / "model_card.json"
    if not card_path.is_file():
        return DEFAULT_PREPROCESS, DEFAULT_PREPROCESS_SIZE
    card = json.loads(card_path.read_text(encoding="utf-8"))
    prep = card.get("preprocess") or {}
    mode = prep.get("test") or prep.get("val") or prep.get("train") or DEFAULT_PREPROCESS
    size = int(prep.get("size") or DEFAULT_PREPROCESS_SIZE)
    return mode, size


def describe_label(name: str) -> str:
    if name == "non_flipped":
        return "upright"
    if name == "flipped":
        return "upside-down (180°)"
    return name


@torch.no_grad()
def predict(
    model,
    processor,
    image_path: Path,
    device,
    *,
    preprocess: str | None,
    size: int,
):
    img = Image.open(image_path).convert("RGB")
    img = apply_preprocess(img, preprocess, size=size)
    pv = processor(
        images=img,
        do_resize=not processor_skip_resize(preprocess),
        return_tensors="pt",
    )["pixel_values"].to(device)
    logits = model(pv)
    probs = torch.softmax(logits, dim=1).squeeze(0).cpu()
    pred = int(probs.argmax())
    return pred, probs.tolist()


def main() -> None:
    ap = argparse.ArgumentParser(
        description="Binary page orientation: non_flipped (upright) vs flipped (180°)."
    )
    ap.add_argument(
        "--checkpoint",
        type=Path,
        default=Path("final_model.pt"),
        help="Weights file (default: final_model.pt in cwd)",
    )
    ap.add_argument("--image", type=Path, nargs="+", required=True)
    ap.add_argument(
        "--preprocess",
        default=None,
        help="none | center_crop | resize_letterbox (default: from model_card.json or resize_letterbox)",
    )
    ap.add_argument(
        "--preprocess-size",
        type=int,
        default=None,
        help="PIL preprocess size before DINO processor (default: from model_card.json or 448)",
    )
    ap.add_argument("--model-id", default=DINOV3_MODEL_ID)
    args = ap.parse_args()

    card_default_mode, card_default_size = load_model_card_defaults(args.checkpoint)
    preprocess = args.preprocess or card_default_mode
    size = args.preprocess_size if args.preprocess_size is not None else card_default_size
    if preprocess in ("none", ""):
        preprocess = None

    ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    classes = label_order(ckpt)
    idx_to_label = {i: lab for i, lab in enumerate(classes)}

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = DINOv3Classifier(args.model_id, num_classes=len(classes)).to(device)
    model.load_state_dict(ckpt["model_state_dict"])
    model.eval()
    processor = AutoImageProcessor.from_pretrained(args.model_id)

    for path in args.image:
        pred, probs = predict(
            model,
            processor,
            path,
            device,
            preprocess=preprocess,
            size=size,
        )
        name = idx_to_label[pred]
        conf = probs[pred]
        print(f"{path.name}: {name} ({describe_label(name)}, {conf:.3f})")
        for i, lab in enumerate(classes):
            print(f"  {lab:14s} {probs[i]:.3f}")


if __name__ == "__main__":
    main()