| """Contains the UNet decoder. |
| |
| 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_decoder import BaseDecoder |
|
|
|
|
| class UNetDecoder(BaseDecoder): |
| """Decoder of UNet model.""" |
|
|
| def __init__( |
| self, |
| dim_out: int, |
| width: List[int] | int, |
| steps: int = 5, |
| norm_type: NormLayerName = "group_norm", |
| norm_num_groups=8, |
| blocks_per_layer=2, |
| ) -> None: |
| """Initialize UNet Decoder. |
| |
| Args: |
| dim_out: The number of output channels. |
| width: Width of last input feature map from encoder |
| or the width list of all input feature maps from encoder. |
| steps: The number of upsampling 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 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_out = dim_out |
|
|
| self.convs_up = nn.ModuleList() |
|
|
| self.output_dims: list[int] |
| |
| if isinstance(width, int): |
| self.input_dims = [width >> i for i in range(0, steps + 1)] |
| else: |
| self.input_dims = width[::-1][: steps + 1] |
|
|
| for i_step in range(steps): |
| input_width = self.input_dims[i_step] |
| current_width = self.input_dims[i_step + 1] |
| convs_up_i = nn.Sequential( |
| nn.Upsample(scale_factor=2), |
| residual_block_2d( |
| input_width * (1 if i_step == 0 else 2), |
| 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_up.append(convs_up_i) |
| input_width = 2 * current_width |
| current_width //= 2 |
|
|
| last_width = self.input_dims[-1] |
| self.conv_out = nn.Sequential( |
| norm_layer_2d(last_width * 2, norm_type, num_groups=norm_num_groups), |
| nn.ReLU(), |
| nn.Conv2d(last_width * 2, dim_out, 1), |
| norm_layer_2d(dim_out, norm_type, num_groups=norm_num_groups), |
| nn.ReLU(), |
| ) |
|
|
| def forward(self, features: list[torch.Tensor]) -> torch.Tensor: |
| """Apply UNet to image. |
| |
| Args: |
| features: The input multi-level feature map from encoder. |
| |
| Returns: |
| The output feature map. |
| """ |
| i_feature_layer = len(features) - 1 |
| out = self.convs_up[0](features[i_feature_layer]) |
| i_feature_layer -= 1 |
| for conv_up in self.convs_up[1:]: |
| out = conv_up(torch.cat([out, features[i_feature_layer]], dim=1)) |
| i_feature_layer -= 1 |
| out = self.conv_out(torch.cat([out, features[i_feature_layer]], dim=1)) |
|
|
| return out |
|
|
