Spaces:
Running on Zero
Running on Zero
| """ | |
| Multi-scale (N-channel) power-transform depth encoding, designed -- like | |
| `multiscale_disparity_torch.py` -- to reuse pretrained RGB VAEs without wasting | |
| channel capacity. This is the power-transform counterpart of the multi-scale | |
| disparity codec. | |
| With three channels and the recommended preset (lam = -3, c = (2, 6, 18)): | |
| y_R = 1 - 2 / (1 + d / (|lam| * c_1))^2 c_1 = 2 (near) | |
| y_G = 1 - 2 / (1 + d / (|lam| * c_2))^2 c_2 = 6 (mid) | |
| y_B = 1 - 2 / (1 + d / (|lam| * c_3))^2 c_3 = 18 (far) | |
| Each channel is an independent power-transform bijection at a different | |
| scale. Decoding inverts each channel and blends the three depth estimates, | |
| w_i = (1 - y_i^2)^2, d_hat = sum_i w_i * d_i / sum_i w_i. | |
| Why (1 - y_i^2)^2? | |
| ------------------ | |
| For the disparity codec, this weight happens to be exactly proportional to | |
| d^2 / Var(d_hat), which minimises AbsRel. For the power transform that | |
| identity is no longer exact -- the AbsRel-optimal weight is | |
| w_i^opt proportional to (1 - y_i)^2 * (sqrt(2) - sqrt(1 - y_i))^2. | |
| In our 5-dataset VAE-noise simulation (sigma_y in [0.01, 0.05]) the two weights | |
| were within 0.5%, so we keep the simpler (1 - y^2)^2 as the default and expose | |
| the optimal form as an opt-in (`weight = "opt"`). | |
| Why c = (2, 6, 18) ? | |
| -------------------- | |
| Power transform's tail decays as d^-3 (vs d^-2 for disparity), so each channel | |
| covers a *narrower* band of depth. Empirically, ratio r = 3 between adjacent | |
| c's outperforms r = 4 (which we use for disparity) by ~5% AbsRel; r = 2 over- | |
| laps too much, r = 4 leaves precision dips between channels. | |
| In our 5-dataset experiment (NYU, ETH3D, DIODE, Sintel, KITTI), this preset | |
| gives ~10% lower AbsRel than 3-channel disparity (2, 8, 32) under the same | |
| VAE noise level -- at the cost of a slightly less-clean optimal weight. | |
| """ | |
| from __future__ import annotations | |
| from typing import Literal, Optional, Sequence, Tuple, Union | |
| import torch | |
| from torch import Tensor | |
| from .power_transform import depth_to_signed_unit, signed_unit_to_depth | |
| # --------------------------------------------------------------------------- | |
| # Default preset. Picked from VAE-noise sweeps over NYU + ETH3D + DIODE + | |
| # Sintel + KITTI; see the docstring for the rationale. | |
| # --------------------------------------------------------------------------- | |
| DEFAULT_LAM: float = -3.0 | |
| DEFAULT_C_RGB: Tuple[float, float, float] = (2.0, 6.0, 18.0) | |
| WeightKind = Literal["iyy2", "opt"] | |
| # --------------------------------------------------------------------------- | |
| # Functional API | |
| # --------------------------------------------------------------------------- | |
| def multiscale_power_encode( | |
| d: Tensor, # shape (..., 1, H, W) or (..., H, W) | |
| c_values: Sequence[float] = DEFAULT_C_RGB, | |
| lam: float = DEFAULT_LAM, | |
| 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 power-transform tensor. | |
| Args: | |
| d : depth tensor with a singleton channel dim (or no channel | |
| dim); any leading batch dims are fine. | |
| c_values : per-channel scale parameters. len == output channels. | |
| lam : shared shape parameter for all channels; must be < -1. | |
| mask : optional bool mask, broadcastable to `d`. | |
| invalid_fill : value at invalid pixels (in [-1, 1) space). | |
| return_mask : also return the effective valid mask. | |
| channel_dim : where to stack the N channels (default -3, so a (..., 1, | |
| H, W) input becomes (..., N, H, W)). | |
| Returns: | |
| y : same as `d` but with N channels on `channel_dim`, in [-1, 1)^N. | |
| """ | |
| if lam >= -1.0: | |
| raise ValueError(f"lam must be < -1, got {lam}") | |
| if any(c <= 0 for c in c_values): | |
| raise ValueError(f"all c must be > 0, got {tuple(c_values)}") | |
| # Collapse a singleton channel dim on d, if present, so per-channel | |
| # encoders see a clean (..., H, W) tensor. | |
| if d.dim() >= -channel_dim and d.shape[channel_dim] == 1: | |
| d_bare = d.squeeze(channel_dim) | |
| if mask is not None and mask.shape == d.shape: | |
| mask_bare: Optional[Tensor] = mask.squeeze(channel_dim) | |
| else: | |
| mask_bare = mask | |
| else: | |
| d_bare = d | |
| mask_bare = mask | |
| channels = [] | |
| valid_per_ch = [] | |
| for c in c_values: | |
| y_i, v_i = depth_to_signed_unit( | |
| d_bare, lam=lam, c=float(c), | |
| mask=mask_bare, invalid_fill=invalid_fill, return_mask=True, | |
| ) | |
| channels.append(y_i) | |
| valid_per_ch.append(v_i) | |
| y = torch.stack(channels, dim=channel_dim) | |
| valid = torch.stack(valid_per_ch, dim=channel_dim) | |
| if return_mask: | |
| return y, valid | |
| return y | |
| def multiscale_power_decode( | |
| y: Tensor, # shape (..., N, H, W) | |
| c_values: Sequence[float] = DEFAULT_C_RGB, | |
| lam: float = DEFAULT_LAM, | |
| inf_cap: Optional[float] = 1e4, | |
| eps_boundary: float = 1e-4, | |
| weight: WeightKind = "iyy2", | |
| channel_dim: int = -3, | |
| ) -> Tensor: | |
| """Decode an N-channel power-transform tensor back to single-channel depth. | |
| Each channel is inverted independently using `signed_unit_to_depth` | |
| and the final estimate is the confidence-weighted average, | |
| d_hat = sum_i w_i * d_i / sum_i w_i | |
| Args: | |
| y : encoded tensor, shape (..., N, H, W). | |
| c_values : per-channel scale parameters. | |
| lam : shared shape parameter, must match the encoder. | |
| inf_cap : clamp on the final depth (None disables). | |
| eps_boundary : keep |y| <= 1 - eps so the inverse stays finite. | |
| weight : "iyy2" -> w = (1 - y^2)^2 (cheap, near-optimal) | |
| "opt" -> AbsRel-optimal weight derived for lam = -3. | |
| channel_dim : channel axis (default -3). | |
| """ | |
| 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" | |
| ) | |
| y_c = y.clamp(min=-1.0 + eps_boundary, max=1.0 - eps_boundary) | |
| # Per-channel depth estimates, one slice at a time so each call to | |
| # `signed_unit_to_depth` sees the right `c` (lam is shared). | |
| chans = [] | |
| for i, c in enumerate(c_values): | |
| y_i = y_c.select(channel_dim, i) | |
| d_i = signed_unit_to_depth(y_i, lam=lam, c=float(c), inf_cap=None) | |
| chans.append(d_i) | |
| d_per_ch = torch.stack(chans, dim=channel_dim) # (..., N, H, W) | |
| # --- confidence weights -------------------------------------------------- | |
| if weight == "iyy2": | |
| w = (1.0 - y_c.pow(2)).pow(2) | |
| elif weight == "opt": | |
| # w_i^opt proportional to (1 - y)^2 * (sqrt(2) - sqrt(1 - y))^2 | |
| # Derived from w proportional to d^2 / Var(d_hat) with lam = -3. | |
| if lam != -3.0: | |
| raise NotImplementedError( | |
| "weight='opt' is currently only derived for lam = -3" | |
| ) | |
| one_my = (1.0 - y_c).clamp(min=eps_boundary) | |
| w = one_my.pow(2) * (2.0 ** 0.5 - one_my.sqrt()).pow(2) | |
| else: | |
| raise ValueError(f"unknown weight kind {weight!r}") | |
| w = w.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 DepthMultiScaleDisparityTransform) | |
| # --------------------------------------------------------------------------- | |
| class DepthMultiScalePowerTransform: | |
| """N-channel power-transform codec. Plain container, not an nn.Module. | |
| Example | |
| ------- | |
| >>> codec = DepthMultiScalePowerTransform() # lam=-3, c=(2,6,18) | |
| >>> y = codec.encode(depth, mask=mask) # (..., 3, H, W) | |
| >>> d_hat = codec.decode(y_pred) # (..., H, W) | |
| """ | |
| def __init__( | |
| self, | |
| c: Sequence[float] = DEFAULT_C_RGB, | |
| lam: float = DEFAULT_LAM, | |
| inf_cap: Optional[float] = 1e4, | |
| invalid_fill: float = 0.0, | |
| weight: WeightKind = "iyy2", | |
| channel_dim: int = -3, | |
| ): | |
| if lam >= -1.0: | |
| raise ValueError(f"lam must be < -1, got {lam}") | |
| self.c_values = tuple(float(_c) for _c in c) | |
| if any(_c <= 0 for _c in self.c_values): | |
| raise ValueError(f"all c must be > 0, got {self.c_values}") | |
| self.lam = float(lam) | |
| self.inf_cap = inf_cap | |
| self.invalid_fill = float(invalid_fill) | |
| self.weight: WeightKind = weight | |
| 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_power_encode( | |
| d, c_values=self.c_values, lam=self.lam, mask=mask, | |
| invalid_fill=self.invalid_fill, return_mask=return_mask, | |
| channel_dim=self.channel_dim, | |
| ) | |
| def decode(self, y): | |
| return multiscale_power_decode( | |
| y, c_values=self.c_values, lam=self.lam, | |
| inf_cap=self.inf_cap, weight=self.weight, | |
| 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"DepthMultiScalePowerTransform(" | |
| f"lam={self.lam:.4f}, c={self.c_values}, " | |
| f"weight={self.weight!r}, inf_cap={self.inf_cap}, " | |
| f"invalid_fill={self.invalid_fill})" | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # Self-test: round-trip + VAE-noise comparison vs single-channel and disparity | |
| # --------------------------------------------------------------------------- | |
| if __name__ == "__main__": | |
| import math | |
| torch.manual_seed(0) | |
| codec = DepthMultiScalePowerTransform() | |
| print(codec) | |
| # Round-trip sanity (no noise) -- weighted blend converges to true depth | |
| # because each channel's inverse is exact. | |
| 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)}") | |
| print(f"y range per channel : " | |
| f"R [{y[:, 0].min():+.3f}, {y[:, 0].max():+.3f}] " | |
| f"G [{y[:, 1].min():+.3f}, {y[:, 1].max():+.3f}] " | |
| f"B [{y[:, 2].min():+.3f}, {y[:, 2].max():+.3f}]") | |
| # ---- VAE-noise comparison vs single-channel power & 3-ch disparity ------ | |
| def lognormal(mean_m, sigma, n=200_000, cap=None): | |
| x = (torch.randn(n) * sigma + math.log(mean_m)).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), | |
| } | |
| # 3-ch disparity (c = 2, 8, 32) for an apples-to-apples baseline. | |
| from multiscale_disparity_torch import DepthMultiScaleDisparityTransform | |
| disp_codec = DepthMultiScaleDisparityTransform() # (2, 8, 32) | |
| pwr_codec = DepthMultiScalePowerTransform() # (2, 6, 18) | |
| pwr_codec_o = DepthMultiScalePowerTransform(weight="opt") # AbsRel-optimal | |
| print("\nVAE-noise AbsRel (sigma_y = 0.01):") | |
| print(f"{'dataset':<8} {'3-ch disp':>10} {'3-ch pwr (1-y2)2':>18} " | |
| f"{'3-ch pwr opt':>14} {'pwr vs disp':>12}") | |
| print("-" * 72) | |
| sigma = 0.01 | |
| for ds_name, d_arr in datasets.items(): | |
| d_arr = d_arr.reshape(1, 1, -1, 1) | |
| y_d = disp_codec.encode(d_arr) + torch.randn( | |
| 1, 3, d_arr.shape[-2], 1) * sigma | |
| ar_d = ((disp_codec.decode(y_d) - d_arr.squeeze(1)).abs() | |
| / d_arr.squeeze(1)).mean().item() | |
| y_p = pwr_codec.encode(d_arr) + torch.randn( | |
| 1, 3, d_arr.shape[-2], 1) * sigma | |
| ar_p = ((pwr_codec.decode(y_p) - d_arr.squeeze(1)).abs() | |
| / d_arr.squeeze(1)).mean().item() | |
| y_po = pwr_codec_o.encode(d_arr) + torch.randn( | |
| 1, 3, d_arr.shape[-2], 1) * sigma | |
| ar_po = ((pwr_codec_o.decode(y_po) - d_arr.squeeze(1)).abs() | |
| / d_arr.squeeze(1)).mean().item() | |
| imp = (ar_d - ar_p) / ar_d * 100.0 | |
| print(f"{ds_name:<8} {ar_d:>10.4f} {ar_p:>18.4f} {ar_po:>14.4f} " | |
| f"{imp:>11.1f}%") | |
| # ---- Per-channel inspection at one specific depth ------------------------ | |
| d_pt = torch.tensor([[[1.0, 6.0, 30.0, 100.0]]]) # (1, 1, 4) | |
| y_pt = pwr_codec.encode(d_pt.unsqueeze(-1)) # (1, 3, 4, 1) | |
| print("\nPer-channel y at d = [1, 6, 30, 100] m (c = (2, 6, 18)):") | |
| for i, c in enumerate(pwr_codec.c_values): | |
| row = y_pt[0, i, :, 0].tolist() | |
| print(f" ch{i} (c={c:>4.1f}) y = " | |
| + " ".join(f"{v:+.3f}" for v in row)) | |