Spaces:
Running on Zero
Running on Zero
File size: 4,232 Bytes
5e5d775 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """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
|