data/depth_encoding/disparity_transform.py

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

    @classmethod
    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)

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