#!/usr/bin/env python3 """Run the two-stage model over a dataset and write predictions for eval.py. python3 predict.py --labels clockface-external/labels.jsonl \ --root clockface-external --stage2 checkpoints/v2mnv3/stage2_best.pt \ --out preds.jsonl Each prediction carries `agreement_minutes`: how far apart the hour hand's reading and the minute hand's reading are. On a real clock they agree, so disagreement is a confidence signal that costs nothing to produce. With --stage1 the dial is located first and the crop comes from that. Without it, the whole image is used, which is what you want only when the clock already fills the frame. """ from __future__ import annotations import argparse import json import os import numpy as np import torch from PIL import Image from jointdecode import joint_decode from model import ClockNetCls, decode_cls 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) def to_tensor(img): x = torch.from_numpy(np.asarray(img, dtype=np.float32).copy() / 255.0).permute(2, 0, 1) return (x - MEAN) / STD 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, 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) def main(): ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) 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", 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("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) locator = DialLocator().to(device).eval() locator.load_state_dict(blob["model"]) loc_res = blob["args"].get("res", 256) rows = [json.loads(l) for l in open(args.labels) if l.strip()] out = open(args.out, "w") n = 0 with torch.no_grad(): for i in range(0, len(rows), args.batch): chunk = rows[i:i + args.batch] crops, ids = [], [] for r in chunk: path = os.path.join(args.root, r["file"]) if not os.path.exists(path): continue img = Image.open(path).convert("RGB") 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: small = to_tensor(img.resize((loc_res, loc_res))).unsqueeze(0).to(device) p = locator(small)[0].float().cpu() cx, cy, r_ = p[0].item(), p[1].item(), float(np.exp(p[2].item())) crop = crop_dial(img, cx, cy, r_, res, args.margin) else: crop = img.resize((res, res)) crops.append(to_tensor(crop)) ids.append(r["id"]) if not crops: continue x = torch.stack(crops).to(device) 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 out.write(json.dumps({ "id": rid, "time": f"{hh}:{int(round(mm)) % 60:02d}", "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}") if __name__ == "__main__": main()