""" PyTorch implementation of Barron's power transform for metric depth encoding: f(d; lam, c) = 1 - (1 + d / (|lam| * c)) ** (lam + 1) = -expm1( (lam + 1) * log1p( d / (|lam| * c) ) ) The second form is what we actually evaluate: it is numerically stable for every `lam < -1` and avoids `0 ** 0` / `0 ** neg` pathologies that hit the naive closed form. The paper's depth encoding uses `(lam, c) = (-3, 10/3)` so that `f(d) = 1 - 1 / (1 + d/10)^2`; keep these as defaults. Valid range for `lam`: `lam < -1` (strictly). * At `lam = -1` the family degenerates (all outputs are 0). * For `lam > -1` the image is no longer contained in `[0, 1)`, so the depth encoding is ill-defined. Practical range for `lam`: * `lam ~= -3` (the paper's default) gives roughly uniform resolution per octave of depth over ~[0.3 m, 100 m]. * Useful range is roughly `lam in [-10, -1.5]`. Past `lam ~ -15` the transform saturates so aggressively that `y = 2f - 1` cannot resolve large depths even in float64 (every `d > c` maps to within one ulp of 1); this is a property of the mapping, not of the implementation. Functions --------- power_transform(d, lam, c, mask) : [0, inf) -> [0, 1) inverse_power_transform(f, lam, c, ...) : [0, 1) -> [0, inf) depth_to_signed_unit(d, lam, c, mask, ...) : [0, inf) -> [-1, 1) signed_unit_to_depth(y, lam, c, ...) : [-1, 1) -> [0, inf) `lam` and `c` may each be either a Python float, a 0-d tensor, or a tensor broadcastable against `d` (e.g. a per-sample `lam` of shape `[B, 1, 1, 1]`). Class ----- DepthPowerTransform : plain container that bundles `(lam, c, inf_cap, invalid_fill)` and exposes `encode` / `decode` (also callable via `()`). It is NOT an `nn.Module` -- this is purely a normalization utility, so it has no buffers / parameters and doesn't need `.to(device)`. Device/dtype follow the input tensor automatically. """ from __future__ import annotations from typing import Optional, Tuple, Union import torch from torch import Tensor # A scalar "shape parameter" may be passed as either a Python float or a # Tensor (so it can be a learnable `nn.Parameter`, a per-sample value, ...). ScalarLike = Union[float, Tensor] # --------------------------------------------------------------------------- # Small utility: bring a ScalarLike onto the same dtype / device as a ref # tensor while preserving its shape (so broadcasting Just Works). # --------------------------------------------------------------------------- def _as_like(x: ScalarLike, ref: Tensor) -> Tensor: if isinstance(x, Tensor): return x.to(dtype=ref.dtype, device=ref.device) return torch.as_tensor(x, dtype=ref.dtype, device=ref.device) # --------------------------------------------------------------------------- # Core forward / inverse # --------------------------------------------------------------------------- def power_transform( d: Tensor, lam: ScalarLike = -3.0, c: ScalarLike = 10.0 / 3.0, mask: Optional[Tensor] = None, invalid_fill: float = 0.0, ) -> Tensor: """Map metric depth `d` in `[0, inf)` to normalized `f` in `[0, 1)`. Computed as f = -expm1( (lam + 1) * log1p( d / (|lam| * c) ) ) which is stable for any `lam < -1` and well-defined for any `d >= 0`. Args: d : depth tensor, any shape (e.g. (1, H, W), (B, 1, H, W)). lam, c : power-transform parameters; float, 0-d tensor, or tensor broadcastable to `d` (lets you learn them or set them per-sample). mask : optional bool tensor broadcastable to `d`. If given, pixels with `mask == False` or `d <= 0` are written as `invalid_fill` in the output. If `None`, we just clamp `d` to be non-negative. invalid_fill : scalar value written at invalid pixels (default 0.0). Returns: f : same shape as `d`, values in `[0, 1)`. """ lam_t = _as_like(lam, d) c_t = _as_like(c, d) d_safe = d.clamp(min=0.0) # |lam| * c is always > 0 in the regime we care about (lam < -1, c > 0), # but we still guard against accidental zero. denom = lam_t.abs() * c_t denom = denom.clamp(min=torch.finfo(d.dtype).tiny) f = -torch.expm1((lam_t + 1.0) * torch.log1p(d_safe / denom)) if mask is not None: valid = mask.to(torch.bool) & (d > 0) f = torch.where(valid, f, torch.full_like(f, invalid_fill)) return f def inverse_power_transform( f: Tensor, lam: ScalarLike = -3.0, c: ScalarLike = 10.0 / 3.0, inf_cap: Optional[float] = None, _one_minus_f: Optional[Tensor] = None, ) -> Tensor: """Map normalized `f` in `[0, 1)` back to metric depth `d` in `[0, inf)`. Computed as d = |lam| * c * expm1( log(1 - f) / (lam + 1) ) `_one_minus_f` is an optional precomputed `(1 - f)` supplied by upstream callers that already have it in higher precision. This matters because when `f` is close to 1, computing `1 - f` directly in float32 throws away much of the precision in `f`; `signed_unit_to_depth` passes `(1 - y)/2` through this path. At `f = 1` the depth is `+inf`; pass an `inf_cap` (e.g. 1e4) to keep gradients finite during VAE training. """ lam_t = _as_like(lam, f) c_t = _as_like(c, f) if _one_minus_f is None: one_minus_f = (1.0 - f).clamp(min=0.0, max=1.0) else: one_minus_f = _one_minus_f.clamp(min=0.0, max=1.0) # Avoid exact zero so `log(0) = -inf -> exp(-inf) = 0` is a finite path. eps = torch.finfo(one_minus_f.dtype).tiny one_minus_f = one_minus_f.clamp(min=eps) denom_lam = (lam_t + 1.0) # Guard against lam_t -> -1 producing a 0-division at inference time. denom_lam = torch.where( denom_lam.abs() < eps, torch.full_like(denom_lam, -eps), denom_lam, ) u = torch.expm1(torch.log(one_minus_f) / denom_lam) d = lam_t.abs() * c_t * u if inf_cap is not None: d = d.clamp(max=inf_cap) return d # --------------------------------------------------------------------------- # VAE-friendly: [0, inf) <-> [-1, 1) with optional mask # --------------------------------------------------------------------------- def depth_to_signed_unit( d: Tensor, lam: ScalarLike = -3.0, c: ScalarLike = 10.0 / 3.0, mask: Optional[Tensor] = None, invalid_fill: float = 0.0, return_mask: bool = False, ) -> Union[Tensor, Tuple[Tensor, Tensor]]: """Encode metric depth in `[0, inf)` to `y` in `[-1, 1)` via y = 2 * power_transform(d, lam, c) - 1 Invalid pixels (either `mask == False` or `d <= 0`) are filled with `invalid_fill`. Using `invalid_fill = 0` (the neutral centre of `[-1, 1]`) and passing the mask to the loss separately keeps real `d = 0` pixels (which encode to `y = -1`) distinct from "unknown" pixels. Args: d : depth tensor, any shape. lam, c : scalars or tensors broadcastable to `d`. mask : optional bool tensor, same broadcastable shape as `d`. True = valid, False = invalid. If None, every `d > 0` is considered valid. invalid_fill : value written at invalid pixels. return_mask : if True, also return the effective valid-mask. Returns: y : same shape as `d`, in `[-1, 1)`. valid (opt.) : bool tensor of same shape as `d`. """ f = power_transform(d, lam=lam, c=c) # raw [0,1) y = 2.0 * f - 1.0 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 signed_unit_to_depth( y: Tensor, lam: ScalarLike = -3.0, c: ScalarLike = 10.0 / 3.0, inf_cap: Optional[float] = 1e4, ) -> Tensor: """Decode `y` in `[-1, 1)` back to metric depth `d` in `[0, inf)`. Internally uses `(1 - f) = (1 - y) / 2` so that no precision is lost when `y` is close to 1 (large depth). """ y_c = y.clamp(min=-1.0, max=1.0) one_minus_f = 0.5 * (1.0 - y_c) f = 1.0 - one_minus_f # for API consistency return inverse_power_transform( f, lam=lam, c=c, inf_cap=inf_cap, _one_minus_f=one_minus_f ) # --------------------------------------------------------------------------- # Class wrapper: define once, use everywhere # --------------------------------------------------------------------------- class DepthPowerTransform: """Plain container bundling the power-transform depth codec's settings. Pure normalization utility -- no buffers, no parameters, not an `nn.Module`. The stored `lam` / `c` are Python floats; the actual forward / inverse computations happen entirely inside `depth_to_signed_unit` / `signed_unit_to_depth`, and device & dtype are picked up automatically from whatever tensor you pass in. Example ------- >>> codec = DepthPowerTransform() # lam=-3, c=10/3 >>> y = codec.encode(depth, mask=mask) # [0, inf) -> [-1, 1) >>> y = codec(depth, mask=mask) # __call__ == encode >>> y, valid = codec(depth, mask=mask, return_mask=True) >>> d_hat = codec.decode(y_pred) # [-1, 1) -> [0, inf) """ def __init__( self, lam: float = -3.0, c: float = 10.0 / 3.0, inf_cap: Optional[float] = 1e4, invalid_fill: float = 0.0, ): if lam >= -1.0: raise ValueError(f"lam must be < -1, got {lam}") if c <= 0.0: raise ValueError(f"c must be > 0, got {c}") self.lam = float(lam) 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]]: """Depth in `[0, inf)` -> y in `[-1, 1)`. See `depth_to_signed_unit`.""" return depth_to_signed_unit( d, lam=self.lam, c=self.c, mask=mask, invalid_fill=self.invalid_fill, return_mask=return_mask, ) def decode(self, y: Tensor) -> Tensor: """y in `[-1, 1)` -> depth in `[0, inf)`. See `signed_unit_to_depth`.""" return signed_unit_to_depth( y, lam=self.lam, 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"DepthPowerTransform(lam={self.lam:.4f}, c={self.c:.4f}, " f"inf_cap={self.inf_cap}, invalid_fill={self.invalid_fill})") # --------------------------------------------------------------------------- # Self-test: round-trip, gradient w.r.t. lam, learnable-lam example # --------------------------------------------------------------------------- if __name__ == "__main__": torch.manual_seed(0) # 1) round-trip on a random (1, H, W) depth map with a mask d = torch.rand(1, 32, 32) * 100.0 mask = d > 0.5 y, valid = depth_to_signed_unit(d, mask=mask, return_mask=True) d_hat = signed_unit_to_depth(y) print(f"shape : {tuple(d.shape)}") print(f"y range : [{y.min().item():+.4f}, {y.max().item():+.4f}]") print(f"y at invalid pixels: {y[~valid].unique().tolist()} (expect [0.0])") print(f"round-trip max err : {(d - d_hat)[valid].abs().max():.2e} m") # 2) match the legacy closed form at (lam, c) = (-3, 10/3) for sanity d2 = torch.linspace(0.0, 200.0, 500, dtype=torch.float64) f_legacy = 1.0 - 1.0 / (1.0 + d2 / 10.0) ** 2 f_new = power_transform(d2, lam=-3.0, c=10.0 / 3.0) print(f"\nmax|f_new - f_legacy| = {(f_new - f_legacy).abs().max():.2e}") # 3) vary `lam` freely: round-trip for several values (float64 for headroom) print("\nround-trip for varying lam (float64, should all be ~0):") for lam in [-1.5, -2.0, -3.0, -5.0, -10.0]: d3 = torch.rand(1, 64, 64, dtype=torch.float64) * 200.0 y3 = depth_to_signed_unit(d3, lam=lam, c=10.0 / 3.0) d3h = signed_unit_to_depth(y3, lam=lam, c=10.0 / 3.0) err = (d3 - d3h).abs().max().item() print(f" lam = {lam:>6.2f} max|d - d_hat| = {err:.2e} m") # 4) per-sample `lam`: different shape parameter for each image in batch d4 = torch.rand(4, 1, 32, 32, dtype=torch.float64) * 100.0 lam4 = torch.tensor([-1.5, -3.0, -5.0, -10.0]).view(4, 1, 1, 1) y4 = depth_to_signed_unit(d4, lam=lam4) d4h = signed_unit_to_depth(y4, lam=lam4) print(f"\nper-sample lam round-trip max err: " f"{(d4 - d4h).abs().max():.2e} m") # 5) `lam` as a learnable nn.Parameter: gradient flows through it lam_p = torch.nn.Parameter(torch.tensor(-3.0)) c_p = torch.nn.Parameter(torch.tensor(10.0 / 3.0)) d5 = (torch.rand(2, 1, 16, 16) * 50.0).requires_grad_(True) y5 = depth_to_signed_unit(d5, lam=lam_p, c=c_p) loss = (y5 - 0.1).pow(2).mean() loss.backward() print(f"\nlearnable lam | loss = {loss.item():.4f}" f" d/dlam = {lam_p.grad.item():+.4f}" f" d/dc = {c_p.grad.item():+.4f}" f" max|d/dd| = {d5.grad.abs().max():.4f}") # 6) DepthPowerTransform class: define once, use everywhere print("\n--- DepthPowerTransform class ---") codec = DepthPowerTransform() print(codec) d6 = torch.rand(2, 1, 48, 48) * 80.0 m6 = d6 > 0.2 y6, v6 = codec(d6, mask=m6, return_mask=True) # __call__ == encode d6h = codec.decode(y6) print(f" round-trip err : {(d6 - d6h)[v6].abs().max():.2e} m") # Non-default parameters. codec_b = DepthPowerTransform(lam=-4.0, c=5.0, inf_cap=1e3, invalid_fill=-1.0) print(codec_b) # Device/dtype follow the input tensor; no `.to()` / `.double()` needed. d8 = torch.rand(1, 64, 64, dtype=torch.float64) * 200.0 y8 = codec(d8) print(f" float64 in/out : y.dtype={y8.dtype}, " f"round-trip = {(d8 - codec.decode(y8)).abs().max():.2e} m")