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