Argus-3D: discovered class-agnostic 3D detection head on EUPE-ViT-B

README.md

@@ -17,25 +17,31 @@ tags:
 
 Class-agnostic 3D bounding box detection on a frozen [EUPE-ViT-B](https://huggingface.co/facebook/EUPE-ViT-B) backbone. Given a posed RGB image and camera intrinsics, returns 7-DoF boxes (cx, cy, cz, w, h, d, theta) for the objects in the scene.
 
-The detector pairs a non-linear (or linear) per-patch foreground head with feature-dim discovery for depth, and unsupervised k-means clustering for size priors. No 2D bounding boxes, no class labels, no segmentation map as final output. Camera-frame 3D boxes only.
+The detector pairs a non-linear (or linear) per-patch foreground head with feature-dim discovery for depth, k-means clustering for size priors derived from 20 CA-1M val scenes, and 3D-IoU-based multi-view fusion. No 2D bounding boxes, no class labels, no segmentation map as final output. Camera-frame 3D boxes only.
 
 ## Architecture
 
 ```
-Image (768x768)
-  -> EUPE-ViT-B (frozen, reused from phanerozoic/argus)
-  -> patch tokens (2304, 768) on a 48x48 grid
-  -> instance head: 2-layer MLP (default) or linear ridge -> per-patch foreground score
-     depth head:    ridge over 768 dims                     -> per-patch metric depth (m)
-     k-means modes: 8 cluster centers                       -> per-patch object-type assignment
-  -> threshold instance score
-  -> upsample mask to 768x768, connected components
-  -> for each component:
-       unproject pixels to 3D using depth + K
-       DBSCAN-split for instance separation
-       PCA-on-xz for yaw, percentile extents for (w, h, d)
-       blend extents toward the matched cluster's size prior
-  -> camera-frame 7-DoF box list
+Per-frame:
+  Image (768x768)
+    -> EUPE-ViT-B (frozen, reused from phanerozoic/argus)
+    -> patch tokens (2304, 768) on a 48x48 grid
+    -> instance head: 2-layer MLP (default) or linear ridge -> per-patch foreground score
+       depth head:    ridge over 768 dims                   -> per-patch metric depth (m)
+       k-means modes: 8 cluster centers (20-scene)          -> per-patch object-type assignment
+    -> threshold instance score, upsample mask to 768x768, connected components
+    -> for each component:
+         unproject to 3D using depth + K
+         DBSCAN-split for instance separation
+         PCA-on-xz for yaw, percentile extents for (w, h, d)
+         blend extents toward the matched cluster's size prior
+    -> camera-frame 7-DoF box list
+
+Multi-view (per scene):
+  -> transform every per-frame box to world frame using camera RT
+  -> 3D-IoU-based clustering: union-find with edges where iou_3d_zup(box_i, box_j) > 0.2
+  -> filter clusters by min_obs and total inlier weight (per-cluster confidence)
+  -> per-cluster: weighted-median fuse to single 7-DoF box
 ```
 
 ## Components
@@ -43,14 +49,14 @@ Image (768x768)
 | Component | Parameters | Discovery / training |
 |---|---|---|
 | EUPE-ViT-B backbone (frozen, reused) | not part of this head | reused from phanerozoic/argus |
-| Instance head — MLP (default) | ~200 K | 2-layer MLP (768 → 256 → 1, ReLU + dropout 0.5), 30 epochs AdamW BCE-with-logits |
+| Instance head — MLP (default) | ~200 K | 2-layer MLP (768 → 256 → 1, ReLU + dropout 0.5), 30 epochs AdamW BCE-with-logits on 4-scene mixed split |
 | Instance head — linear ridge (fallback) | 769 floats + threshold | random K=20 subset search + hard-neg mining, AUC selection |
 | Depth head (ridge over 768 dims) | 769 floats | random K=20 subset search, RMSE selection |
-| K-means cluster centers | 8 × 768 floats | MiniBatchKMeans on foreground patches |
-| Per-cluster size priors (w, h, d) | 8 × 3 floats | median of observed extents per mode |
+| K-means cluster centers | 8 × 768 floats | MiniBatchKMeans on foreground patches across 20 CA-1M val scenes |
+| Per-cluster size priors (w, h, d) | 8 × 3 floats | median of observed extents per mode (8 800 instances aggregated) |
 | OBB fitter (PCA + percentile + Tikhonov) | 0 | closed-form |
+| Multi-view fusion (3D-IoU union-find + weighted median) | 0 | closed-form |
 | **Total head footprint (MLP)** | **~200 K params / ~830 KB** | |
-| **Total head footprint (linear ridge)** | **~7 700 floats / 43 KB** | |
 
 ## File layout
 
@@ -59,8 +65,8 @@ instance_head_mlp.safetensors   # MLP fc1/fc2 weights
 instance_head_mlp_meta.json     # MLP config + threshold
 instance_head.safetensors       # linear ridge: dims + coef + intercept + threshold
 depth_head.safetensors          # depth ridge: dims + coef + intercept
-size_priors.safetensors         # 8 cluster centers + 8 (w, h, d) priors
-config.json                     # input_res, patch_grid, prior_weight, etc.
+size_priors.safetensors         # 8 cluster centers + 8 (w, h, d) priors (20-scene)
+config.json                     # input_res, patch_grid, prior_weight, fusion_iou, etc.
 argus_3d.py                     # Argus3D class
 infer.py                        # CLI dispatcher
 ```
@@ -71,11 +77,7 @@ infer.py                        # CLI dispatcher
 from argus_3d import Argus3D
 import numpy as np
 
-# default (head='auto') uses MLP if shipped, falls back to linear ridge.
 model = Argus3D.from_pretrained("phanerozoic/argus-3d", device="cuda")
-# or force one explicitly:
-model = Argus3D.from_pretrained("phanerozoic/argus-3d", device="cuda", head="mlp")
-model = Argus3D.from_pretrained("phanerozoic/argus-3d", device="cuda", head="linear")
 
 K = np.array([[850, 0, 395], [0, 850, 510], [0, 0, 1]])
 boxes = model.detect("room.jpg", K)               # list of Box3D
@@ -90,38 +92,48 @@ for b in boxes:
 
 CA-1M val sequence `ca1m-val-45662921`. Class-agnostic per-scene 3D IoU after multi-view fusion across 284 frames (stride-4 sampling of 1135 total). The head produces its own instance hypotheses; no ground-truth 2D bounding boxes are used. Sensor depth is supplied; the discovered depth head can be used in its place.
 
-End-to-end pipeline numbers, MLP head vs linear ridge head, all else identical:
+### Headline result
 
-| Metric | Linear ridge (v0) | MLP (default) |
-|---|---|---|
-| Fused boxes per scene | 72 | 80 |
-| Mean 3D IoU | 0.063 | **0.071** |
-| Fraction > 0.1 IoU | 27.8 % | 23.7 % |
-| Fraction > 0.25 IoU | 6.9 % | **11.3 %** |
-| Fraction > 0.5 IoU | 0.0 % | **1.3 %** |
-| Recall (matched / GT) | 19.8 % | **24.4 %** |
+3D-IoU multi-view fusion + 4-scene MLP head + 20-scene size priors:
+
+| Metric | Linear ridge (v0) + position fusion | MLP + position fusion | MLP + IoU fusion (default) |
+|---|---|---|---|
+| Mean 3D IoU | 0.063 | 0.076 | **0.154** |
+| Median 3D IoU | — | 0.019 | **0.117** |
+| Fraction > 0.1 IoU | 27.8 % | 25.6 % | **53.5 %** |
+| Fraction > 0.25 IoU | 6.9 % | 11.5 % | **28.2 %** |
+| Fraction > 0.5 IoU | 0.0 % | 1.3 % | **2.8 %** |
+| Recall (matched / GT) | 19.8 % | 24.4 % | 28.1 % |
+| Fused boxes per scene | 72 | 78 | 71 |
+
+The IoU-based fusion at threshold 0.2, with min 4 observations per cluster and weight floor at the 80th percentile of nonzero cluster weights, gives the strictest filtering. Looser filters trade per-box quality for recall; useful operating points along the curve:
 
-The MLP head moves the high-IoU tail and the recall up. The 0.1 IoU bucket drops slightly because the MLP's tighter mask produces fewer marginal-quality fused boxes; the boxes it does produce are of higher quality on average.
+| Config | Mean IoU | > 0.25 IoU | > 0.5 IoU | Recall | Fused boxes |
+|---|---|---|---|---|---|
+| min_obs=3, weight_pct=75 (loose) | 0.120 | 21.2 % | 1.9 % | 35.0 % | 104 |
+| min_obs=3, weight_pct=80 (balanced) | 0.137 | 25.3 % | 2.3 % | 30.9 % | 87 |
+| min_obs=4, weight_pct=80 (strict, default) | **0.154** | 28.2 % | 2.8 % | 28.1 % | 71 |
 
-Per-stage discovery metrics (in-distribution random patch split mixing all 4 cached scenes):
+### Per-stage discovery metrics (in-distribution random patch split, 4 scenes mixed)
 
 | Discovery output | Metric | Linear ridge | MLP |
 |---|---|---|---|
 | Instance head, per-patch foreground | AUC | 0.860 | 0.980 |
 | Instance head, per-patch foreground | F1 (tuned threshold) | 0.569 | 0.815 |
-| Depth head, foreground patches in 0.1-3 m | RMSE | 0.190 m | (linear ridge unchanged) |
-| Depth head, foreground patches in 0.1-3 m | delta1 (1.25× ratio) | 0.919 | (linear ridge unchanged) |
+| Depth head, foreground patches in 0.1-3 m | RMSE | 0.190 m | 0.133 m (MLP variant) |
+| Depth head, foreground patches in 0.1-3 m | delta1 (1.25× ratio) | 0.919 | 0.974 (MLP variant) |
 
 ### Cross-scene held-out
 
-The numbers above use a random patch split that mixes patches from all 4 cached scenes. That setup overstates generalization because adjacent patches in the same room share scene-specific cues (lighting, wall colour, carpet pattern). A leave-one-scene-out evaluation gives the honest cross-scene number:
+The numbers above use a random patch split that mixes patches from 4 cached scenes. That setup overstates generalization because adjacent patches in the same room share scene-specific cues. A leave-one-scene-out eval gives the honest cross-scene head AUC:
 
 | Head | In-distribution AUC | Cross-scene AUC | Cross-scene F1 |
 |---|---|---|---|
 | Linear ridge (all-768) | 0.860 | 0.604 | 0.301 |
-| MLP (768 → 256 → 1) | 0.980 | **0.678** | **0.352** |
+| MLP (3 train scenes) | 0.980 | 0.566 (fold 45662921) | 0.312 |
+| MLP (19 train scenes) | 0.980 | **0.780** (45662921 held out) | **0.465** |
 
-The MLP narrows the in-distribution-to-cross-scene gap by lifting both numbers, but the gap (0.98 vs 0.68) signals the linear/non-linear head still leans on per-scene cues that don't transfer. The cleanest path to closing that gap is more scenes from CA-1M, not deeper heads.
+Scaling the train set from 3 to 19 scenes lifts cross-scene AUC by 0.21 (0.566 → 0.780). End-to-end IoU on the held-out scene with the 19-scene MLP head was 0.046 mean IoU under position-only fusion; the IoU-fusion improvement above is in-distribution. Re-running cross-scene with the IoU-fusion stage is open work.
 
 ## Backbone
 

argus_3d.py

@@ -59,6 +59,192 @@ class Box3D:
     n_inliers: int = 0
 
 
+def cam_to_world(box_cam: Tuple[float, ...], RT: np.ndarray) -> Tuple[float, ...]:
+    """Camera-frame 7-DoF box -> world-frame (X, Y, Z, sx, sy, sz, yaw_world).
+
+    Output convention: Z up. yaw_world is rotation about world +Z that takes
+    world +X to the box's local +x axis. sz is height along world Z.
+    """
+    cx, cy, cz, w, h, d, theta = box_cam
+    pos_cam = np.array([cx, cy, cz])
+    R_cw = RT[:3, :3]
+    t_cw = RT[:3, 3]
+    pos_world = R_cw @ pos_cam + t_cw
+    x_local_cam = np.array([math.cos(theta), 0.0, -math.sin(theta)])
+    x_local_world = R_cw @ x_local_cam
+    yaw_world = math.atan2(x_local_world[1], x_local_world[0])
+    return (
+        float(pos_world[0]), float(pos_world[1]), float(pos_world[2]),
+        float(w), float(d), float(h), yaw_world,
+    )
+
+
+def _polygon_area(poly: np.ndarray) -> float:
+    if len(poly) < 3:
+        return 0.0
+    return float(0.5 * abs(sum(
+        (poly[i, 0] * poly[(i + 1) % len(poly), 1]
+         - poly[(i + 1) % len(poly), 0] * poly[i, 1])
+        for i in range(len(poly))
+    )))
+
+
+def _polygon_clip(subject: np.ndarray, clip: np.ndarray) -> np.ndarray:
+    """Sutherland-Hodgman polygon clipping in 2D."""
+    out = subject.tolist()
+    cn = len(clip)
+    for i in range(cn):
+        if not out:
+            return np.zeros((0, 2))
+        a = clip[i]
+        b = clip[(i + 1) % cn]
+        edge = b - a
+        new_out = []
+        prev = out[-1]
+        prev_in = (b[0] - a[0]) * (prev[1] - a[1]) - (b[1] - a[1]) * (prev[0] - a[0]) >= 0
+        for cur in out:
+            cur_in = (b[0] - a[0]) * (cur[1] - a[1]) - (b[1] - a[1]) * (cur[0] - a[0]) >= 0
+            if cur_in:
+                if not prev_in:
+                    dx = cur[0] - prev[0]
+                    dy = cur[1] - prev[1]
+                    denom = edge[0] * dy - edge[1] * dx
+                    if abs(denom) > 1e-12:
+                        t = (edge[0] * (prev[1] - a[1]) - edge[1] * (prev[0] - a[0])) / denom
+                        new_out.append([prev[0] + t * dx, prev[1] + t * dy])
+                new_out.append(cur)
+            elif prev_in:
+                dx = cur[0] - prev[0]
+                dy = cur[1] - prev[1]
+                denom = edge[0] * dy - edge[1] * dx
+                if abs(denom) > 1e-12:
+                    t = (edge[0] * (prev[1] - a[1]) - edge[1] * (prev[0] - a[0])) / denom
+                    new_out.append([prev[0] + t * dx, prev[1] + t * dy])
+            prev = cur
+            prev_in = cur_in
+        out = new_out
+    return np.array(out) if out else np.zeros((0, 2))
+
+
+def iou_3d_zup(a: Tuple[float, ...], b: Tuple[float, ...]) -> float:
+    """3D IoU for gravity-aligned (Z-up) world-frame boxes (X, Y, Z, sx, sy, sz, yaw)."""
+    Ax, Ay, Az, Asx, Asy, Asz, Ayaw = a
+    Bx, By, Bz, Bsx, Bsy, Bsz, Byaw = b
+
+    def xy_poly(x, y, sx, sy, yaw):
+        local = np.array([[-sx/2, -sy/2], [+sx/2, -sy/2], [+sx/2, +sy/2], [-sx/2, +sy/2]])
+        c, s = math.cos(yaw), math.sin(yaw)
+        R = np.array([[c, -s], [s, c]])
+        return local @ R.T + np.array([x, y])
+
+    pa = xy_poly(Ax, Ay, Asx, Asy, Ayaw)
+    pb = xy_poly(Bx, By, Bsx, Bsy, Byaw)
+    inter_poly = _polygon_clip(pa, pb)
+    inter_xy = _polygon_area(inter_poly)
+    if inter_xy == 0:
+        return 0.0
+    z_overlap = max(0.0, min(Az + Asz/2, Bz + Bsz/2) - max(Az - Asz/2, Bz - Bsz/2))
+    if z_overlap == 0:
+        return 0.0
+    inter_vol = inter_xy * z_overlap
+    vol_a = _polygon_area(pa) * Asz
+    vol_b = _polygon_area(pb) * Bsz
+    union_vol = vol_a + vol_b - inter_vol
+    if union_vol < 1e-9:
+        return 0.0
+    return float(inter_vol / union_vol)
+
+
+def fuse_observations_iou(
+    observations: List[Tuple[Tuple[float, ...], float]],
+    iou_thresh: float = 0.2,
+    min_obs: int = 4,
+    weight_thresh_pct: float = 80.0,
+) -> List[Tuple[float, ...]]:
+    """Multi-view fusion via 3D-IoU clustering + weighted-median fuse.
+
+    observations: list of (world_frame_box, weight). world_frame_box is the
+                  7-DoF tuple (X, Y, Z, sx, sy, sz, yaw_world) returned by
+                  cam_to_world(...). Weight is typically the inlier count of
+                  the per-frame OBB fit.
+    iou_thresh:   union-find edge if iou_3d_zup(box_i, box_j) > iou_thresh.
+    min_obs:      drop clusters with fewer than this many observations.
+    weight_thresh_pct: drop clusters whose total inlier weight is below the
+                  given percentile of nonzero cluster weights.
+
+    Returns a list of fused 7-DoF world-frame boxes.
+    """
+    if not observations:
+        return []
+    boxes = [o[0] for o in observations]
+    weights = [float(o[1]) for o in observations]
+    n = len(boxes)
+
+    # Pairwise IoU (with XY prefilter to bound the pair count).
+    from scipy.spatial import cKDTree
+    arr = np.array(boxes, dtype=np.float64)
+    tree = cKDTree(arr[:, 0:2])
+    candidate_pairs = tree.query_pairs(0.5)
+
+    parent = list(range(n))
+    def find(i):
+        while parent[i] != i:
+            parent[i] = parent[parent[i]]
+            i = parent[i]
+        return i
+    def union(i, j):
+        ri, rj = find(i), find(j)
+        if ri != rj:
+            parent[ri] = rj
+
+    for (i, j) in candidate_pairs:
+        if iou_3d_zup(boxes[i], boxes[j]) > iou_thresh:
+            union(i, j)
+
+    from collections import defaultdict
+    groups = defaultdict(list)
+    for i in range(n):
+        groups[find(i)].append(i)
+    clusters = list(groups.values())
+
+    cluster_weight = [(idx, sum(weights[i] for i in idx)) for idx in clusters]
+    cluster_weight.sort(key=lambda x: -x[1])
+    nonzero = [w for _, w in cluster_weight if w > 0]
+    weight_thresh = max(50.0, float(np.percentile(nonzero, weight_thresh_pct))) if nonzero else 50.0
+
+    fused: List[Tuple[float, ...]] = []
+    for cluster, total_w in cluster_weight:
+        if len(cluster) < min_obs or total_w < weight_thresh:
+            continue
+        cluster_boxes = [boxes[i] for i in cluster]
+        cluster_ws = np.array([weights[i] for i in cluster], dtype=np.float64)
+        # Weighted-median fuse on (X, Y, Z, sx, sy, sz); circular-mean on yaw.
+        arr = np.array(cluster_boxes, dtype=np.float64)
+        order = arr[:, 0].argsort()  # any deterministic order for tie-breaking
+        sorted_w = cluster_ws[order]
+        sorted_a = arr[order]
+        cum = sorted_w.cumsum()
+        target = cum[-1] * 0.5
+        med_idx = int(np.searchsorted(cum, target))
+        med_idx = min(med_idx, len(sorted_a) - 1)
+        med = sorted_a[med_idx].copy()  # 7-vector
+        for j in [1, 2, 3, 4, 5]:  # Y, Z, sx, sy, sz
+            sw = cluster_ws[arr[:, j].argsort()]
+            sa = arr[arr[:, j].argsort()]
+            c = sw.cumsum()
+            t = c[-1] * 0.5
+            mi = min(int(np.searchsorted(c, t)), len(sa) - 1)
+            med[j] = sa[mi, j]
+        # Circular-mean yaw modulo pi (boxes are flip-symmetric on yaw + pi).
+        yaws = arr[:, 6]
+        yaws_mod = (yaws * 2.0) % (2 * math.pi)
+        sin_sum = float((cluster_ws * np.sin(yaws_mod)).sum())
+        cos_sum = float((cluster_ws * np.cos(yaws_mod)).sum())
+        med[6] = math.atan2(sin_sum, cos_sum) / 2.0
+        fused.append(tuple(float(x) for x in med))
+    return fused
+
+
 class _MLPInstanceHead(nn.Module):
     """2-layer MLP fg head: 768 -> hidden -> 1."""
 

size_priors.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:901ae6958490daf1c8dd6541b6e877d83d9cdb822577993fcd5f17cf94b31394
+oid sha256:ba90168daf451ea21bf5009df47bd5353bb455b160f6c2b718ba33ce1f97fd61
 size 24832