| |
| """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.26 is the best-mean-F1 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"} |
|
|
| |
| |
| 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) |
| |
| 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 |
| 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()) |
|
|
