Add twostage.py

twostage.py

@@ -0,0 +1,120 @@
+"""Two-stage clock reading: find the dial, then read it.
+
+A single-stage CNN has to locate the face and measure two angles at once, from
+an image where the clock may fill 30% of the frame. Splitting those apart gives
+each stage an easier job, and the supervision for stage 1 is free because the
+renderer knows exactly where it put the dial.
+
+    stage 1  DialLocator  full image -> (cx, cy, r) of the dial
+    stage 2  ClockReader  a crop around the dial -> two angular distributions
+
+Stage 2 is trained on ground-truth crops with jitter, so it tolerates the error
+stage 1 will actually make at inference.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+
+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.
+
+    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.
+    """
+
+    def __init__(self, width=16):
+        super().__init__()
+        w = width
+        self.net = nn.Sequential(
+            ConvBlock(3, w, stride=2),
+            ConvBlock(w, w * 2, stride=2), ConvBlock(w * 2, w * 2),
+            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),
+        )
+
+    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)
+
+
+def crop_box(cx, cy, r, margin=1.25):
+    """Square box around a dial, in normalised coordinates."""
+    half = r * margin
+    return cx - half, cy - half, cx + half, cy + half
+
+
+def crop_dial(img, cx, cy, r, out_res, margin=1.25, rng=None, jitter=0.0):
+    """Crop a PIL image around the dial and resize.
+
+    With jitter > 0 the box is perturbed, so stage 2 sees the kind of imperfect
+    box stage 1 will hand it rather than only perfect ones.
+    """
+    W, H = img.size
+    if rng is not None and jitter > 0:
+        cx += rng.uniform(-jitter, jitter) * r
+        cy += rng.uniform(-jitter, jitter) * r
+        r *= 1.0 + rng.uniform(-jitter, jitter)
+    x0, y0, x1, y1 = crop_box(cx * W, cy * H, r * W, margin)
+    return img.crop((int(x0), int(y0), int(x1), int(y1))).resize((out_res, out_res))
+
+
+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)
+
+
+def locator_loss(pred, tgt):
+    """Centre in plain units, radius in log units, weighted so neither dominates."""
+    centre = ((pred[:, :2] - tgt[:, :2]) ** 2).mean()
+    radius = ((pred[:, 2] - tgt[:, 2]) ** 2).mean()
+    return centre * 4.0 + radius, centre.detach(), radius.detach()
+
+
+class PretrainedReader(nn.Module):
+    """ImageNet-pretrained backbone + two angular heads.
+
+    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.
+    """
+
+    def __init__(self, bins=180, arch="mobilenet_v3_small", pretrained=True):
+        super().__init__()
+        import torchvision.models as tvm
+        self.bins = bins
+        if arch == "mobilenet_v3_small":
+            net = tvm.mobilenet_v3_small(weights=tvm.MobileNet_V3_Small_Weights.DEFAULT
+                                         if pretrained else None)
+            self.features = net.features
+            feat_dim = 576
+        elif arch == "resnet18":
+            net = tvm.resnet18(weights=tvm.ResNet18_Weights.DEFAULT if pretrained else None)
+            self.features = nn.Sequential(*list(net.children())[:-2])
+            feat_dim = 512
+        else:
+            raise ValueError(arch)
+        self.pool = nn.AdaptiveAvgPool2d(2)          # keep some spatial layout
+        self.trunk = nn.Sequential(nn.Flatten(), nn.Linear(feat_dim * 4, 512),
+                                   nn.SiLU(inplace=True), nn.Dropout(0.1))
+        self.hour_head = nn.Linear(512, bins)
+        self.minute_head = nn.Linear(512, bins)
+
+    def forward(self, x):
+        f = self.trunk(self.pool(self.features(x)))
+        return self.hour_head(f), self.minute_head(f)