Spaces:

hugging-apps
/

geoug

Running on Zero

File size: 8,870 Bytes

5e5d775

"""
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

    @property
    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}%")