mapdet3d/utils/align.py

"""Alignment."""

from __future__ import annotations

import math
from typing import Optional, Tuple, Union

import numpy as np
import torch


def align_depth_least_square(
    gt_arr: np.ndarray,
    pred_arr: np.ndarray,
    valid_mask_arr: np.ndarray,
    return_scale_shift=True,
    max_resolution=None,
):
    ori_shape = pred_arr.shape  # input shape

    gt = gt_arr.squeeze()  # [H, W]
    pred = pred_arr.squeeze()
    valid_mask = valid_mask_arr.squeeze()

    # Downsample
    if max_resolution is not None:
        scale_factor = np.min(max_resolution / np.array(ori_shape[-2:]))
        if scale_factor < 1:
            downscaler = torch.nn.Upsample(
                scale_factor=scale_factor, mode="nearest"
            )
            gt = downscaler(torch.as_tensor(gt).unsqueeze(0)).numpy()
            pred = downscaler(torch.as_tensor(pred).unsqueeze(0)).numpy()
            valid_mask = (
                downscaler(torch.as_tensor(valid_mask).unsqueeze(0).float())
                .bool()
                .numpy()
            )

    assert (
        gt.shape == pred.shape == valid_mask.shape
    ), f"{gt.shape}, {pred.shape}, {valid_mask.shape}"

    gt_masked = gt[valid_mask].reshape((-1, 1))
    pred_masked = pred[valid_mask].reshape((-1, 1))

    # numpy solver
    _ones = np.ones_like(pred_masked)
    A = np.concatenate([pred_masked, _ones], axis=-1)
    X = np.linalg.lstsq(A, gt_masked, rcond=None)[0]
    scale, shift = X

    aligned_pred = pred_arr * scale + shift

    # restore dimensions
    aligned_pred = aligned_pred.reshape(ori_shape)

    if return_scale_shift:
        return aligned_pred, scale, shift
    else:
        return aligned_pred


def _pad_inf(x_: torch.Tensor):
    return torch.cat(
        [
            torch.full_like(x_[..., :1], -torch.inf),
            x_,
            torch.full_like(x_[..., :1], torch.inf),
        ],
        dim=-1,
    )


def _pad_cumsum(cumsum: torch.Tensor):
    return torch.cat(
        [torch.zeros_like(cumsum[..., :1]), cumsum, cumsum[..., -1:]], dim=-1
    )


def _compute_residual(a: torch.Tensor, xyw: torch.Tensor, trunc: float):
    return (
        a.mul(xyw[..., 0])
        .sub_(xyw[..., 1])
        .abs_()
        .mul_(xyw[..., 2])
        .clamp_max_(trunc)
        .sum(dim=-1)
    )


def align(
    x: torch.Tensor,
    y: torch.Tensor,
    w: torch.Tensor,
    trunc: Optional[Union[float, torch.Tensor]] = None,
    eps: float = 1e-7,
) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:
    """
    If trunc is None, solve `min sum_i w_i * |a * x_i - y_i|`, otherwise solve `min sum_i min(trunc, w_i * |a * x_i - y_i|)`.

    w_i must be >= 0.

    ### Parameters:
    - `x`: tensor of shape (..., n)
    - `y`: tensor of shape (..., n)
    - `w`: tensor of shape (..., n)
    - `trunc`: optional, float or tensor of shape (..., n) or None

    ### Returns:
    - `a`: tensor of shape (...), differentiable
    - `loss`: tensor of shape (...), value of loss function at `a`, detached
    - `index`: tensor of shape (...), where a = y[idx] / x[idx]
    """
    if trunc is None:
        x, y, w = torch.broadcast_tensors(x, y, w)
        sign = torch.sign(x)
        x, y = x * sign, y * sign
        y_div_x = y / x.clamp_min(eps)
        y_div_x, argsort = y_div_x.sort(dim=-1)

        wx = torch.gather(x * w, dim=-1, index=argsort)
        derivatives = 2 * wx.cumsum(dim=-1) - wx.sum(dim=-1, keepdim=True)
        search = torch.searchsorted(
            derivatives, torch.zeros_like(derivatives[..., :1]), side="left"
        ).clamp_max(derivatives.shape[-1] - 1)

        a = y_div_x.gather(dim=-1, index=search).squeeze(-1)
        index = argsort.gather(dim=-1, index=search).squeeze(-1)
        loss = (w * (a[..., None] * x - y).abs()).sum(dim=-1)

    else:
        # Reshape to (batch_size, n) for simplicity
        x, y, w = torch.broadcast_tensors(x, y, w)
        batch_shape = x.shape[:-1]
        batch_size = math.prod(batch_shape)
        x, y, w = (
            x.reshape(-1, x.shape[-1]),
            y.reshape(-1, y.shape[-1]),
            w.reshape(-1, w.shape[-1]),
        )

        sign = torch.sign(x)
        x, y = x * sign, y * sign
        wx, wy = w * x, w * y
        xyw = torch.stack(
            [x, y, w], dim=-1
        )  # Stacked for convenient gathering

        y_div_x = A = y / x.clamp_min(eps)
        B = (wy - trunc) / wx.clamp_min(eps)
        C = (wy + trunc) / wx.clamp_min(eps)
        with torch.no_grad():
            # Caculate prefix sum by orders of A, B, C
            A, A_argsort = A.sort(dim=-1)
            Q_A = torch.cumsum(
                torch.gather(wx, dim=-1, index=A_argsort), dim=-1
            )
            A, Q_A = _pad_inf(A), _pad_cumsum(
                Q_A
            )  # Pad [-inf, A1, ..., An, inf] and [0, Q1, ..., Qn, Qn] to handle edge cases.

            B, B_argsort = B.sort(dim=-1)
            Q_B = torch.cumsum(
                torch.gather(wx, dim=-1, index=B_argsort), dim=-1
            )
            B, Q_B = _pad_inf(B), _pad_cumsum(Q_B)

            C, C_argsort = C.sort(dim=-1)
            Q_C = torch.cumsum(
                torch.gather(wx, dim=-1, index=C_argsort), dim=-1
            )
            C, Q_C = _pad_inf(C), _pad_cumsum(Q_C)

            # Caculate left and right derivative of A
            j_A = torch.searchsorted(A, y_div_x, side="left").sub_(1)
            j_B = torch.searchsorted(B, y_div_x, side="left").sub_(1)
            j_C = torch.searchsorted(C, y_div_x, side="left").sub_(1)
            left_derivative = (
                2 * torch.gather(Q_A, dim=-1, index=j_A)
                - torch.gather(Q_B, dim=-1, index=j_B)
                - torch.gather(Q_C, dim=-1, index=j_C)
            )
            j_A = torch.searchsorted(A, y_div_x, side="right").sub_(1)
            j_B = torch.searchsorted(B, y_div_x, side="right").sub_(1)
            j_C = torch.searchsorted(C, y_div_x, side="right").sub_(1)
            right_derivative = (
                2 * torch.gather(Q_A, dim=-1, index=j_A)
                - torch.gather(Q_B, dim=-1, index=j_B)
                - torch.gather(Q_C, dim=-1, index=j_C)
            )

            # Find extrema
            is_extrema = (left_derivative < 0) & (right_derivative >= 0)
            is_extrema[..., 0] |= ~is_extrema.any(
                dim=-1
            )  # In case all derivatives are zero, take the first one as extrema.
            where_extrema_batch, where_extrema_index = torch.where(is_extrema)

            # Calculate objective value at extrema
            extrema_a = y_div_x[
                where_extrema_batch, where_extrema_index
            ]  # (num_extrema,)
            MAX_ELEMENTS = (
                4096**2
            )  # Split into small batches to avoid OOM in case there are too many extrema.(~1G)
            SPLIT_SIZE = MAX_ELEMENTS // x.shape[-1]
            extrema_value = torch.cat(
                [
                    _compute_residual(
                        extrema_a_split[:, None],
                        xyw[extrema_i_split, :, :],
                        trunc,
                    )
                    for extrema_a_split, extrema_i_split in zip(
                        extrema_a.split(SPLIT_SIZE),
                        where_extrema_batch.split(SPLIT_SIZE),
                    )
                ]
            )  # (num_extrema,)

            # Find minima among corresponding extrema
            minima, indices = scatter_min(
                size=batch_size,
                dim=0,
                index=where_extrema_batch,
                src=extrema_value,
            )  # (batch_size,)
            index = where_extrema_index[indices]

        a = torch.gather(y, dim=-1, index=index[..., None]) / torch.gather(
            x, dim=-1, index=index[..., None]
        ).clamp_min(eps)
        a = a.reshape(batch_shape)
        loss = minima.reshape(batch_shape)
        index = index.reshape(batch_shape)

    return a, loss, index


def scatter_min(
    size: int, dim: int, index: torch.LongTensor, src: torch.Tensor
) -> torch.return_types.min:
    "Scatter the minimum value along the given dimension of `input` into `src` at the indices specified in `index`."
    shape = src.shape[:dim] + (size,) + src.shape[dim + 1 :]
    minimum = torch.full(
        shape, float("inf"), dtype=src.dtype, device=src.device
    ).scatter_reduce(
        dim=dim, index=index, src=src, reduce="amin", include_self=False
    )
    minimum_where = torch.where(
        src == torch.gather(minimum, dim=dim, index=index)
    )
    indices = torch.full(shape, -1, dtype=torch.long, device=src.device)
    indices[
        (*minimum_where[:dim], index[minimum_where], *minimum_where[dim + 1 :])
    ] = minimum_where[dim]
    return torch.return_types.min((minimum, indices))