Spaces:
Running on Zero
Running on Zero
| """ | |
| Disparity (Cayley / Möbius) depth encoding: | |
| y = (d - c) / (d + c) = tanh( (1/2) * log(d / c) ) | |
| Bijection [0, inf) -> [-1, 1) with ONE parameter `c`: | |
| * c is the "pivot" depth that maps to y = 0. | |
| * d = 0 -> y = -1 | |
| * d = c -> y = 0 | |
| * d -> inf -> y -> +1 | |
| Inverse: | |
| d = c * (1 + y) / (1 - y) | |
| Why this encoding for multi-dataset training | |
| -------------------------------------------- | |
| The L1 distance in y-space factors as | |
| |y_p - y_g| = 2c * |d_p - d_g| / ((c + d_p)(c + d_g)) | |
| which interpolates smoothly between | |
| * |d_p - d_g| / c (near field, d << c) | |
| * 2c * |1/d_p - 1/d_g| (far field, d >> c) | |
| i.e. the effective metric-depth loss is flat in close-up and transitions to | |
| the `1/depth` loss used by DepthPro / modern disparity-based depth learners | |
| at far range. That prior ("near depths deserve more precision") is baked | |
| into the encoding itself, with no loss surgery required. | |
| This is a special case of the log-logistic transform (`beta = 1`). It is | |
| called out as its own module because for multi-dataset (NYU + KITTI + ETH3D | |
| + DIODE + Sintel) training, `beta = 1` is the principled choice that aligns | |
| the L1-on-y loss with the 1/depth prior. | |
| See also | |
| -------- | |
| loglogistic_transform_torch.py : generalisation with a tunable `beta`. | |
| power_transform_torch.py : Barron's power transform family. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| from typing import Iterable, Mapping, Optional, Tuple, Union | |
| import torch | |
| from torch import Tensor | |
| # --------------------------------------------------------------------------- | |
| # Functional API | |
| # --------------------------------------------------------------------------- | |
| def disparity_encode( | |
| d: Tensor, | |
| c: float = 5.0, | |
| mask: Optional[Tensor] = None, | |
| invalid_fill: float = 0.0, | |
| return_mask: bool = False, | |
| ) -> Union[Tensor, Tuple[Tensor, Tensor]]: | |
| """Encode metric depth `d` in `[0, inf)` to `y` in `[-1, 1)`. | |
| y = (d - c) / (d + c) | |
| Computed in the `(d - c) / (d + c)` algebraic form directly (stable for | |
| any `d >= 0`, `c > 0`). | |
| Args: | |
| d : depth tensor, any shape (e.g. (1, H, W), (B, 1, H, W)). | |
| c : pivot depth (the d that maps to y = 0). | |
| mask : optional bool tensor broadcastable to `d`. Pixels where | |
| `mask == False` or `d <= 0` are written as | |
| `invalid_fill`. If `None`, just `d > 0` is required. | |
| invalid_fill : value at invalid pixels. | |
| return_mask : if True, also return the effective valid-mask. | |
| """ | |
| y = (d - c) / (d + c) | |
| d_valid = d > 0 | |
| valid = d_valid if mask is None else (mask.to(torch.bool) & d_valid) | |
| y = torch.where(valid, y, torch.full_like(y, invalid_fill)) | |
| if return_mask: | |
| return y, valid | |
| return y | |
| def disparity_decode( | |
| y: Tensor, | |
| c: float = 5.0, | |
| inf_cap: Optional[float] = 1e4, | |
| ) -> Tensor: | |
| """Decode `y` in `[-1, 1)` back to metric depth `d` in `[0, inf)`. | |
| d = c * (1 + y) / (1 - y) | |
| `inf_cap` clips the output at a finite maximum (defends against | |
| `y -> 1` where d would blow up). | |
| """ | |
| eps = torch.finfo(y.dtype).eps | |
| y_c = y.clamp(min=-1.0 + eps, max=1.0 - eps) | |
| d = c * (1.0 + y_c) / (1.0 - y_c) | |
| if inf_cap is not None: | |
| d = d.clamp(max=inf_cap) | |
| return d | |
| # --------------------------------------------------------------------------- | |
| # Parameter fitting | |
| # --------------------------------------------------------------------------- | |
| def fit_c_geometric_mean( | |
| depth_samples_by_dataset: Mapping[str, Tensor], | |
| masks_by_dataset: Optional[Mapping[str, Tensor]] = None, | |
| min_depth: float = 1e-3, | |
| ) -> float: | |
| """Pick `c` so it matches the geometric mean of the per-dataset medians. | |
| Rationale: Cayley is log-symmetric around `c`, so for cross-dataset | |
| training we want `c` to sit roughly at the log-center of all dataset | |
| medians. The geometric mean of the medians is exactly that. | |
| """ | |
| log_medians = [] | |
| for name, d in depth_samples_by_dataset.items(): | |
| m = None | |
| if masks_by_dataset is not None and name in masks_by_dataset: | |
| m = masks_by_dataset[name] | |
| if m is not None: | |
| d = d[m.to(torch.bool) & (d > min_depth)] | |
| else: | |
| d = d[d > min_depth] | |
| if d.numel() == 0: | |
| continue | |
| log_medians.append(math.log(float(d.to(torch.float64).median()))) | |
| if not log_medians: | |
| raise ValueError("No valid depth samples.") | |
| return math.exp(sum(log_medians) / len(log_medians)) | |
| # --------------------------------------------------------------------------- | |
| # Convenience class | |
| # --------------------------------------------------------------------------- | |
| class DepthDisparityTransform: | |
| """Plain container for the Cayley / disparity depth codec. | |
| Example | |
| ------- | |
| >>> codec = DepthDisparityTransform() # c=5 (NYU+KITTI+...) | |
| >>> y = codec.encode(depth, mask=mask) # [0, inf) -> [-1, 1) | |
| >>> d_hat = codec.decode(y_pred) # [-1, 1) -> [0, inf) | |
| Ready-made multi-dataset preset: | |
| >>> codec = DepthDisparityTransform.for_multidataset() # c = 5.0 | |
| """ | |
| def __init__( | |
| self, | |
| c: float = 5.0, | |
| inf_cap: Optional[float] = 1e4, | |
| invalid_fill: float = 0.0, | |
| ): | |
| if c <= 0.0: | |
| raise ValueError(f"c must be > 0, got {c}") | |
| self.c = float(c) | |
| self.inf_cap = inf_cap | |
| self.invalid_fill = float(invalid_fill) | |
| def encode( | |
| self, | |
| d: Tensor, | |
| mask: Optional[Tensor] = None, | |
| return_mask: bool = False, | |
| ) -> Union[Tensor, Tuple[Tensor, Tensor]]: | |
| return disparity_encode( | |
| d, c=self.c, mask=mask, | |
| invalid_fill=self.invalid_fill, return_mask=return_mask, | |
| ) | |
| def decode(self, y: Tensor) -> Tensor: | |
| return disparity_decode(y, c=self.c, inf_cap=self.inf_cap) | |
| def __call__( | |
| self, | |
| d: Tensor, | |
| valid_mask: Optional[Tensor] = None, | |
| return_mask: bool = False, | |
| ) -> Union[Tensor, Tuple[Tensor, Tensor]]: | |
| return self.encode(d, mask=valid_mask, return_mask=return_mask) | |
| def __repr__(self) -> str: | |
| return (f"DepthDisparityTransform(c={self.c:.4f}, " | |
| f"inf_cap={self.inf_cap}, invalid_fill={self.invalid_fill})") | |
| # -- presets ------------------------------------------------------------ | |
| def for_multidataset(cls, **kwargs) -> "DepthDisparityTransform": | |
| """Preset for training on NYU + KITTI + ETH3D + DIODE + Sintel. | |
| The value c = 5.0 m is the approximate geometric mean of those | |
| datasets' medians (NYU ~2-3, ETH3D ~5, DIODE ~8, Sintel ~10, | |
| KITTI ~15). With c = 5, every dataset's useful depth range maps to | |
| a well-spread segment of [-1, 1): | |
| NYU 0.3-10 m -> y in [-0.89, +0.33] | |
| ETH3D 0.3-70 m -> y in [-0.89, +0.87] | |
| DIODE 0.3-350m -> y in [-0.89, +0.97] | |
| Sintel 0.5-100m-> y in [-0.82, +0.90] | |
| KITTI 1-80 m -> y in [-0.67, +0.88] | |
| """ | |
| return cls(c=5.0, **kwargs) | |
| def fit( | |
| cls, | |
| depth_samples_by_dataset: Mapping[str, Tensor], | |
| masks_by_dataset: Optional[Mapping[str, Tensor]] = None, | |
| **kwargs, | |
| ) -> "DepthDisparityTransform": | |
| """Fit `c` to the geometric mean of per-dataset medians.""" | |
| c = fit_c_geometric_mean(depth_samples_by_dataset, masks_by_dataset) | |
| return cls(c=c, **kwargs) | |
| # --------------------------------------------------------------------------- | |
| # Self-test + plot: encoded-y coverage on a simulated 5-dataset pool | |
| # --------------------------------------------------------------------------- | |
| if __name__ == "__main__": | |
| torch.manual_seed(0) | |
| # --- basic round-trip -------------------------------------------------- | |
| codec = DepthDisparityTransform.for_multidataset() | |
| print(codec) | |
| d = torch.rand(1, 64, 64, dtype=torch.float64) * 200.0 | |
| y = codec(d) | |
| d_hat = codec.decode(y) | |
| print(f" round-trip max err : {(d - d_hat).abs().max():.2e} m") | |
| # --- simulate 5 datasets ----------------------------------------------- | |
| def lognormal(mean_m, sigma, n=50_000): | |
| return torch.distributions.LogNormal(math.log(mean_m), sigma).sample([n]) | |
| datasets = { | |
| "NYU": lognormal(2.5, 0.5), # indoor | |
| "ETH3D": lognormal(5.0, 0.8), | |
| "DIODE": lognormal(8.0, 1.2), | |
| "Sintel": lognormal(10.0, 0.6), | |
| "KITTI": lognormal(15.0, 0.7), | |
| } | |
| # Clamp to reasonable per-dataset ranges | |
| caps = {"NYU": 10.0, "ETH3D": 70.0, "DIODE": 350.0, | |
| "Sintel": 100.0, "KITTI": 80.0} | |
| for name in datasets: | |
| datasets[name] = datasets[name].clamp(0.1, caps[name]) | |
| # --- sanity: coverage of each dataset under the preset ----------------- | |
| print("\ncoverage of c=5 codec:") | |
| for name, d_ds in datasets.items(): | |
| y_ds = codec(d_ds) | |
| p05, p50, p95 = torch.quantile(y_ds, torch.tensor([0.05, 0.5, 0.95])) | |
| print(f" {name:<7} y5%={p05:+.3f} y50%={p50:+.3f} y95%={p95:+.3f}") | |
| # --- plot: y-histogram comparison across encodings --------------------- | |
| try: | |
| import os, tempfile | |
| os.environ.setdefault("MPLCONFIGDIR", tempfile.mkdtemp(prefix="mpl-")) | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt | |
| # Encodings to compare | |
| def enc_power(d, lam=-3.0, c=10.0 / 3.0): | |
| return 2.0 * (1.0 - (1.0 + d / (abs(lam) * c)) ** (lam + 1.0)) - 1.0 | |
| def enc_loglog(d, mu=5.0, beta=2.0): | |
| return torch.tanh(0.5 * beta * torch.log(d.clamp(min=1e-6) / mu)) | |
| def enc_disp(d, c=5.0): | |
| return (d - c) / (d + c) | |
| encodings = { | |
| "power (c=10/3, indoor)": lambda d: enc_power(d, lam=-3.0, c=10.0 / 3.0), | |
| "power (c=20, driving)": lambda d: enc_power(d, lam=-3.0, c=20.0), | |
| "loglogistic (mu=5, beta=2)": lambda d: enc_loglog(d, mu=5.0, beta=2.0), | |
| "disparity (c=5) <-- proposal": lambda d: enc_disp(d, c=5.0), | |
| } | |
| fig, axes = plt.subplots( | |
| len(encodings), 1, figsize=(9, 2.2 * len(encodings)), sharex=True) | |
| bins = torch.linspace(-1, 1, 41) | |
| colors = { | |
| "NYU": "#1f77b4", | |
| "ETH3D": "#2ca02c", | |
| "DIODE": "#ff7f0e", | |
| "Sintel": "#d62728", | |
| "KITTI": "#9467bd", | |
| } | |
| for ax, (name, enc) in zip(axes, encodings.items()): | |
| for ds_name, d_ds in datasets.items(): | |
| y_ds = enc(d_ds) | |
| ax.hist(y_ds.numpy(), bins=bins.numpy(), density=True, | |
| histtype="step", linewidth=2.0, | |
| label=ds_name, color=colors[ds_name]) | |
| ax.set_ylabel("density") | |
| ax.set_title(name, fontsize=10) | |
| ax.set_xlim(-1.05, 1.05) | |
| ax.grid(True, alpha=0.3) | |
| ax.legend(fontsize=8, ncol=5, loc="upper center") | |
| axes[-1].set_xlabel("encoded y in [-1, 1]") | |
| fig.suptitle("y-distribution coverage across 5 datasets under different encodings", | |
| y=1.02, fontsize=12) | |
| plt.tight_layout() | |
| out = "multidataset_encoding_coverage.png" | |
| plt.savefig(out, dpi=140, bbox_inches="tight") | |
| print(f"\nSaved coverage plot to {out}") | |
| except Exception as e: | |
| print(f"\n(Skipped plot: {e})") | |