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

README.md

@@ -0,0 +1,93 @@
+# Argus-3D
+
+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 head is built by feature-dim discovery and unsupervised clustering, not gradient training. 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: ridge over 768 dims  -> 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
+```
+
+## Components
+
+| Component | Parameters | Discovery method |
+|---|---|---|
+| EUPE-ViT-B backbone (frozen, reused) | not part of this head | reused from phanerozoic/argus |
+| Instance head (ridge over 768 dims) | 769 floats + 1 threshold | random K=20 subset search + hard-negative mining, AUC selection |
+| Depth head (ridge over 768 dims) | 769 floats | random K=20 subset search, RMSE selection |
+| K-means cluster centers | 8 x 768 floats | MiniBatchKMeans on foreground patches |
+| Per-cluster size priors (w, h, d) | 8 x 3 floats | median of observed extents per mode |
+| OBB fitter (PCA + percentile + Tikhonov) | 0 | closed-form |
+| **Total head footprint** | **~7,700 floats / 43 KB** | |
+
+## File layout
+
+```
+instance_head.safetensors   # ridge dims + coef + intercept + threshold
+depth_head.safetensors      # ridge dims + coef + intercept
+size_priors.safetensors     # 8 cluster centers + 8 (w, h, d) priors
+config.json                 # input_res, patch_grid, prior_weight, etc.
+argus_3d.py                 # Argus3D class
+infer.py                    # CLI dispatcher
+```
+
+## Usage
+
+```python
+from argus_3d import Argus3D
+import numpy as np
+
+model = Argus3D.from_pretrained("phanerozoic/argus-3d", device="cuda")
+
+K = np.array([[850, 0, 395], [0, 850, 510], [0, 0, 1]])
+boxes = model.detect("room.jpg", K)               # list of Box3D
+boxes = model.detect("room.jpg", K, depth=d)      # supply RGBD sensor depth
+out   = model.perceive("room.jpg", K)             # fg score map + depth map + boxes
+
+for b in boxes:
+    print(b.cx, b.cy, b.cz, b.w, b.h, b.d, b.theta)
+```
+
+## Eval
+
+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.
+
+| Metric | Value |
+|---|---|
+| Fused boxes per scene | 72 |
+| Mean 3D IoU | 0.063 |
+| Fraction > 0.1 IoU | 27.8% |
+| Fraction > 0.25 IoU | 6.9% |
+| Recall (matched / GT) | 19.8% |
+
+Per-stage discovery metrics:
+
+| Discovery output | Metric | Value |
+|---|---|---|
+| Instance head, per-patch foreground | AUC | 0.860 |
+| Instance head, per-patch foreground | F1 (tuned threshold) | 0.569 |
+| Depth head, foreground patches in 0.1-3 m | RMSE | 0.190 m |
+| Depth head, foreground patches in 0.1-3 m | delta1 (1.25x ratio) | 0.919 |
+
+## Backbone
+
+EUPE-ViT-B from Meta FAIR (arXiv:2603.22387) via [phanerozoic/argus](https://huggingface.co/phanerozoic/argus). The backbone is frozen and not modified by this repo.
+
+## License
+
+FAIR Research License (non-commercial), inherited via the EUPE-ViT-B backbone.

argus_3d.py

@@ -0,0 +1,346 @@
+"""Argus3D: class-agnostic 3D bounding box detection on EUPE-ViT-B.
+
+A single class-agnostic 3D detector head built by discovery, not gradient
+training. Given a posed RGB image and camera intrinsics, returns a list of
+7-DoF bounding boxes (cx, cy, cz, w, h, d, theta) in camera frame.
+
+Components (all loaded from safetensors at from_pretrained time):
+  - frozen EUPE-ViT-B backbone (reused from phanerozoic/argus)
+  - instance_head: per-patch foreground ridge, 769 floats + 1 threshold
+  - depth_head:    per-patch metric depth ridge, 769 floats
+  - size_priors:   8 k-means cluster centers (768 each) plus 8 (w, h, d) priors
+  - zero-parameter OBB fitter, DBSCAN instance separation, Tikhonov prior blend
+  - optional multi-view fusion utilities
+
+Use:
+    model = Argus3D.from_pretrained('phanerozoic/argus-3d').cuda().eval()
+    boxes = model.detect(image, K)               # camera-frame 7-DoF boxes
+    boxes = model.detect(image, K, depth=depth)  # supply sensor depth (RGBD)
+    out   = model.perceive(image, K)             # foreground mask + depth + boxes
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from PIL import Image
+
+DEFAULT_INPUT_RES = 768
+DEFAULT_PATCH_GRID = 48
+DEFAULT_PATCH_SIZE = DEFAULT_INPUT_RES // DEFAULT_PATCH_GRID  # 16
+
+
+@dataclass
+class Box3D:
+    cx: float
+    cy: float
+    cz: float
+    w: float
+    h: float
+    d: float
+    theta: float
+    score: float = 1.0
+    n_inliers: int = 0
+
+
+class Argus3D(nn.Module):
+    """Class-agnostic 3D detector on a frozen EUPE-ViT-B backbone."""
+
+    def __init__(
+        self,
+        instance_dims: torch.Tensor,
+        instance_coef: torch.Tensor,
+        instance_intercept: float,
+        instance_threshold: float,
+        depth_dims: torch.Tensor,
+        depth_coef: torch.Tensor,
+        depth_intercept: float,
+        cluster_centers: torch.Tensor,
+        size_priors: torch.Tensor,
+        prior_weight: float = 80.0,
+        input_res: int = DEFAULT_INPUT_RES,
+        patch_grid: int = DEFAULT_PATCH_GRID,
+    ):
+        super().__init__()
+        self.register_buffer("instance_dims", instance_dims.long())
+        self.register_buffer("instance_coef", instance_coef.float())
+        self.instance_intercept = float(instance_intercept)
+        self.instance_threshold = float(instance_threshold)
+        self.register_buffer("depth_dims", depth_dims.long())
+        self.register_buffer("depth_coef", depth_coef.float())
+        self.depth_intercept = float(depth_intercept)
+        self.register_buffer("cluster_centers", cluster_centers.float())
+        self.register_buffer("size_priors", size_priors.float())
+        self.prior_weight = prior_weight
+        self.input_res = input_res
+        self.patch_grid = patch_grid
+        self.patch_size = input_res // patch_grid
+        self.backbone: Optional[nn.Module] = None  # set externally
+
+    def attach_backbone(self, backbone: nn.Module) -> "Argus3D":
+        """Attach the frozen EUPE-ViT-B backbone (the trunk is not shipped)."""
+        self.backbone = backbone
+        for p in self.backbone.parameters():
+            p.requires_grad = False
+        return self
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        repo_or_dir: str,
+        backbone_repo: str = "phanerozoic/argus",
+        device: str = "cpu",
+    ) -> "Argus3D":
+        """Load heads from safetensors in `repo_or_dir`. Backbone is loaded
+        from `backbone_repo`'s frozen EUPE-ViT-B trunk."""
+        from safetensors.torch import load_file
+
+        if os.path.isdir(repo_or_dir):
+            base = repo_or_dir
+        else:
+            from huggingface_hub import snapshot_download
+            base = snapshot_download(repo_or_dir)
+
+        inst = load_file(os.path.join(base, "instance_head.safetensors"))
+        depth = load_file(os.path.join(base, "depth_head.safetensors"))
+        priors = load_file(os.path.join(base, "size_priors.safetensors"))
+        with open(os.path.join(base, "config.json"), "r") as f:
+            cfg = json.load(f)
+
+        model = cls(
+            instance_dims=inst["dims"],
+            instance_coef=inst["coef"],
+            instance_intercept=float(inst["intercept"].item()),
+            instance_threshold=float(inst["threshold"].item()),
+            depth_dims=depth["dims"],
+            depth_coef=depth["coef"],
+            depth_intercept=float(depth["intercept"].item()),
+            cluster_centers=priors["cluster_centers"],
+            size_priors=priors["priors_whd"],
+            prior_weight=cfg.get("prior_weight", 80.0),
+            input_res=cfg.get("input_res", DEFAULT_INPUT_RES),
+            patch_grid=cfg.get("patch_grid", DEFAULT_PATCH_GRID),
+        )
+
+        # Backbone reuse: load the frozen EUPE-ViT-B from argus.
+        from transformers import AutoModel
+        argus = AutoModel.from_pretrained(backbone_repo, trust_remote_code=True)
+        model.attach_backbone(argus.backbone.to(device).eval())
+        return model.to(device)
+
+    @staticmethod
+    def _imagenet_normalize(x: torch.Tensor) -> torch.Tensor:
+        mean = torch.tensor([0.485, 0.456, 0.406], device=x.device).view(1, 3, 1, 1)
+        std = torch.tensor([0.229, 0.224, 0.225], device=x.device).view(1, 3, 1, 1)
+        return (x - mean) / std
+
+    def _prepare_image(self, image) -> torch.Tensor:
+        if isinstance(image, str):
+            pil = Image.open(image).convert("RGB")
+        elif isinstance(image, np.ndarray):
+            pil = Image.fromarray(image).convert("RGB")
+        elif isinstance(image, Image.Image):
+            pil = image.convert("RGB")
+        else:
+            raise TypeError(f"unsupported image type: {type(image)}")
+        pil_in = pil.resize((self.input_res, self.input_res), Image.BILINEAR)
+        arr = np.array(pil_in).astype(np.float32) / 255.0
+        device = next(self.buffers()).device
+        return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).to(device), pil.size
+
+    @torch.no_grad()
+    def patch_features(self, image) -> Tuple[torch.Tensor, Tuple[int, int]]:
+        if self.backbone is None:
+            raise RuntimeError("backbone not attached; call attach_backbone or from_pretrained")
+        img_t, orig_size = self._prepare_image(image)
+        feats = self.backbone.forward_features(self._imagenet_normalize(img_t))
+        return feats["x_norm_patchtokens"][0], orig_size  # (N, 768)
+
+    @torch.no_grad()
+    def fg_score_map(self, image) -> Tuple[np.ndarray, Tuple[int, int]]:
+        """Per-patch foreground score map at patch_grid x patch_grid."""
+        patch_tokens, orig_size = self.patch_features(image)
+        f = patch_tokens.index_select(1, self.instance_dims)
+        scores = (f @ self.instance_coef + self.instance_intercept).cpu().numpy()
+        return scores.reshape(self.patch_grid, self.patch_grid), orig_size
+
+    @torch.no_grad()
+    def depth_map(self, image) -> Tuple[np.ndarray, Tuple[int, int]]:
+        """Per-patch metric depth at patch_grid x patch_grid (meters)."""
+        patch_tokens, orig_size = self.patch_features(image)
+        f = patch_tokens.index_select(1, self.depth_dims)
+        depths = (f @ self.depth_coef + self.depth_intercept).cpu().numpy()
+        return depths.reshape(self.patch_grid, self.patch_grid), orig_size
+
+    @torch.no_grad()
+    def detect(
+        self,
+        image,
+        K: np.ndarray,
+        depth: Optional[np.ndarray] = None,
+        return_internals: bool = False,
+    ) -> List[Box3D]:
+        """Camera-frame 7-DoF box list for the image.
+
+        K: 3x3 intrinsic matrix matching `image`'s native resolution.
+        depth: optional (H, W) sensor depth in meters at image resolution. If
+        None, the discovered depth head is used.
+        """
+        from PIL import Image as PILImage
+
+        patch_tokens, (orig_W, orig_H) = self.patch_features(image)
+        feat_inst = patch_tokens.index_select(1, self.instance_dims)
+        scores = (feat_inst @ self.instance_coef + self.instance_intercept).cpu().numpy()
+        score_grid = scores.reshape(self.patch_grid, self.patch_grid)
+
+        # K rescaled to input_res
+        sx, sy = self.input_res / orig_W, self.input_res / orig_H
+        K_scaled = K.astype(np.float64).copy()
+        K_scaled[0, 0] *= sx
+        K_scaled[0, 2] *= sx
+        K_scaled[1, 1] *= sy
+        K_scaled[1, 2] *= sy
+        K_inv = np.linalg.inv(K_scaled)
+
+        # Depth at input_res
+        if depth is None:
+            feat_depth = patch_tokens.index_select(1, self.depth_dims)
+            d_grid = (feat_depth @ self.depth_coef + self.depth_intercept).cpu().numpy()
+            d_grid = d_grid.reshape(self.patch_grid, self.patch_grid).astype(np.float32)
+            d_pil = PILImage.fromarray(d_grid, mode="F")
+            depth_full = np.array(
+                d_pil.resize((self.input_res, self.input_res), PILImage.BILINEAR)
+            )
+        else:
+            d_pil = PILImage.fromarray(depth.astype(np.float32), mode="F")
+            depth_full = np.array(
+                d_pil.resize((self.input_res, self.input_res), PILImage.BILINEAR)
+            )
+
+        # Cluster assignment per patch (nearest center).
+        diff = patch_tokens.unsqueeze(1) - self.cluster_centers.unsqueeze(0)
+        dist = (diff * diff).sum(dim=2)
+        cluster_assign = dist.argmin(dim=1).cpu().numpy()
+
+        # Foreground mask at full image resolution.
+        score_full_pil = PILImage.fromarray(score_grid.astype(np.float32), mode="F")
+        score_full = np.array(
+            score_full_pil.resize((self.input_res, self.input_res), PILImage.BILINEAR)
+        )
+        fg_full = score_full > self.instance_threshold
+
+        try:
+            from scipy.ndimage import label
+            labeled, n_comp = label(fg_full)
+        except ImportError:
+            labeled = fg_full.astype(np.int32)
+            n_comp = 1
+
+        boxes: List[Box3D] = []
+        stride = 4
+        for cid in range(1, n_comp + 1):
+            comp = labeled == cid
+            if comp.sum() < 100:
+                continue
+            ys_full, xs_full = np.where(comp)
+            ys = ys_full[::stride]
+            xs = xs_full[::stride]
+            if len(ys) < 20:
+                continue
+            d_arr = depth_full[ys, xs].astype(np.float64)
+            valid = (d_arr > 0.1) & (d_arr < 10.0) & np.isfinite(d_arr)
+            if valid.sum() < 20:
+                continue
+            ys_v, xs_v, d_v = ys[valid], xs[valid], d_arr[valid]
+            rays = K_inv @ np.stack([xs_v.astype(np.float64), ys_v.astype(np.float64),
+                                     np.ones_like(xs_v, dtype=np.float64)], axis=0)
+            pts = np.stack([rays[0] * d_v, rays[1] * d_v, d_v], axis=-1)
+
+            try:
+                from sklearn.cluster import DBSCAN
+                db = DBSCAN(eps=0.15, min_samples=15).fit(pts)
+                labels_pts = db.labels_
+            except ImportError:
+                labels_pts = np.zeros(len(pts), dtype=np.int32)
+
+            uniq = [u for u in set(labels_pts) if u >= 0] or [-1]
+            for u in uniq:
+                sel = labels_pts == u if u >= 0 else np.ones(len(pts), dtype=bool)
+                if sel.sum() < 20:
+                    continue
+                pts_sub = pts[sel]
+                obb = self._fit_obb(pts_sub)
+                if obb is None:
+                    continue
+                cx, cy, cz, w, h, d, theta = obb
+                if cz <= 0.2 or cz > 10.0:
+                    continue
+
+                # Per-cluster size prior blend.
+                comp_patch_lin = (ys_v[sel] // self.patch_size) * self.patch_grid + (xs_v[sel] // self.patch_size)
+                comp_patch_lin = comp_patch_lin.astype(np.int32)
+                clusters_here = cluster_assign[comp_patch_lin]
+                if len(clusters_here):
+                    mode = int(np.bincount(clusters_here).argmax())
+                    w_p, h_p, d_p = self.size_priors[mode].cpu().numpy()
+                    n = int(sel.sum())
+                    denom = n + self.prior_weight
+                    w = (n * w + self.prior_weight * float(w_p)) / denom
+                    h = (n * h + self.prior_weight * float(h_p)) / denom
+                    d = (n * d + self.prior_weight * float(d_p)) / denom
+
+                if max(w, h, d) > 2.5 or min(w, h, d) < 0.05 or w * h * d > 4.0:
+                    continue
+                boxes.append(Box3D(cx, cy, cz, w, h, d, theta, n_inliers=int(sel.sum())))
+        return boxes
+
+    @torch.no_grad()
+    def perceive(self, image, K: np.ndarray, depth: Optional[np.ndarray] = None) -> Dict:
+        score_map, _ = self.fg_score_map(image)
+        depth_map_pred, _ = self.depth_map(image)
+        boxes = self.detect(image, K, depth=depth)
+        return {
+            "fg_score_map": score_map,
+            "depth_map": depth_map_pred,
+            "boxes": boxes,
+        }
+
+    @staticmethod
+    def _fit_obb(pts: np.ndarray, trim_pct: float = 2.0):
+        if len(pts) < 8:
+            return None
+        xz = pts[:, [0, 2]]
+        center_xz = xz.mean(axis=0)
+        xz_centered = xz - center_xz
+        cov = xz_centered.T @ xz_centered / max(1, len(xz_centered) - 1)
+        eigvals, eigvecs = np.linalg.eigh(cov)
+        principal = eigvecs[:, np.argmax(eigvals)]
+        theta = math.atan2(-principal[1], principal[0])
+
+        c, s = math.cos(-theta), math.sin(-theta)
+        x_loc = c * pts[:, 0] + s * pts[:, 2]
+        z_loc = -s * pts[:, 0] + c * pts[:, 2]
+        y = pts[:, 1]
+        lo, hi = trim_pct, 100.0 - trim_pct
+        x_lo, x_hi = float(np.percentile(x_loc, lo)), float(np.percentile(x_loc, hi))
+        y_lo, y_hi = float(np.percentile(y, lo)), float(np.percentile(y, hi))
+        z_lo, z_hi = float(np.percentile(z_loc, lo)), float(np.percentile(z_loc, hi))
+        w = x_hi - x_lo
+        h = y_hi - y_lo
+        d = z_hi - z_lo
+        if w <= 0 or h <= 0 or d <= 0:
+            return None
+        cx_loc = (x_lo + x_hi) / 2.0
+        cz_loc = (z_lo + z_hi) / 2.0
+        cy = (y_lo + y_hi) / 2.0
+        c2, s2 = math.cos(theta), math.sin(theta)
+        cx = c2 * cx_loc + s2 * cz_loc
+        cz = -s2 * cx_loc + c2 * cz_loc
+        return (cx, cy, cz, w, h, d, theta)

config.json

@@ -0,0 +1,13 @@
+{
+  "input_res": 768,
+  "patch_grid": 48,
+  "patch_size": 16,
+  "feature_dim": 768,
+  "prior_weight": 80.0,
+  "n_clusters": 8,
+  "depth_range_m": [
+    0.1,
+    10.0
+  ],
+  "fg_threshold": 0.31326109170913696
+}

depth_head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0f9b589df8de62e4e4a707956e82143c0ec685d4f5a8feaf98934db7f8afebad
+size 9420

infer.py

@@ -0,0 +1,77 @@
+"""CLI for argus-3d.
+
+Usage:
+  python infer.py detect <image> --K K.json [--depth depth.png]
+  python infer.py perceive <image> --K K.json
+"""
+import argparse
+import json
+import sys
+
+import numpy as np
+from PIL import Image
+
+from argus_3d import Argus3D
+
+
+def _load_K(path: str) -> np.ndarray:
+    with open(path, "r") as f:
+        return np.array(json.load(f), dtype=np.float64)
+
+
+def _load_depth(path: str) -> np.ndarray:
+    arr = np.array(Image.open(path))
+    if arr.dtype == np.uint16:
+        return arr.astype(np.float32) / 1000.0
+    return arr.astype(np.float32)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    sub = ap.add_subparsers(dest="cmd", required=True)
+    p_det = sub.add_parser("detect")
+    p_det.add_argument("image")
+    p_det.add_argument("--K", required=True)
+    p_det.add_argument("--depth", default=None)
+    p_det.add_argument("--repo", default="phanerozoic/argus-3d")
+    p_det.add_argument("--device", default="cuda")
+    p_per = sub.add_parser("perceive")
+    p_per.add_argument("image")
+    p_per.add_argument("--K", required=True)
+    p_per.add_argument("--depth", default=None)
+    p_per.add_argument("--repo", default="phanerozoic/argus-3d")
+    p_per.add_argument("--device", default="cuda")
+    args = ap.parse_args()
+
+    model = Argus3D.from_pretrained(args.repo, device=args.device)
+    K = _load_K(args.K)
+    depth = _load_depth(args.depth) if args.depth else None
+
+    if args.cmd == "detect":
+        boxes = model.detect(args.image, K, depth=depth)
+        out = [
+            {
+                "box": [b.cx, b.cy, b.cz, b.w, b.h, b.d, b.theta],
+                "score": b.score,
+                "n_inliers": b.n_inliers,
+            }
+            for b in boxes
+        ]
+        json.dump(out, sys.stdout, indent=2)
+    else:
+        out = model.perceive(args.image, K, depth=depth)
+        out["fg_score_map"] = out["fg_score_map"].tolist()
+        out["depth_map"] = out["depth_map"].tolist()
+        out["boxes"] = [
+            {
+                "box": [b.cx, b.cy, b.cz, b.w, b.h, b.d, b.theta],
+                "score": b.score,
+                "n_inliers": b.n_inliers,
+            }
+            for b in out["boxes"]
+        ]
+        json.dump(out, sys.stdout, indent=2)
+
+
+if __name__ == "__main__":
+    main()

instance_head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f71ae2a496a02a1727355073a43946d8eae0ad2ff18d3668955b8b2cd5164ac3
+size 9496

size_priors.safetensors

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