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