Spaces:
Running on Zero
Running on Zero
| """ | |
| Multi-scale (N-channel) disparity depth encoding, designed to reuse pretrained | |
| RGB VAEs (SD-VAE, Flux-VAE, etc.) without wasting channel capacity. | |
| With three channels: | |
| y_R = (d - c_1) / (d + c_1) c_1 = 2 (near) | |
| y_G = (d - c_2) / (d + c_2) c_2 = 8 (mid) | |
| y_B = (d - c_3) / (d + c_3) c_3 = 32 (far) | |
| Each channel is an independent disparity bijection at a different pivot. | |
| Their resolution curves have "knees" at c_1, c_2, c_3, respectively, and | |
| together give near-flat high-precision coverage across ~2 decades of depth. | |
| Decoding uses each channel's own inverse and blends the three depth estimates | |
| weighted by w_i = (1 - y_i^2)^2 (saturation-aware confidence). | |
| """ | |
| from __future__ import annotations | |
| from typing import Optional, Sequence, Tuple, Union | |
| import torch | |
| from torch import Tensor | |
| # --------------------------------------------------------------------------- | |
| # Default preset for NYU + ETH3D + DIODE + Sintel + KITTI | |
| # --------------------------------------------------------------------------- | |
| DEFAULT_C_RGB: Tuple[float, float, float] = (2.0, 8.0, 32.0) | |
| # --------------------------------------------------------------------------- | |
| # Functional API | |
| # --------------------------------------------------------------------------- | |
| def multiscale_disparity_encode( | |
| d: Tensor, # shape (..., 1, H, W) or (..., H, W) | |
| c_values: Sequence[float] = DEFAULT_C_RGB, | |
| mask: Optional[Tensor] = None, | |
| invalid_fill: float = 0.0, | |
| return_mask: bool = False, | |
| channel_dim: int = -3, | |
| ) -> Union[Tensor, Tuple[Tensor, Tensor]]: | |
| """Encode depth to an N-channel disparity tensor. | |
| Args: | |
| d : depth tensor with a singleton channel dim (or no channel | |
| dim); any leading batch dims are fine. | |
| c_values : per-channel pivot depths. len == output channels. | |
| mask : optional bool mask, same shape as `d`. | |
| invalid_fill : value at invalid pixels. | |
| return_mask : also return the effective valid mask. | |
| channel_dim : where to stack the N channels (default -3, i.e. | |
| producing (..., N, H, W) from (..., 1, H, W)). | |
| Returns: | |
| y : same as `d` but with N channels on `channel_dim`. | |
| """ | |
| # Collapse a singleton channel dim on d, if present, so broadcasting works. | |
| if d.dim() >= -channel_dim and d.shape[channel_dim] == 1: | |
| d_bare = d.squeeze(channel_dim) | |
| else: | |
| d_bare = d | |
| channels = [] | |
| for c in c_values: | |
| y = (d_bare - c) / (d_bare + c) | |
| channels.append(y) | |
| y = torch.stack(channels, dim=channel_dim) | |
| d_valid = (d_bare > 0).unsqueeze(channel_dim).expand_as(y) | |
| if mask is not None: | |
| if mask.shape == d.shape and d.shape[channel_dim] == 1: | |
| mask_bare = mask.squeeze(channel_dim) | |
| else: | |
| mask_bare = mask | |
| mask_bc = mask_bare.to(torch.bool).unsqueeze(channel_dim).expand_as(y) | |
| valid = mask_bc & d_valid | |
| else: | |
| valid = d_valid | |
| y = torch.where(valid, y, torch.full_like(y, invalid_fill)) | |
| if return_mask: | |
| return y, valid | |
| return y | |
| def multiscale_disparity_decode( | |
| y: Tensor, # shape (..., N, H, W) | |
| c_values: Sequence[float] = DEFAULT_C_RGB, | |
| inf_cap: Optional[float] = 1e4, | |
| eps_boundary: float = 1e-4, | |
| channel_dim: int = -3, | |
| ) -> Tensor: | |
| """Decode an N-channel disparity tensor back to a single-channel depth. | |
| Each channel is inverted independently; d_i = c_i * (1 + y_i) / (1 - y_i). | |
| The final estimate is the confidence-weighted average, | |
| w_i = (1 - y_i^2)^2, d_hat = sum_i w_i d_i / sum_i w_i | |
| which assigns near-zero weight to any saturated channel and full weight to | |
| whichever channel is most sensitive at the true depth. | |
| """ | |
| if y.shape[channel_dim] != len(c_values): | |
| raise ValueError( | |
| f"y has {y.shape[channel_dim]} channels but " | |
| f"c_values has {len(c_values)} entries" | |
| ) | |
| c_t = torch.as_tensor(c_values, dtype=y.dtype, device=y.device) | |
| view = [1] * y.dim() | |
| view[channel_dim] = len(c_values) | |
| c_t = c_t.view(view) | |
| y_c = y.clamp(min=-1.0 + eps_boundary, max=1.0 - eps_boundary) | |
| d_per_ch = c_t * (1.0 + y_c) / (1.0 - y_c) # (..., N, H, W) | |
| # Confidence: (1 - y^2)^2 -- peaks at y=0, vanishes at y=+-1. | |
| w = (1.0 - y_c * y_c).pow(2).clamp(min=1e-12) | |
| d_hat = (w * d_per_ch).sum(dim=channel_dim) / w.sum(dim=channel_dim) | |
| if inf_cap is not None: | |
| d_hat = d_hat.clamp(max=inf_cap) | |
| return d_hat | |
| # --------------------------------------------------------------------------- | |
| # Container class (parallels DepthDisparityTransform) | |
| # --------------------------------------------------------------------------- | |
| class DepthMultiScaleDisparityTransform: | |
| """N-channel disparity codec. Plain container, not an nn.Module. | |
| Example | |
| ------- | |
| >>> codec = DepthMultiScaleDisparityTransform() # c=(2, 8, 32) | |
| >>> y = codec.encode(depth, mask=mask) # (..., 3, H, W) | |
| >>> d_hat = codec.decode(y_pred) # (..., H, W) | |
| """ | |
| def __init__( | |
| self, | |
| c_values: Sequence[float] = DEFAULT_C_RGB, | |
| inf_cap: Optional[float] = 1e4, | |
| invalid_fill: float = 0.0, | |
| channel_dim: int = -3, | |
| ): | |
| self.c_values = tuple(float(c) for c in c_values) | |
| if any(c <= 0 for c in self.c_values): | |
| raise ValueError(f"all c must be > 0, got {self.c_values}") | |
| self.inf_cap = inf_cap | |
| self.invalid_fill = float(invalid_fill) | |
| self.channel_dim = channel_dim | |
| def num_channels(self) -> int: | |
| return len(self.c_values) | |
| def encode(self, d, mask=None, return_mask=False): | |
| return multiscale_disparity_encode( | |
| d, c_values=self.c_values, mask=mask, | |
| invalid_fill=self.invalid_fill, return_mask=return_mask, | |
| channel_dim=self.channel_dim, | |
| ) | |
| def decode(self, y): | |
| return multiscale_disparity_decode( | |
| y, c_values=self.c_values, | |
| inf_cap=self.inf_cap, channel_dim=self.channel_dim, | |
| ) | |
| def __call__(self, d, valid_mask=None, return_mask=False): | |
| return self.encode(d, mask=valid_mask, return_mask=return_mask) | |
| def __repr__(self): | |
| return (f"DepthMultiScaleDisparityTransform(c={self.c_values}, " | |
| f"inf_cap={self.inf_cap}, invalid_fill={self.invalid_fill})") | |
| # --------------------------------------------------------------------------- | |
| # Self-test: round-trip + VAE-noise accuracy vs single-channel | |
| # --------------------------------------------------------------------------- | |
| if __name__ == "__main__": | |
| import math | |
| torch.manual_seed(0) | |
| codec = DepthMultiScaleDisparityTransform() | |
| print(codec) | |
| # Round-trip sanity | |
| d = torch.rand(2, 1, 64, 64, dtype=torch.float64) * 200.0 | |
| y = codec(d) | |
| d_hat = codec.decode(y) | |
| print(f"round-trip max err : {(d.squeeze(1) - d_hat).abs().max():.2e} m") | |
| print(f"y shape : {tuple(y.shape)}") | |
| # Simulate VAE noise + compare to single-channel disparity (c=5) | |
| def lognormal(mean_m, sigma, n=200_000, cap=None): | |
| x = torch.randn(n) * sigma + math.log(mean_m) | |
| x = x.exp() | |
| return x if cap is None else x.clamp(0.1, cap) | |
| datasets = { | |
| "NYU": lognormal(2.5, 0.5, cap=10.0), | |
| "ETH3D": lognormal(5.0, 0.8, cap=70.0), | |
| "DIODE": lognormal(8.0, 1.2, cap=350.0), | |
| "Sintel": lognormal(10.0, 0.6, cap=100.0), | |
| "KITTI": lognormal(15.0, 0.7, cap=80.0), | |
| } | |
| # Single-channel disparity (c=5), for comparison | |
| def enc_single(d, c=5.0): | |
| return (d - c) / (d + c) | |
| def dec_single(y, c=5.0, eps=1e-4): | |
| y = y.clamp(-1 + eps, 1 - eps) | |
| return c * (1 + y) / (1 - y) | |
| print("\nVAE-noise comparison (sigma_y = 0.01):") | |
| print(f"{'dataset':<8} {'single-ch AbsRel':>18} {'3-ch AbsRel':>14} {'improvement':>12}") | |
| print("-" * 60) | |
| sigma = 0.01 | |
| for ds_name, d in datasets.items(): | |
| d = d.reshape(1, 1, -1, 1) | |
| # 1 channel | |
| y_s = enc_single(d.squeeze(1)) | |
| y_s = y_s + torch.randn_like(y_s) * sigma | |
| d_hat_s = dec_single(y_s) | |
| abs_rel_s = ((d_hat_s - d.squeeze(1)).abs() / d.squeeze(1)).mean().item() | |
| # 3 channels (independent noise per channel) | |
| y_m = codec.encode(d) | |
| y_m = y_m + torch.randn_like(y_m) * sigma | |
| d_hat_m = codec.decode(y_m) | |
| abs_rel_m = ((d_hat_m - d.squeeze(1)).abs() / d.squeeze(1)).mean().item() | |
| imp = (abs_rel_s - abs_rel_m) / abs_rel_s * 100.0 | |
| print(f"{ds_name:<8} {abs_rel_s:>18.4f} {abs_rel_m:>14.4f} {imp:>11.1f}%") | |