Add novel heads: Optimal Transport (det), Info Bottleneck (seg), Harmonic Depth (depth)

heads/__init__.py

@@ -7,6 +7,7 @@ from .wavelet.head import Wavelet
 from .log_linear.head import LogLinear
 from .ordinal_regression.head import OrdinalRegression
 from .multiscale_gradient.head import MultiscaleGradient
+from .harmonic.head import HarmonicDepth
 
 REGISTRY = {
     "linear_probe": LinearProbe,
@@ -16,6 +17,7 @@ REGISTRY = {
     "log_linear": LogLinear,
     "ordinal_regression": OrdinalRegression,
     "multiscale_gradient": MultiscaleGradient,
+    "harmonic": HarmonicDepth,
 }
 
 ALL_NAMES = list(REGISTRY.keys())

heads/harmonic/__init__.py

@@ -0,0 +1 @@
+from .head import HarmonicDepth

heads/harmonic/head.py

@@ -0,0 +1,74 @@
+"""Harmonic Depth: cofiber edge detection + boundary depth prediction + PDE solve.
+
+770 parameters. Depth at edge locations is predicted by a single linear layer.
+Depth at non-edge locations is solved via the discrete Laplace equation
+(iterative Jacobi relaxation). The cofiber energy identifies edges — locations
+where the feature changes across scales, which correspond to depth discontinuities.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def cofiber_energy(spatial):
+    """Compute per-location cofiber energy (L2 norm of high-freq residual)."""
+    omega = F.avg_pool2d(spatial, 2)
+    sigma_omega = F.interpolate(omega, size=spatial.shape[2:], mode="bilinear", align_corners=False)
+    cofib = spatial - sigma_omega
+    return cofib.pow(2).sum(dim=1, keepdim=True)
+
+
+def jacobi_solve(depth, mask, n_iters=50):
+    """Solve discrete Laplace equation at non-masked locations.
+    At mask=1 (edge) locations, depth is fixed. At mask=0 locations,
+    depth = average of 4 neighbors. Iterated to convergence."""
+    for _ in range(n_iters):
+        # Pad with replicate boundary conditions
+        padded = F.pad(depth, (1, 1, 1, 1), mode="replicate")
+        avg = (padded[:, :, :-2, 1:-1] + padded[:, :, 2:, 1:-1] +
+               padded[:, :, 1:-1, :-2] + padded[:, :, 1:-1, 2:]) / 4
+        # Update only non-edge locations
+        depth = mask * depth + (1 - mask) * avg
+    return depth
+
+
+class HarmonicDepth(nn.Module):
+    """770 parameters. Depth solved as a harmonic function with learned boundary values."""
+    name = "harmonic"
+    needs_intermediates = False
+
+    def __init__(self, feat_dim=768, min_depth=0.001, max_depth=10.0,
+                 edge_percentile=0.2, n_iters=50):
+        super().__init__()
+        self.min_depth = min_depth
+        self.max_depth = max_depth
+        self.edge_percentile = edge_percentile
+        self.n_iters = n_iters
+        # Boundary depth predictor: single linear layer at edge locations
+        self.depth_proj = nn.Conv2d(feat_dim, 1, 1)
+
+    def forward(self, spatial, inter=None):
+        B, C, H, W = spatial.shape
+
+        # Cofiber energy identifies edges
+        energy = cofiber_energy(spatial)
+
+        # Threshold: top edge_percentile of locations are edges
+        flat_energy = energy.reshape(B, -1)
+        k = max(1, int(H * W * self.edge_percentile))
+        threshold = flat_energy.topk(k, dim=1).values[:, -1:].reshape(B, 1, 1, 1)
+        edge_mask = (energy >= threshold).float()
+
+        # Predict depth at edge locations
+        boundary_depth = self.depth_proj(spatial)
+        boundary_depth = boundary_depth.clamp(self.min_depth, self.max_depth)
+
+        # Initialize: boundary values at edges, mean depth elsewhere
+        mean_depth = (boundary_depth * edge_mask).sum(dim=(2, 3), keepdim=True) / edge_mask.sum(dim=(2, 3), keepdim=True).clamp(min=1)
+        depth = edge_mask * boundary_depth + (1 - edge_mask) * mean_depth
+
+        # Solve Laplace equation at non-edge locations
+        depth = jacobi_solve(depth, edge_mask, self.n_iters)
+
+        return depth.clamp(self.min_depth, self.max_depth)