Sync twostage.py for 263702

twostage.py

@@ -18,19 +18,28 @@ import math
 
 import torch
 import torch.nn as nn
+from PIL import Image
 
 from model import ConvBlock, ClockNetCls, decode_cls, soft_targets, soft_argmax_angle  # noqa: F401
 
 
 class DialLocator(nn.Module):
-    """Predicts the dial's centre and radius, normalised to image width.
+    """Predicts the dial's centre and its 2x2 shape matrix.
 
-    Radius is predicted in log space: dial sizes span a factor of three across
-    the corpus, and a plain linear target makes the large ones dominate the loss.
+    A centre and a radius describe a circle, and a dial viewed off-axis is an
+    ellipse. The 2x2 matrix carries foreshortening, rotation and scale together,
+    which is what makes rectification possible: invert it and the dial comes
+    back frontal and upright.
+
+    Output is (cx, cy, a, b, c, d) where the matrix maps dial-local coordinates
+    to image offsets.
     """
 
-    def __init__(self, width=16):
+    def __init__(self, width=16, outputs=6):
         super().__init__()
+        # outputs=6 is (cx, cy) + the 2x2 matrix. outputs=3 is the older
+        # (cx, cy, log r) head, kept so existing checkpoints still load.
+        self.outputs = outputs
         w = width
         self.net = nn.Sequential(
             ConvBlock(3, w, stride=2),
@@ -38,15 +47,40 @@ class DialLocator(nn.Module):
             ConvBlock(w * 2, w * 4, stride=2), ConvBlock(w * 4, w * 4),
             ConvBlock(w * 4, w * 8, stride=2), ConvBlock(w * 8, w * 8, stride=2),
             nn.AdaptiveAvgPool2d(4), nn.Flatten(),
-            nn.Linear(w * 8 * 16, 128), nn.SiLU(inplace=True), nn.Linear(128, 3),
+            nn.Linear(w * 8 * 16, 128), nn.SiLU(inplace=True), nn.Linear(128, outputs),
         )
 
     def forward(self, x):
         out = self.net(x)
         cx = torch.sigmoid(out[:, 0])          # centre stays inside the frame
         cy = torch.sigmoid(out[:, 1])
-        log_r = out[:, 2]
-        return torch.stack([cx, cy, log_r], dim=1)
+        return torch.cat([cx.unsqueeze(1), cy.unsqueeze(1), out[:, 2:]], dim=1)
+
+
+def rectify_dial(img, dial, out_res, margin=1.15):
+    """Warp a tilted dial back to a frontal, upright circle.
+
+    The dial is a circle in the world and an ellipse in the image. Its 2x2
+    matrix M maps dial-local coordinates to image offsets, so M inverse undoes
+    the foreshortening AND the in-image rotation in one step, handing stage 2 a
+    canonical face instead of a squashed one at an arbitrary angle.
+
+    PIL's AFFINE transform maps OUTPUT coordinates back into the INPUT, which is
+    exactly the direction we already have, so no inversion is needed: output
+    pixel -> dial-local -> input pixel.
+    """
+    W, H = img.size
+    M = dial["M"]
+    cx, cy = dial["cx"] * W, dial["cy"] * H
+    # pixels per unit of dial-local coordinate in the output
+    k = out_res / (2.0 * margin)
+    a, b = M[0][0] * W / k, M[0][1] * W / k
+    d, e = M[1][0] * H / k, M[1][1] * H / k
+    half = out_res / 2.0
+    c = cx - (a + b) * half
+    f = cy - (d + e) * half
+    return img.transform((out_res, out_res), Image.AFFINE, (a, b, c, d, e, f),
+                         resample=Image.BILINEAR)
 
 
 def crop_box(cx, cy, r, margin=1.25):
@@ -71,30 +105,47 @@ def crop_dial(img, cx, cy, r, out_res, margin=1.25, rng=None, jitter=0.0):
 
 
 def locator_targets(dials, device):
-    cx = torch.tensor([d["cx"] for d in dials], dtype=torch.float32, device=device)
-    cy = torch.tensor([d["cy"] for d in dials], dtype=torch.float32, device=device)
-    r = torch.tensor([d["r_max"] for d in dials], dtype=torch.float32, device=device)
-    return torch.stack([cx, cy, torch.log(r.clamp_min(1e-4))], dim=1)
+    rows = []
+    for d in dials:
+        M = d["M"]
+        rows.append([d["cx"], d["cy"], M[0][0], M[0][1], M[1][0], M[1][1]])
+    return torch.tensor(rows, dtype=torch.float32, device=device)
 
 
 def locator_loss(pred, tgt):
-    """Centre in plain units, radius in log units, weighted so neither dominates."""
+    """Centre and matrix, weighted so neither dominates.
+
+    Matrix entries run about 0.1 to 0.8 and the centre runs 0 to 1, so the two
+    are already on comparable scales; the centre is weighted up because getting
+    it wrong moves the whole crop while a slightly wrong matrix only skews it.
+    """
     centre = ((pred[:, :2] - tgt[:, :2]) ** 2).mean()
-    radius = ((pred[:, 2] - tgt[:, 2]) ** 2).mean()
-    return centre * 4.0 + radius, centre.detach(), radius.detach()
+    shape = ((pred[:, 2:6] - tgt[:, 2:6]) ** 2).mean()
+    return centre * 4.0 + shape, centre.detach(), shape.detach()
 
 
 class PretrainedReader(nn.Module):
-    """ImageNet-pretrained backbone + two angular heads.
+    """ImageNet-pretrained backbone + two angular heads and a whole-time head.
 
     A from-scratch CNN memorises 400 crops but sits on a plateau at 10.39
     (= 2*ln(180), a uniform distribution) on 10,800. Yang/Xie/Zisserman used
     ImageNet-pretrained ResNet-50 for the same task; the features that make a
     dial's hands salient are exactly the low-level edge and orientation filters
     ImageNet training already provides.
+
+    The third head exists because of what the two angular heads get wrong. On
+    real photographs 27.5% of readings put the minute right and the hour wrong,
+    split between hands read in the wrong roles and hours simply missed. Both
+    angular heads describe a hand, so a face whose hands are hard to tell apart
+    corrupts both of them at once and no amount of decoding recovers the hour --
+    measured, in jointdecode.py. `time_head` is asked instead for the time
+    itself over the whole 12-hour ring, coarsely, with no notion of a hand. It
+    cannot express a role confusion because it never assigns roles, and it only
+    has to be good enough to break the twelve-way tie: the minute hand, which
+    the model already reads well, supplies the precision.
     """
 
-    def __init__(self, bins=180, arch="mobilenet_v3_small", pretrained=True):
+    def __init__(self, bins=180, arch="mobilenet_v3_small", pretrained=True, time_bins=144):
         super().__init__()
         import torchvision.models as tvm
         self.bins = bins
@@ -107,6 +158,10 @@ class PretrainedReader(nn.Module):
             net = tvm.resnet18(weights=tvm.ResNet18_Weights.DEFAULT if pretrained else None)
             self.features = nn.Sequential(*list(net.children())[:-2])
             feat_dim = 512
+        elif arch == "resnet50":
+            net = tvm.resnet50(weights=tvm.ResNet50_Weights.DEFAULT if pretrained else None)
+            self.features = nn.Sequential(*list(net.children())[:-2])
+            feat_dim = 2048
         else:
             raise ValueError(arch)
         self.pool = nn.AdaptiveAvgPool2d(2)          # keep some spatial layout
@@ -114,7 +169,69 @@ class PretrainedReader(nn.Module):
                                    nn.SiLU(inplace=True), nn.Dropout(0.1))
         self.hour_head = nn.Linear(512, bins)
         self.minute_head = nn.Linear(512, bins)
+        self.time_bins = time_bins
+        self.time_head = nn.Linear(512, time_bins) if time_bins else None
 
     def forward(self, x):
         f = self.trunk(self.pool(self.features(x)))
-        return self.hour_head(f), self.minute_head(f)
+        t = self.time_head(f) if self.time_head is not None else None
+        return self.hour_head(f), self.minute_head(f), t
+
+
+class CocoDialDetector:
+    """Stage 1: find the dial with a COCO-pretrained detector.
+
+    This replaces DialLocator. Measured on 200 held-out real photographs, same
+    reader, three ways of producing a crop:
+
+        whole image, no locator      MAE 120.6 min   median 88.3
+        DialLocator (2.2 MB)         MAE 120.7 min   median 67.9
+        this (78 MB)                 MAE  73.6 min   median 18.4
+
+    DialLocator was trained on renders and contributes nothing over having no
+    locator at all -- 120.7 against 120.6. Distilling this detector into it
+    could not have helped either: a student cannot beat its teacher, and the
+    training would have run through the renderer whose domain gap is the thing
+    being escaped.
+
+    It costs 75 MB. That buys halving the error, so the single-digit-MB target
+    loses.
+    """
+
+    COCO_CLOCK = 85
+
+    def __init__(self, arch="mobilenet", score_thresh=0.3, device=None, det_res=640):
+        import torchvision.models.detection as tvd
+        if arch == "mobilenet":
+            # ResNet50-FPN trips the M2 GPU watchdog under sustained load
+            # (kIOGPUCommandBufferCallbackErrorImpactingInteractivity).
+            w = tvd.FasterRCNN_MobileNet_V3_Large_FPN_Weights.DEFAULT
+            self.net = tvd.fasterrcnn_mobilenet_v3_large_fpn(weights=w)
+        else:
+            w = tvd.FasterRCNN_ResNet50_FPN_V2_Weights.DEFAULT
+            self.net = tvd.fasterrcnn_resnet50_fpn_v2(weights=w)
+        self.device = device or torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+        self.net = self.net.eval().to(self.device)
+        self.score_thresh = score_thresh
+        self.det_res = det_res
+
+    @torch.no_grad()
+    def locate(self, img):
+        """(cx, cy, r) normalised to image width, or None when nothing is found."""
+        import numpy as np
+        scale = min(1.0, self.det_res / max(img.size))
+        small = (img.resize((max(32, round(img.width * scale)),
+                             max(32, round(img.height * scale)))) if scale < 1.0 else img)
+        x = torch.from_numpy(np.asarray(small.convert("RGB"), dtype=np.float32) / 255.0)
+        x = x.permute(2, 0, 1).to(self.device)
+        out = self.net([x])[0]
+        best, best_score = None, 0.0
+        for box, label, score in zip(out["boxes"], out["labels"], out["scores"]):
+            s = float(score)
+            if int(label) == self.COCO_CLOCK and s >= self.score_thresh and s > best_score:
+                best, best_score = [float(v) / scale for v in box], s
+        if best is None:
+            return None
+        x0, y0, x1, y1 = best
+        return ((x0 + x1) / 2 / img.width, (y0 + y1) / 2 / img.height,
+                max(x1 - x0, y1 - y0) / 2 / img.width, best_score)