Sync predict.py for 263702

predict.py

@@ -24,8 +24,9 @@ import numpy as np
 import torch
 from PIL import Image
 
+from jointdecode import joint_decode
 from model import ClockNetCls, decode_cls
-from twostage import DialLocator, PretrainedReader, crop_dial
+from twostage import CocoDialDetector, DialLocator, PretrainedReader, crop_dial
 
 MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
 STD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
@@ -40,9 +41,13 @@ def build_reader(ckpt, device):
     blob = torch.load(ckpt, map_location="cpu", weights_only=False)
     a = blob["args"]
     backbone = a.get("backbone", "scratch")
+    # checkpoints written before the whole-time head have no weights for it,
+    # so take the head's size from the file rather than from today's default
+    tb = blob["model"].get("time_head.weight")
     model = (ClockNetCls(width=a.get("width", 32), bins=a.get("bins", 180))
              if backbone == "scratch" else
-             PretrainedReader(bins=a.get("bins", 180), arch=backbone))
+             PretrainedReader(bins=a.get("bins", 180), arch=backbone,
+                              time_bins=0 if tb is None else tb.shape[0]))
     model.load_state_dict(blob["model"])
     return model.to(device).eval(), a.get("res", 256)
 
@@ -53,16 +58,34 @@ def main():
     ap.add_argument("--labels", required=True)
     ap.add_argument("--root", required=True, help="directory the 'file' fields are relative to")
     ap.add_argument("--stage2", required=True)
-    ap.add_argument("--stage1")
+    ap.add_argument("--stage1", help="a trained DialLocator checkpoint (deprecated)")
+    ap.add_argument("--coco", action="store_true",
+                    help="locate the dial with a COCO-pretrained detector. Measured at "
+                         "MAE 73.6 min against 120.7 for the trained locator on the same "
+                         "200 real photographs.")
+    ap.add_argument("--coco-arch", default="mobilenet", choices=["mobilenet", "resnet50"])
     ap.add_argument("--use-gt-dial", action="store_true",
                     help="crop with the label's own dial geometry (synthetic only)")
     ap.add_argument("--margin", type=float, default=1.25)
     ap.add_argument("--batch", type=int, default=32)
     ap.add_argument("--out", required=True)
+    ap.add_argument("--decode", choices=["joint", "independent"], default="joint",
+                    help="joint scores every time the clock could show and keeps the "
+                         "best; independent reads each head alone (the old behaviour)")
+    ap.add_argument("--dump-logits", help="write raw head distributions here, so "
+                    "decoders can be compared without re-running the network")
+    ap.add_argument("--time-weight", type=float, default=1.0,
+                    help="how loudly the whole-time head votes in the joint decode")
+    ap.add_argument("--cpu", action="store_true",
+                    help="stay off the GPU, so a prediction run can share the "
+                         "machine with a training run")
     args = ap.parse_args()
+    logit_dump = open(args.dump_logits, "w") if args.dump_logits else None
 
-    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+    device = torch.device("cpu" if args.cpu else
+                          "mps" if torch.backends.mps.is_available() else "cpu")
     reader, res = build_reader(args.stage2, device)
+    coco = CocoDialDetector(args.coco_arch, device=device) if args.coco else None
     locator = None
     if args.stage1:
         blob = torch.load(args.stage1, map_location="cpu", weights_only=False)
@@ -82,7 +105,11 @@ def main():
                 if not os.path.exists(path):
                     continue
                 img = Image.open(path).convert("RGB")
-                if args.use_gt_dial and r.get("render", {}).get("dial"):
+                if coco is not None:
+                    got = coco.locate(img)
+                    crop = (crop_dial(img, got[0], got[1], got[2], res, 1.15)
+                            if got else img.resize((res, res)))
+                elif args.use_gt_dial and r.get("render", {}).get("dial"):
                     d = r["render"]["dial"]
                     crop = crop_dial(img, d["cx"], d["cy"], d["r_max"], res, args.margin)
                 elif locator is not None:
@@ -97,8 +124,20 @@ def main():
             if not crops:
                 continue
             x = torch.stack(crops).to(device)
-            hl, ml = reader(x)
-            t, hour_only, dis, conf = decode_cls(hl.float().cpu(), ml.float().cpu())
+            out_heads = reader(x)
+            hl, ml, tl = (out_heads if len(out_heads) == 3
+                          else (out_heads[0], out_heads[1], None))
+            hl, ml = hl.float().cpu(), ml.float().cpu()
+            tl = tl.float().cpu() if tl is not None else None
+            t, hour_only, dis, conf = decode_cls(hl, ml)
+            if args.decode == "joint":
+                t, margin, swapped, consistency = joint_decode(hl, ml, tl,
+                                                               args.time_weight)
+            if logit_dump is not None:
+                for j, rid in enumerate(ids):
+                    logit_dump.write(json.dumps({"id": rid,
+                        "hour": [round(v, 4) for v in hl[j].tolist()],
+                        "minute": [round(v, 4) for v in ml[j].tolist()]}) + "\n")
             for j, rid in enumerate(ids):
                 hh = int(t[j].item() // 60) or 12
                 mm = t[j].item() - (t[j].item() // 60) * 60
@@ -108,9 +147,15 @@ def main():
                     "minutes": round(t[j].item(), 3),
                     "agreement_minutes": round(dis[j].item(), 3),
                     "sharpness": round(float(conf[j].mean()), 4),
+                    **({"margin": round(margin[j].item(), 4),
+                        "hands_swapped": bool(swapped[j]),
+                        "consistency": round(consistency[j].item(), 4)}
+                       if args.decode == "joint" else {}),
                 }) + "\n")
                 n += 1
     out.close()
+    if logit_dump:
+        logit_dump.close()
     print(f"wrote {n} predictions -> {args.out}")