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