Upload infer.py with huggingface_hub

infer.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+"""Run the Tibetan modern-book layout detector (RF-DETR-L fine-tune) on one or
+more page images, applying the recommended *per-class* confidence thresholds.
+
+The model is a 4-class RF-DETR-L (header, text-area, footnote, footer),
+fine-tuned on the same `tam2col` labels as BDRC's primary RT-DETR-l release
+(see the model card / blog post). Like that model, the detector is
+deliberately recall-happy on the small marginal header/footer boxes, so the
+single best operating point differs by class. The thresholds below are each
+class's own max-F1 confidence from a native per-class sweep on the held-out
+test set (same methodology as the primary RT-DETR-l release); a single global
+0.30 is the best compromise if you need one number for all classes.
+
+Usage:
+    python infer.py --checkpoint rfdetr_tibetan_book_layout.pth --source page.jpg
+    python infer.py --checkpoint rfdetr_tibetan_book_layout.pth --source pages/ --out preds
+"""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+IMG_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tif", ".tiff"}
+
+# Per-class max-F1 operating points (see model card). Use --global-conf 0.30
+# instead if you prefer one number for all classes.
+CLASS_THRESHOLDS = {0: 0.46, 1: 0.32, 2: 0.26, 3: 0.52}
+CONF_FLOOR = min(CLASS_THRESHOLDS.values())
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__,
+                                 formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--checkpoint", required=True, help="path to the .pth checkpoint")
+    ap.add_argument("--source", required=True, help="image file or folder")
+    ap.add_argument("--out", default=None,
+                    help="optional folder to write YOLO-format .txt labels")
+    ap.add_argument("--shape", type=int, default=1024)
+    ap.add_argument("--global-conf", type=float, default=None,
+                    help="use ONE threshold for all classes instead of per-class")
+    args = ap.parse_args()
+
+    from rfdetr import RFDETRLarge
+    from PIL import Image
+
+    thresholds = ({c: args.global_conf for c in CLASS_THRESHOLDS}
+                  if args.global_conf is not None else CLASS_THRESHOLDS)
+    floor = min(thresholds.values())
+
+    model = RFDETRLarge.from_checkpoint(args.checkpoint)
+    # class_names on the checkpoint: ['none', 'header', 'text-area', 'footnote', 'footer']
+    names = {0: "header", 1: "text-area", 2: "footnote", 3: "footer"}
+
+    src = Path(args.source)
+    imgs = sorted(p for p in src.iterdir() if p.suffix.lower() in IMG_EXTS) \
+        if src.is_dir() else [src]
+
+    out_dir = Path(args.out) if args.out else None
+    if out_dir:
+        out_dir.mkdir(parents=True, exist_ok=True)
+
+    n_img = n_kept = 0
+    for ip in imgs:
+        n_img += 1
+        with Image.open(ip) as im:
+            W, H = im.size
+        det = model.predict(str(ip), threshold=floor, shape=(args.shape, args.shape))
+        lines = []
+        if det is not None and len(det) > 0:
+            for box, cls_id, score in zip(det.xyxy, det.class_id, det.confidence):
+                cls = int(cls_id) - 1  # class 0 on the checkpoint is background
+                if cls < 0 or cls > 3 or score < thresholds.get(cls, floor):
+                    continue
+                x1, y1, x2, y2 = box.tolist()
+                cx, cy = ((x1 + x2) / 2) / W, ((y1 + y2) / 2) / H
+                w, h = (x2 - x1) / W, (y2 - y1) / H
+                lines.append((cls, float(score), cx, cy, w, h))
+        n_kept += len(lines)
+        print(f"{ip.stem}: {len(lines)} boxes")
+        for cls, score, cx, cy, w, h in lines:
+            print(f"    {names[cls]:10} conf={score:.3f}  "
+                  f"cx={cx:.3f} cy={cy:.3f} w={w:.3f} h={h:.3f}")
+        if out_dir:
+            (out_dir / f"{ip.stem}.txt").write_text(
+                "".join(f"{c} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}\n"
+                        for c, _, cx, cy, w, h in lines))
+
+    print(f"\n{n_img} images, {n_kept} boxes kept (thresholds: {thresholds})")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())