Spaces:
Running on Zero
Running on Zero
| """Robust alignment routines used by the public depth metrics.""" | |
| import math | |
| from typing import Optional, Tuple | |
| import torch | |
| def _scatter_min( | |
| size: int, index: torch.LongTensor, src: torch.Tensor | |
| ) -> torch.return_types.min: | |
| minimum = torch.full( | |
| (size,), float("inf"), dtype=src.dtype, device=src.device | |
| ).scatter_reduce(0, index, src, reduce="amin", include_self=False) | |
| locations = torch.where(src == minimum[index])[0] | |
| indices = torch.full((size,), -1, dtype=torch.long, device=src.device) | |
| indices[index[locations]] = locations | |
| return torch.return_types.min((minimum, indices)) | |
| def _align_scale( | |
| source: torch.Tensor, | |
| target: torch.Tensor, | |
| weight: torch.Tensor, | |
| eps: float = 1e-7, | |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor]: | |
| """Minimize ``sum(weight * abs(scale * source - target))``.""" | |
| source, target, weight = torch.broadcast_tensors(source, target, weight) | |
| sign = torch.sign(source) | |
| source, target = source * sign, target * sign | |
| target_over_source, order = (target / source.clamp_min(eps)).sort(dim=-1) | |
| weighted_source = torch.gather(source * weight, dim=-1, index=order) | |
| derivatives = ( | |
| 2 * weighted_source.cumsum(dim=-1) | |
| - weighted_source.sum(dim=-1, keepdim=True) | |
| ) | |
| location = torch.searchsorted( | |
| derivatives, torch.zeros_like(derivatives[..., :1]), side="left" | |
| ).clamp_max(derivatives.shape[-1] - 1) | |
| scale = target_over_source.gather(dim=-1, index=location).squeeze(-1) | |
| index = order.gather(dim=-1, index=location).squeeze(-1) | |
| loss = (weight * (scale[..., None] * source - target).abs()).sum(dim=-1) | |
| return scale, loss, index | |
| def align_depth_scale( | |
| source: torch.Tensor, target: torch.Tensor, weight: torch.Tensor | |
| ) -> torch.Tensor: | |
| scale, _, _ = _align_scale(source, target, weight) | |
| return scale | |
| def align_depth_affine( | |
| source: torch.Tensor, target: torch.Tensor, weight: torch.Tensor | |
| ) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """Robust L1 depth alignment with a shared scale and shift.""" | |
| batch_shape, count = source.shape[:-1], source.shape[-1] | |
| batch_size = math.prod(batch_shape) | |
| source = source.reshape(batch_size, count) | |
| target = target.reshape(batch_size, count) | |
| weight = weight.reshape(batch_size, count) | |
| anchor_batch, anchor_index = torch.where(weight > 0) | |
| if anchor_index.numel() == 0: | |
| raise ValueError("Cannot align depth maps without valid pixels") | |
| with torch.no_grad(): | |
| source_anchor = source[anchor_batch, anchor_index] | |
| target_anchor = target[anchor_batch, anchor_index] | |
| source_centered = source[anchor_batch] - source_anchor[:, None] | |
| target_centered = target[anchor_batch] - target_anchor[:, None] | |
| anchor_weight = weight[anchor_batch] | |
| scale, loss, second_index = _align_scale( | |
| source_centered, target_centered, anchor_weight | |
| ) | |
| _, best_anchor = _scatter_min(batch_size, anchor_batch, loss) | |
| first_index = anchor_index[best_anchor] | |
| second_index = second_index[best_anchor] | |
| source_first = source.gather(1, first_index[:, None]).squeeze(1) | |
| target_first = target.gather(1, first_index[:, None]).squeeze(1) | |
| source_second = source.gather(1, second_index[:, None]).squeeze(1) | |
| target_second = target.gather(1, second_index[:, None]).squeeze(1) | |
| scale = (target_second - target_first) / torch.where( | |
| source_second != source_first, | |
| source_second - source_first, | |
| torch.full_like(source_second, 1e-7), | |
| ) | |
| shift = target_first - scale * source_first | |
| return scale.reshape(batch_shape), shift.reshape(batch_shape) | |
| def align_affine_lstsq( | |
| source: torch.Tensor, | |
| target: torch.Tensor, | |
| weight: Optional[torch.Tensor] = None, | |
| ) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """Solve weighted least-squares affine alignment.""" | |
| weight_sqrt = torch.ones_like(source) if weight is None else weight.sqrt() | |
| design = torch.stack( | |
| [weight_sqrt * source, weight_sqrt], dim=-1 | |
| ) | |
| response = (weight_sqrt * target)[..., None] | |
| scale, shift = ( | |
| torch.linalg.lstsq(design, response)[0].squeeze(-1).unbind(-1) | |
| ) | |
| return scale, shift | |