Update inference_classifier.py: letterbox inference defaults and docs

inference_classifier.py

@@ -1,30 +1,139 @@
 #!/usr/bin/env python3
-"""Run DINOv3 script classifier on image paths (classes from checkpoint ``idx_to_label``)."""
+"""Standalone binary page-orientation inference (copied to Hub as ``inference_classifier.py``)."""
 
 from __future__ import annotations
 
 import argparse
-import sys
+import json
 from pathlib import Path
 
 import torch
+import torch.nn as nn
 from PIL import Image
-from transformers import AutoImageProcessor
+from transformers import AutoImageProcessor, AutoModel
 
-ROOT = Path(__file__).resolve().parents[1]
-if str(ROOT) not in sys.path:
-    sys.path.insert(0, str(ROOT))
+DINOV3_MODEL_ID = "facebook/dinov3-vits16-pretrain-lvd1689m"
+DEFAULT_LABELS = ("non_flipped", "flipped")
+DEFAULT_PREPROCESS = "resize_letterbox"
+DEFAULT_PREPROCESS_SIZE = 448
 
-from src.labels import DINOV3_MODEL_ID, MULTICLASS_6_LABELS
-from src.models import DINOv3Classifier
-from src.transforms import apply_preprocess, normalize_preprocess_mode
+
+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):
+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, return_tensors="pt")["pixel_values"].to(device)
+    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())
@@ -32,33 +141,60 @@ def predict(model, processor, image_path: Path, device, *, preprocess: str | Non
 
 
 def main() -> None:
-    ap = argparse.ArgumentParser(description="DINOv3 script-family inference")
-    ap.add_argument("--checkpoint", type=Path, required=True)
+    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="center_crop", help="none | center_crop")
-    ap.add_argument("--preprocess-size", type=int, default=224)
+    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)
-    idx_to_label = {int(k): v for k, v in ckpt.get("idx_to_label", {}).items()}
-    if not idx_to_label:
-        idx_to_label = {i: lab for i, lab in enumerate(MULTICLASS_6_LABELS)}
+    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(idx_to_label)).to(device)
+    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)
-    prep = normalize_preprocess_mode(args.preprocess)
 
     for path in args.image:
         pred, probs = predict(
-            model, processor, path, device, preprocess=prep, size=args.preprocess_size
+            model,
+            processor,
+            path,
+            device,
+            preprocess=preprocess,
+            size=size,
         )
         name = idx_to_label[pred]
         conf = probs[pred]
-        print(f"{path.name}: {name} ({conf:.3f})")
+        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__":