| |
| """Run the TiBLA PP-DocLayout-L Tibetan book layout detector on one or more page |
| images, writing YOLO-format labels. |
| |
| The model is a 4-class PP-DocLayout-L (header, text-area, footnote, footer) |
| fine-tuned on the leak-free v4 `tam2col` split of TiBLAD. It is served from the |
| exported PaddlePaddle inference model in `inference/`. The recommended global |
| operating confidence is ~0.68 (the best-mean-F1 point on the v4 test). |
| |
| Usage: |
| python infer.py --model-dir inference --source page.jpg |
| python infer.py --model-dir inference --source pages/ --out preds --conf 0.68 |
| |
| Requires: paddlepaddle==3.0.0, paddlex. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| IMG_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tif", ".tiff"} |
| |
| NAMES = {"header": 0, "text-area": 1, "footnote": 2, "footer": 3} |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser(description=__doc__, |
| formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--model-dir", default="inference", |
| help="path to the exported inference model dir") |
| 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("--conf", type=float, default=0.68, |
| help="global operating confidence (default 0.68)") |
| args = ap.parse_args() |
|
|
| from PIL import Image |
| from paddlex import create_model |
|
|
| model = create_model(model_name="PP-DocLayout-L", model_dir=args.model_dir) |
|
|
| 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 |
| lines = [] |
| for res in model.predict(str(ip), threshold=args.conf): |
| for box in res["boxes"]: |
| label = box["label"] |
| cls = NAMES.get(label) |
| if cls is None: |
| continue |
| x1, y1, x2, y2 = box["coordinate"] |
| cx, cy = ((x1 + x2) / 2) / W, ((y1 + y2) / 2) / H |
| w, h = (x2 - x1) / W, (y2 - y1) / H |
| lines.append((cls, float(box["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: |
| name = next(k for k, v in NAMES.items() if v == cls) |
| print(f" {name:10} conf={score:.3f} cx={cx:.3f} cy={cy:.3f} " |
| f"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 (conf {args.conf})") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
