127 lines
4.4 KiB
Python
127 lines
4.4 KiB
Python
"""Contains modules for different types of alignment.
|
|
|
|
For licensing see accompanying LICENSE file.
|
|
Copyright (C) 2025 Apple Inc. All Rights Reserved.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import nn
|
|
|
|
from sharp.models.decoders import UNetDecoder
|
|
from sharp.models.encoders import UNetEncoder
|
|
from sharp.utils import math as math_utils
|
|
|
|
from .params import AlignmentParams
|
|
|
|
|
|
def create_alignment(
|
|
params: AlignmentParams, depth_decoder_dim: int | None = None
|
|
) -> nn.Module | None:
|
|
"""Create depth alignment."""
|
|
if depth_decoder_dim is None:
|
|
raise ValueError("Requires depth_decoder_dim for LearnedAlignment.")
|
|
alignment = LearnedAlignment(
|
|
depth_decoder_features=params.depth_decoder_features,
|
|
depth_decoder_dim=depth_decoder_dim,
|
|
steps=params.steps,
|
|
stride=params.stride,
|
|
base_width=params.base_width,
|
|
activation_type=params.activation_type,
|
|
)
|
|
|
|
if params.frozen:
|
|
alignment.requires_grad_(False)
|
|
|
|
return alignment
|
|
|
|
|
|
class LearnedAlignment(nn.Module):
|
|
"""Aligns tensors using a UNet."""
|
|
|
|
def __init__(
|
|
self,
|
|
steps: int = 4,
|
|
stride: int = 8,
|
|
base_width: int = 16,
|
|
depth_decoder_features: bool = False,
|
|
depth_decoder_dim: int = 256,
|
|
activation_type: math_utils.ActivationType = "exp",
|
|
) -> None:
|
|
"""Initialize LearnedAlignment.
|
|
|
|
Args:
|
|
steps: Number of steps in the UNet.
|
|
stride: Effective downsampling of the alignment module.
|
|
base_width: Base width of the UNet.
|
|
depth_decoder_features: Whether to use depth decoder features.
|
|
depth_decoder_dim: Dimension of the depth decoder features.
|
|
activation_type: Activation type for the alignment output.
|
|
"""
|
|
super().__init__()
|
|
self.activation = math_utils.create_activation_pair(activation_type)
|
|
bias_value = self.activation.inverse(torch.tensor(1.0))
|
|
|
|
self.depth_decoder_features = depth_decoder_features
|
|
if depth_decoder_features:
|
|
dim_in = 2 + depth_decoder_dim
|
|
else:
|
|
dim_in = 2
|
|
|
|
def is_power_of_two(n: int) -> bool:
|
|
"""Check if a number is a power of two."""
|
|
if n <= 0:
|
|
return False
|
|
return (n & (n - 1)) == 0
|
|
|
|
if not is_power_of_two(stride):
|
|
raise ValueError(f"Stride {stride} is not a power of two.")
|
|
|
|
steps_decoder = steps - int(math.log2(stride))
|
|
if steps_decoder < 1:
|
|
raise ValueError(f"{steps_decoder} must be greater or equal to 1.")
|
|
widths = [min(base_width << i, 1024) for i in range(steps + 1)]
|
|
self.encoder = UNetEncoder(dim_in=dim_in, width=widths, steps=steps, norm_num_groups=4)
|
|
self.decoder = UNetDecoder(
|
|
dim_out=widths[0], width=widths, steps=steps_decoder, norm_num_groups=4
|
|
)
|
|
self.conv_out = nn.Conv2d(widths[0], 1, 1, bias=True)
|
|
nn.init.zeros_(self.conv_out.weight)
|
|
nn.init.constant_(self.conv_out.bias, bias_value)
|
|
|
|
def forward(
|
|
self,
|
|
tensor_src: torch.Tensor,
|
|
tensor_tgt: torch.Tensor,
|
|
depth_decoder_features: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
"""Compute alignment map."""
|
|
# Since the tensors are usually given by depth which is >= 1.0, we invert
|
|
# the tensors to have them in a reasonable range.
|
|
tensor_src = 1.0 / tensor_src.clamp(min=1e-4)
|
|
tensor_tgt = 1.0 / tensor_tgt.clamp(min=1e-4)
|
|
tensor_input = torch.cat([tensor_src, tensor_tgt], dim=1)
|
|
if self.depth_decoder_features:
|
|
height, width = tensor_src.shape[-2:]
|
|
upsampled_encodings = F.interpolate(
|
|
depth_decoder_features,
|
|
size=(height, width),
|
|
mode="bilinear",
|
|
)
|
|
tensor_input = torch.cat([tensor_input, upsampled_encodings], dim=1)
|
|
features = self.encoder(tensor_input)
|
|
output = self.conv_out(self.decoder(features))
|
|
alignment_map_lowres = self.activation.forward(output)
|
|
if alignment_map_lowres.shape[-2:] != tensor_src.shape[-2]:
|
|
alignment_map = F.interpolate(
|
|
alignment_map_lowres,
|
|
size=tensor_src.shape[-2:],
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
)
|
|
return alignment_map
|