| """Contains backbone models for feature extraction from RGBD input. |
| |
| For licensing see accompanying LICENSE file. |
| Copyright (C) 2025 Apple Inc. All Rights Reserved. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from typing import List |
|
|
| import torch |
| from torch import nn |
|
|
| from sharp.models.blocks import ( |
| NormLayerName, |
| norm_layer_2d, |
| residual_block_2d, |
| ) |
|
|
| from .base_encoder import BaseEncoder |
|
|
|
|
| class UNetEncoder(BaseEncoder): |
| """Encoder of UNet model.""" |
|
|
| def __init__( |
| self, |
| dim_in: int, |
| width: List[int] | int, |
| steps: int = 6, |
| norm_type: NormLayerName = "group_norm", |
| norm_num_groups=8, |
| blocks_per_layer=2, |
| ) -> None: |
| """Initialize UNet Encoder. |
| |
| Args: |
| dim_in: The number of input channels. |
| width: Width multiplicator of intermediate layers or the width list of all layers. |
| steps: The number of downsampling steps. |
| norm_type: Which kind of normalization layer to use. |
| norm_num_groups: How many groups to use for group norm (if relevant). |
| blocks_per_layer: How many residual blocks per layer to use. |
| """ |
| super().__init__() |
|
|
| if blocks_per_layer < 1: |
| raise ValueError("blocks_per_layer must be greater or equal to one.") |
|
|
| self.dim_in = dim_in |
| self.width = width |
| self.num_steps = steps |
|
|
| self.convs_down = nn.ModuleList() |
|
|
| self.output_dims: list[int] |
| |
| if isinstance(width, int): |
| self.output_dims = [width << i for i in range(0, steps + 1)] |
| else: |
| if len(width) != (steps + 1): |
| raise ValueError("Length of width should match the steps for UNetEncoder.") |
| self.output_dims = width |
|
|
| self.conv_in = nn.Sequential( |
| nn.Conv2d(self.dim_in, self.output_dims[0], 3, stride=1, padding=1), |
| norm_layer_2d(self.output_dims[0], norm_type, num_groups=norm_num_groups), |
| nn.ReLU(), |
| ) |
|
|
| for i_step in range(steps): |
| input_width = self.output_dims[i_step] |
| current_width = self.output_dims[i_step + 1] |
| convs_down_i = nn.Sequential( |
| nn.AvgPool2d(2, stride=2), |
| residual_block_2d( |
| input_width, |
| current_width, |
| norm_type=norm_type, |
| norm_num_groups=norm_num_groups, |
| ), |
| *[ |
| residual_block_2d( |
| current_width, |
| current_width, |
| norm_type=norm_type, |
| norm_num_groups=norm_num_groups, |
| ) |
| for _ in range(blocks_per_layer - 1) |
| ], |
| ) |
| self.convs_down.append(convs_down_i) |
|
|
| def forward(self, input: torch.Tensor) -> list[torch.Tensor]: |
| """Apply UNet Encoder to image. |
| |
| Args: |
| input: The input image. |
| |
| Returns: |
| The output multi-level feature map from encoder. |
| """ |
| features = [] |
|
|
| feat_i = self.conv_in(input) |
| features.append(feat_i) |
|
|
| for conv_down in self.convs_down: |
| feat_i = conv_down(feat_i) |
| features.append(feat_i) |
|
|
| return features |
|
|
| @property |
| def out_width(self) -> int: |
| """Compute the output width for UNet decoder.""" |
| return self.output_dims[-1] |
|
|
