model.py

# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# The deconvolution code is based on Simple Baseline.
# (https://github.com/microsoft/human-pose-estimation.pytorch/blob/master/lib/models/pose_resnet.py)
# Modified by Zigang Geng (zigang@mail.ustc.edu.cn).
# ------------------------------------------------------------------------------

import torch
import torch.nn as nn
from timm.models.layers import trunc_normal_, DropPath
from mmcv.cnn import (build_conv_layer, build_norm_layer, build_upsample_layer,
                      constant_init, normal_init)
from omegaconf import OmegaConf
from ldm.util import instantiate_from_config
import torch.nn.functional as F
import sys
import os
current_script_path = os.path.abspath(__file__)
parent_folder_path = os.path.dirname(os.path.dirname(current_script_path))
sys.path.append(parent_folder_path)
parent_folder_path = os.path.dirname(parent_folder_path)
print(parent_folder_path)
# Add the parent folder to sys.path
sys.path.append(parent_folder_path)

from .evpconfig import EVPConfig

from .models import UNetWrapper, TextAdapterRefer, FrozenCLIPEmbedder
from .miniViT import mViT
from .attractor import AttractorLayer, AttractorLayerUnnormed
from .dist_layers import ConditionalLogBinomial
from .localbins_layers import (Projector, SeedBinRegressor, SeedBinRegressorUnnormed)
import os
from transformers import PreTrainedModel
import sys
current_script_path = os.path.abspath(__file__)
parent_folder_path = os.path.dirname(os.path.dirname(current_script_path))
import torchvision.transforms as transforms

# Add the parent folder to sys.path
sys.path.append(parent_folder_path)

def icnr(x, scale=2, init=nn.init.kaiming_normal_):
    """
    Checkerboard artifact free sub-pixel convolution
    https://arxiv.org/abs/1707.02937
    """
    ni,nf,h,w = x.shape
    ni2 = int(ni/(scale**2))
    k = init(torch.zeros([ni2,nf,h,w])).transpose(0, 1)
    k = k.contiguous().view(ni2, nf, -1)
    k = k.repeat(1, 1, scale**2)
    k = k.contiguous().view([nf,ni,h,w]).transpose(0, 1)
    x.data.copy_(k)
    
    
class PixelShuffle(nn.Module):
    """
    Real-Time Single Image and Video Super-Resolution
    https://arxiv.org/abs/1609.05158
    """
    def __init__(self, n_channels, scale):
        super(PixelShuffle, self).__init__()
        self.conv = nn.Conv2d(n_channels, n_channels*(scale**2), kernel_size=1)
        icnr(self.conv.weight)
        self.shuf = nn.PixelShuffle(scale)
        self.relu = nn.ReLU()

    def forward(self,x):
        x = self.shuf(self.relu(self.conv(x)))
        return x
        

class AttentionModule(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(AttentionModule, self).__init__()
        
        # Convolutional Layers
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
        
        # Group Normalization
        self.group_norm = nn.GroupNorm(20, out_channels)
        
        # ReLU Activation
        self.relu = nn.ReLU()
        
        # Spatial Attention
        self.spatial_attention = nn.Sequential(
            nn.Conv2d(in_channels, 1, kernel_size=1),
            nn.Sigmoid()
        )
        
    def forward(self, x):
        # Apply spatial attention
        spatial_attention = self.spatial_attention(x)
        x = x * spatial_attention
        
        # Apply convolutional layer
        x = self.conv1(x)
        x = self.group_norm(x)
        x = self.relu(x)
        
        return x
        
        
class AttentionDownsamplingModule(nn.Module):
    def __init__(self, in_channels, out_channels, scale_factor=2):
        super(AttentionDownsamplingModule, self).__init__()
        
        # Spatial Attention
        self.spatial_attention = nn.Sequential(
            nn.Conv2d(in_channels, 1, kernel_size=1),
            nn.Sigmoid()
        )
        
        # Channel Attention
        self.channel_attention = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, in_channels // 8, kernel_size=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels // 8, in_channels, kernel_size=1),
            nn.Sigmoid()
        )
        
        # Convolutional Layers
        if scale_factor == 2:
            self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
        elif scale_factor == 4:
            self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1)
            
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1)
        
        # Group Normalization
        self.group_norm = nn.GroupNorm(20, out_channels)
        
        # ReLU Activation
        self.relu = nn.ReLU(inplace=True)
        
    def forward(self, x):
        # Apply spatial attention
        spatial_attention = self.spatial_attention(x)
        x = x * spatial_attention
        
        # Apply channel attention
        channel_attention = self.channel_attention(x)
        x = x * channel_attention
        
        # Apply convolutional layers
        x = self.conv1(x)
        x = self.group_norm(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.group_norm(x)
        x = self.relu(x)
        
        return x


class AttentionUpsamplingModule(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(AttentionUpsamplingModule, self).__init__()
        
        # Spatial Attention for outs[2]
        self.spatial_attention = nn.Sequential(
            nn.Conv2d(in_channels, 1, kernel_size=1),
            nn.Sigmoid()
        )
        
        # Channel Attention for outs[2]
        self.channel_attention = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, in_channels // 8, kernel_size=1),
            nn.ReLU(),
            nn.Conv2d(in_channels // 8, in_channels, kernel_size=1),
            nn.Sigmoid()
        )
        
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
        
        # Group Normalization
        self.group_norm = nn.GroupNorm(20, out_channels)
        
        # ReLU Activation
        self.relu = nn.ReLU()
        self.upscale = PixelShuffle(in_channels, 2)
        
    def forward(self, x):
        # Apply spatial attention
        spatial_attention = self.spatial_attention(x)
        x = x * spatial_attention
        
        # Apply channel attention
        channel_attention = self.channel_attention(x)
        x = x * channel_attention
        
        # Apply convolutional layers
        x = self.conv1(x)
        x = self.group_norm(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.group_norm(x)
        x = self.relu(x)
        
        # Upsample
        x = self.upscale(x)
        
        return x
        
        
class ConvLayer(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ConvLayer, self).__init__()
        
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1),
            nn.GroupNorm(20, out_channels),
            nn.ReLU(),
        )
        
    def forward(self, x):
        x = self.conv1(x)
        
        return x
        
        
class InverseMultiAttentiveFeatureRefinement(nn.Module):
    def __init__(self, in_channels_list):
        super(InverseMultiAttentiveFeatureRefinement, self).__init__()
        
        self.layer1 = AttentionModule(in_channels_list[0], in_channels_list[0])
        self.layer2 = AttentionDownsamplingModule(in_channels_list[0], in_channels_list[0]//2, scale_factor = 2)
        self.layer3 = ConvLayer(in_channels_list[0]//2 + in_channels_list[1], in_channels_list[1])
        self.layer4 = AttentionDownsamplingModule(in_channels_list[1], in_channels_list[1]//2, scale_factor = 2)
        self.layer5 = ConvLayer(in_channels_list[1]//2 + in_channels_list[2], in_channels_list[2])
        self.layer6 = AttentionDownsamplingModule(in_channels_list[2], in_channels_list[2]//2, scale_factor = 2)
        self.layer7 = ConvLayer(in_channels_list[2]//2 + in_channels_list[3], in_channels_list[3])
        
        '''
        self.layer8 = AttentionUpsamplingModule(in_channels_list[3], in_channels_list[3])
        self.layer9 = ConvLayer(in_channels_list[2] + in_channels_list[3], in_channels_list[2])
        self.layer10 = AttentionUpsamplingModule(in_channels_list[2], in_channels_list[2])
        self.layer11 = ConvLayer(in_channels_list[1] + in_channels_list[2], in_channels_list[1])
        self.layer12 = AttentionUpsamplingModule(in_channels_list[1], in_channels_list[1])
        self.layer13 = ConvLayer(in_channels_list[0] + in_channels_list[1], in_channels_list[0])
        '''
    def forward(self, inputs):
        x_c4, x_c3, x_c2, x_c1 = inputs
        x_c4 = self.layer1(x_c4)
        x_c4_3 = self.layer2(x_c4)
        x_c3 = torch.cat([x_c4_3, x_c3], dim=1)
        x_c3 = self.layer3(x_c3)
        x_c3_2 = self.layer4(x_c3)
        x_c2 = torch.cat([x_c3_2, x_c2], dim=1)
        x_c2 = self.layer5(x_c2)
        x_c2_1 = self.layer6(x_c2)
        x_c1 = torch.cat([x_c2_1, x_c1], dim=1)
        x_c1 = self.layer7(x_c1)
        '''
        x_c1_2 = self.layer8(x_c1)
        x_c2 = torch.cat([x_c1_2, x_c2], dim=1)
        x_c2 = self.layer9(x_c2)
        x_c2_3 = self.layer10(x_c2)
        x_c3 = torch.cat([x_c2_3, x_c3], dim=1)
        x_c3 = self.layer11(x_c3)
        x_c3_4 = self.layer12(x_c3)
        x_c4 = torch.cat([x_c3_4, x_c4], dim=1)
        x_c4 = self.layer13(x_c4)
        '''
        return [x_c4, x_c3, x_c2, x_c1]
        

class EVPDepthEncoder(nn.Module):
    def __init__(self, out_dim=1024, ldm_prior=[320, 680, 1320+1280], sd_path=None, text_dim=768,
                 dataset='nyu', caption_aggregation=False
                ):
        super().__init__()


        self.layer1 = nn.Sequential(
            nn.Conv2d(ldm_prior[0], ldm_prior[0], 3, stride=2, padding=1),
            nn.GroupNorm(16, ldm_prior[0]),
            nn.ReLU(),
            nn.Conv2d(ldm_prior[0], ldm_prior[0], 3, stride=2, padding=1),
        )

        self.layer2 = nn.Sequential(
            nn.Conv2d(ldm_prior[1], ldm_prior[1], 3, stride=2, padding=1),
        )

        self.out_layer = nn.Sequential(
            nn.Conv2d(sum(ldm_prior), out_dim, 1),
            nn.GroupNorm(16, out_dim),
            nn.ReLU(),
        )
        
        self.aggregation = InverseMultiAttentiveFeatureRefinement([320, 680, 1320, 1280])

        self.apply(self._init_weights)

        ### stable diffusion layers 

        config = OmegaConf.load('./v1-inference.yaml')
        if sd_path is None:
            if os.path.exists('../checkpoints/v1-5-pruned-emaonly.ckpt'):
                config.model.params.ckpt_path = '../checkpoints/v1-5-pruned-emaonly.ckpt'
            else:
                config.model.params.ckpt_path = None
        else:
            config.model.params.ckpt_path = f'../{sd_path}'

        sd_model = instantiate_from_config(config.model)
        self.encoder_vq = sd_model.first_stage_model

        self.unet = UNetWrapper(sd_model.model, use_attn=True)
        if dataset == 'kitti':
            self.unet = UNetWrapper(sd_model.model, use_attn=True, base_size=384)
    
        del sd_model.cond_stage_model
        del self.encoder_vq.decoder
        del self.unet.unet.diffusion_model.out
        del self.encoder_vq.post_quant_conv.weight
        del self.encoder_vq.post_quant_conv.bias

        for param in self.encoder_vq.parameters():
            param.requires_grad = True
        
        self.text_adapter = TextAdapterRefer(text_dim=text_dim)
        self.alpha = nn.Parameter(torch.ones(text_dim) * 1e-4)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        if caption_aggregation:                
            class_embeddings = torch.load(f'{dataset}_class_embeddings_my_captions.pth', map_location=device)
            #class_embeddings_list = [value['class_embeddings'] for key, value in class_embeddings.items()]
            #stacked_embeddings = torch.stack(class_embeddings_list, dim=0)
            #class_embeddings = torch.mean(stacked_embeddings, dim=0).unsqueeze(0)
            
            if 'aggregated' in class_embeddings:
                class_embeddings = class_embeddings['aggregated'] 
            else:
                clip_model = FrozenCLIPEmbedder(max_length=40,pool=False).to(device)
                class_embeddings_new = [clip_model.encode(value['caption'][0]) for key, value in class_embeddings.items()]
                class_embeddings_new = torch.mean(torch.stack(class_embeddings_new, dim=0), dim=0)
                class_embeddings['aggregated'] = class_embeddings_new
                torch.save(class_embeddings, f'{dataset}_class_embeddings_my_captions.pth')
                class_embeddings = class_embeddings['aggregated']
            self.register_buffer('class_embeddings', class_embeddings)
        else:
            self.class_embeddings = torch.load(f'{dataset}_class_embeddings_my_captions.pth', map_location=device)

            self.clip_model = FrozenCLIPEmbedder(max_length=40,pool=False)
            for param in self.clip_model.parameters():
                param.requires_grad = True
                
        #if dataset == 'kitti':    
        #    self.text_adapter_ = TextAdapterRefer(text_dim=text_dim)
        #    self.gamma_ = nn.Parameter(torch.ones(text_dim) * 1e-4)
        
        self.caption_aggregation = caption_aggregation
        self.dataset = dataset
        
    def _init_weights(self, m):
        if isinstance(m, (nn.Conv2d, nn.Linear)):
            trunc_normal_(m.weight, std=.02)
            nn.init.constant_(m.bias, 0)

    def forward_features(self, feats):
        x =  self.ldm_to_net[0](feats[0])
        for i in range(3):
            if i > 0:
                x = x + self.ldm_to_net[i](feats[i])
            x = self.layers[i](x)
            x = self.upsample_layers[i](x)
        return self.out_conv(x)

    def forward(self, x, class_ids=None, img_paths=None):
        latents = self.encoder_vq.encode(x).mode()
        
        # add division by std
        if self.dataset == 'nyu':
            latents = latents / 5.07543
        elif self.dataset == 'kitti':
            latents = latents / 4.6211
        else:
            print('Please calculate the STD for the dataset!')

        if class_ids is not None:
            if self.caption_aggregation:
                class_embeddings = self.class_embeddings[[0]*len(class_ids.tolist())]#[class_ids.tolist()]
            else:
                class_embeddings = []
                
                for img_path in img_paths:
                    class_embeddings.extend([value['caption'][0] for key, value in self.class_embeddings.items() if key in img_path.replace('//', '/')])
                    
                class_embeddings = self.clip_model.encode(class_embeddings)
        else:
            class_embeddings = self.class_embeddings
        
        c_crossattn = self.text_adapter(latents, class_embeddings, self.alpha)
        t = torch.ones((x.shape[0],), device=x.device).long()

        #if self.dataset == 'kitti':
        #    c_crossattn_last = self.text_adapter_(latents, class_embeddings, self.gamma_)
        #    outs = self.unet(latents, t, c_crossattn=[c_crossattn, c_crossattn_last])
        #else:
        outs = self.unet(latents, t, c_crossattn=[c_crossattn])
        outs = self.aggregation(outs)
        
        feats = [outs[0], outs[1], torch.cat([outs[2], F.interpolate(outs[3], scale_factor=2)], dim=1)]
        x = torch.cat([self.layer1(feats[0]), self.layer2(feats[1]), feats[2]], dim=1)
        return self.out_layer(x)
    
    def get_latent(self, x):
        return self.encoder_vq.encode(x).mode()


class EVPDepth(PreTrainedModel):
    config_class = EVPConfig
    def __init__(self, config, caption_aggregation=True):
        super().__init__(config)
        args = config
        self.max_depth = args.max_depth
        self.min_depth = args.min_depth_eval

        embed_dim = 192
        
        channels_in = embed_dim*8
        channels_out = embed_dim

        if args.dataset == 'nyudepthv2':
            self.encoder = EVPDepthEncoder(out_dim=channels_in, dataset='nyu', caption_aggregation=caption_aggregation)
        else:
            self.encoder = EVPDepthEncoder(out_dim=channels_in, dataset='kitti', caption_aggregation=caption_aggregation)
            
        self.decoder = Decoder(channels_in, channels_out, args)
        self.decoder.init_weights()
        self.mViT = False
        self.custom = False
        
        
        if not self.mViT and not self.custom:
            n_bins = 64
            bin_embedding_dim = 128
            num_out_features = [32, 32, 32, 192]
            min_temp = 0.0212
            max_temp = 50
            btlnck_features = 256
            n_attractors = [16, 8, 4, 1]
            attractor_alpha = 1000
            attractor_gamma = 2
            attractor_kind = "mean"
            attractor_type = "inv"
            self.bin_centers_type = "softplus"
                        
            self.bottle_neck = nn.Sequential(
                nn.Conv2d(channels_in, btlnck_features, kernel_size=3, stride=1, padding=1),
                nn.ReLU(inplace=False),
                nn.Conv2d(btlnck_features, btlnck_features, kernel_size=3, stride=1, padding=1))
            
    
            for m in self.bottle_neck.modules():
                if isinstance(m, nn.Conv2d):
                    normal_init(m, std=0.001, bias=0)
                        
            
            SeedBinRegressorLayer = SeedBinRegressorUnnormed
            Attractor = AttractorLayerUnnormed
            self.seed_bin_regressor = SeedBinRegressorLayer(
                btlnck_features, n_bins=n_bins, min_depth=self.min_depth, max_depth=self.max_depth)
            self.seed_projector = Projector(btlnck_features, bin_embedding_dim)
            self.projectors = nn.ModuleList([
                Projector(num_out, bin_embedding_dim)
                for num_out in num_out_features
            ])
            self.attractors = nn.ModuleList([
                Attractor(bin_embedding_dim, n_bins, n_attractors=n_attractors[i], min_depth=self.min_depth, max_depth=self.max_depth,
                          alpha=attractor_alpha, gamma=attractor_gamma, kind=attractor_kind, attractor_type=attractor_type)
                for i in range(len(num_out_features))
            ])
            
            last_in = 192 + 1
            self.conditional_log_binomial = ConditionalLogBinomial(
                last_in, bin_embedding_dim, n_classes=n_bins, min_temp=min_temp, max_temp=max_temp)
        elif self.mViT and not self.custom:
            n_bins = 256
            self.adaptive_bins_layer = mViT(192, n_query_channels=192, patch_size=16,
                                            dim_out=n_bins,
                                            embedding_dim=192, norm='linear')
            self.conv_out = nn.Sequential(nn.Conv2d(192, n_bins, kernel_size=1, stride=1, padding=0),
                                          nn.Softmax(dim=1))
                             
    
    def forward(self, image, class_ids=None, img_paths=None):
        
        #image = transform(image).unsqueeze(0)
        shape = image.shape
        image = torch.nn.functional.interpolate(image, (440,480), mode='bilinear', align_corners=True)
        x = F.pad(image, (0, 0, 40, 0))

        b, c, h, w = x.shape
        x = x*2.0 - 1.0  # normalize to [-1, 1]
        if h == 480 and w == 480:
            new_x = torch.zeros(b, c, 512, 512, device=x.device)
            new_x[:, :, 0:480, 0:480] = x
            x = new_x
        elif h==352 and w==352:
            new_x = torch.zeros(b, c, 384, 384, device=x.device)
            new_x[:, :, 0:352, 0:352] = x
            x = new_x
        elif h == 512 and w == 512:
            pass
        else:
            print(h,w)
            raise NotImplementedError
        conv_feats = self.encoder(x, class_ids, img_paths)

        if h == 480 or h == 352:
            conv_feats = conv_feats[:, :, :-1, :-1]      
        
        self.decoder.remove_hooks()
        out_depth, out, x_blocks = self.decoder([conv_feats])
        
        if not self.mViT and not self.custom:
            x = self.bottle_neck(conv_feats)
            _, seed_b_centers = self.seed_bin_regressor(x)
            
            if self.bin_centers_type == 'normed' or self.bin_centers_type == 'hybrid2':
                b_prev = (seed_b_centers - self.min_depth) / \
                    (self.max_depth - self.min_depth)
            else:
                b_prev = seed_b_centers
                
            prev_b_embedding = self.seed_projector(x)
            
            for projector, attractor, x in zip(self.projectors, self.attractors, x_blocks):
                b_embedding = projector(x)
                b, b_centers = attractor(
                    b_embedding, b_prev, prev_b_embedding, interpolate=True)
                b_prev = b.clone()
                prev_b_embedding = b_embedding.clone()
            
            rel_cond = torch.sigmoid(out_depth) * self.max_depth
            
            # concat rel depth with last. First interpolate rel depth to last size
            rel_cond = nn.functional.interpolate(
                rel_cond, size=out.shape[2:], mode='bilinear', align_corners=True)
            last = torch.cat([out, rel_cond], dim=1)
                        
            b_embedding = nn.functional.interpolate(
                b_embedding, last.shape[-2:], mode='bilinear', align_corners=True)
            x = self.conditional_log_binomial(last, b_embedding)
            
            # Now depth value is Sum px * cx , where cx are bin_centers from the last bin tensor
            b_centers = nn.functional.interpolate(
                b_centers, x.shape[-2:], mode='bilinear', align_corners=True)
            out_depth = torch.sum(x * b_centers, dim=1, keepdim=True)
                        
        elif self.mViT and not self.custom:
            bin_widths_normed, range_attention_maps = self.adaptive_bins_layer(out)
            out = self.conv_out(range_attention_maps)
            
            bin_widths = (self.max_depth - self.min_depth) * bin_widths_normed  # .shape = N, dim_out
            bin_widths = nn.functional.pad(bin_widths, (1, 0), mode='constant', value=self.min_depth)
            bin_edges = torch.cumsum(bin_widths, dim=1)

            centers = 0.5 * (bin_edges[:, :-1] + bin_edges[:, 1:])
            n, dout = centers.size()
            centers = centers.view(n, dout, 1, 1)

            out_depth = torch.sum(out * centers, dim=1, keepdim=True)
        else:
            out_depth = torch.sigmoid(out_depth) * self.max_depth
        
        pred = out_depth
        pred = pred[:,:,40:,:]
        pred = torch.nn.functional.interpolate(pred, shape[2:], mode='bilinear', align_corners=True)
        pred_d_numpy = pred.squeeze().detach().cpu().numpy()

        return pred_d_numpy


class Decoder(nn.Module):
    def __init__(self, in_channels, out_channels, args):
        super().__init__()
        self.deconv = args.num_deconv
        self.in_channels = in_channels
        
        embed_dim = 192
        
        channels_in = embed_dim*8
        channels_out = embed_dim
        
        self.deconv_layers, self.intermediate_results = self._make_deconv_layer(
            args.num_deconv,
            args.num_filters,
            args.deconv_kernels,
        )
        self.last_layer_depth = nn.Sequential(
            nn.Conv2d(channels_out, channels_out, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=False),
            nn.Conv2d(channels_out, 1, kernel_size=3, stride=1, padding=1))
            
        for m in self.last_layer_depth.modules():
            if isinstance(m, nn.Conv2d):
                normal_init(m, std=0.001, bias=0)
        
        conv_layers = []
        conv_layers.append(
            build_conv_layer(
                dict(type='Conv2d'),
                in_channels=args.num_filters[-1],
                out_channels=out_channels,
                kernel_size=3,
                stride=1,
                padding=1))
        conv_layers.append(
            build_norm_layer(dict(type='BN'), out_channels)[1])
        conv_layers.append(nn.ReLU())
        self.conv_layers = nn.Sequential(*conv_layers)
        
        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)

    def forward(self, conv_feats):
        out = self.deconv_layers(conv_feats[0])
        out = self.conv_layers(out)
        out = self.up(out)
        self.intermediate_results.append(out)
        out = self.up(out)
        out_depth = self.last_layer_depth(out)

        return out_depth, out, self.intermediate_results

    def _make_deconv_layer(self, num_layers, num_filters, num_kernels):
        """Make deconv layers."""
        
        layers = []
        in_planes = self.in_channels
        intermediate_results = []  # List to store intermediate feature maps

        for i in range(num_layers):
            kernel, padding, output_padding = \
                self._get_deconv_cfg(num_kernels[i])

            planes = num_filters[i]
            layers.append(
                build_upsample_layer(
                    dict(type='deconv'),
                    in_channels=in_planes,
                    out_channels=planes,
                    kernel_size=kernel,
                    stride=2,
                    padding=padding,
                    output_padding=output_padding,
                    bias=False))
            layers.append(nn.BatchNorm2d(planes))
            layers.append(nn.ReLU())
            in_planes = planes

            # Add a hook to store the intermediate result
            layers[-1].register_forward_hook(self._hook_fn(intermediate_results))

        return nn.Sequential(*layers), intermediate_results

    def _hook_fn(self, intermediate_results):
        def hook(module, input, output):
            intermediate_results.append(output)
        return hook
        
    def remove_hooks(self):
        self.intermediate_results.clear()
        
    def _get_deconv_cfg(self, deconv_kernel):
        """Get configurations for deconv layers."""
        if deconv_kernel == 4:
            padding = 1
            output_padding = 0
        elif deconv_kernel == 3:
            padding = 1
            output_padding = 1
        elif deconv_kernel == 2:
            padding = 0
            output_padding = 0
        else:
            raise ValueError(f'Not supported num_kernels ({deconv_kernel}).')

        return deconv_kernel, padding, output_padding

    def init_weights(self):
        """Initialize model weights."""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                normal_init(m, std=0.001, bias=0)
            elif isinstance(m, nn.BatchNorm2d):
                constant_init(m, 1)
            elif isinstance(m, nn.ConvTranspose2d):
                normal_init(m, std=0.001)