| |
| """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 model import ClockNetCls, decode_cls |
| from twostage import 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") |
| model = (ClockNetCls(width=a.get("width", 32), bins=a.get("bins", 180)) |
| if backbone == "scratch" else |
| PretrainedReader(bins=a.get("bins", 180), arch=backbone)) |
| 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") |
| 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) |
| args = ap.parse_args() |
|
|
| device = torch.device("mps" if torch.backends.mps.is_available() else "cpu") |
| reader, res = build_reader(args.stage2, device) |
| 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 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) |
| hl, ml = reader(x) |
| t, hour_only, dis, conf = decode_cls(hl.float().cpu(), ml.float().cpu()) |
| 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), |
| }) + "\n") |
| n += 1 |
| out.close() |
| print(f"wrote {n} predictions -> {args.out}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
