third_parties/dsine/models/conv_encoder_decoder/submodules.py

""" submodules
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import geffnet
import os

INPUT_CHANNELS_DICT = {
    0: [1280, 112, 40, 24, 16],
    1: [1280, 112, 40, 24, 16],
    2: [1408, 120, 48, 24, 16],
    3: [1536, 136, 48, 32, 24],
    4: [1792, 160, 56, 32, 24],
    5: [2048, 176, 64, 40, 24],
    6: [2304, 200, 72, 40, 32],
    7: [2560, 224, 80, 48, 32]
}


class Encoder(nn.Module):
    def __init__(self, B=5, pretrained=True):
        """ e.g. B=5 will return EfficientNet-B5
        """
        super(Encoder, self).__init__()
        basemodel_name = 'tf_efficientnet_b%s_ap' % B
        basemodel = geffnet.create_model(basemodel_name, pretrained=False)
        if pretrained:
            ckpt_path = 'ckpts/dsine/tf_efficientnet_b5_ap-9e82fae8.pth'
            if os.path.exists(ckpt_path):
                print(f'Loading geffnet from {ckpt_path}')
                basemodel.load_state_dict(torch.load(ckpt_path, map_location='cpu'))
            else:
                print(f'Warning: {ckpt_path} not found. Initializing randomly.')

        # Remove last layer
        basemodel.global_pool = nn.Identity()
        basemodel.classifier = nn.Identity()
        self.original_model = basemodel

    def forward(self, x):
        features = [x]
        for k, v in self.original_model._modules.items():
            if (k == 'blocks'):
                for ki, vi in v._modules.items():
                    features.append(vi(features[-1]))
            else:
                features.append(v(features[-1]))
        return features


class UpSampleBN(nn.Module):
    def __init__(self, skip_input, output_features, align_corners=True):
        super(UpSampleBN, self).__init__()
        self._net = nn.Sequential(nn.Conv2d(skip_input, output_features, kernel_size=3, stride=1, padding=1),
                                  nn.BatchNorm2d(output_features),
                                  nn.LeakyReLU(),
                                  nn.Conv2d(output_features, output_features, kernel_size=3, stride=1, padding=1),
                                  nn.BatchNorm2d(output_features),
                                  nn.LeakyReLU())
        self.align_corners = align_corners

    def forward(self, x, concat_with):
        up_x = F.interpolate(x, size=[concat_with.size(2), concat_with.size(3)], mode='bilinear', align_corners=self.align_corners)
        f = torch.cat([up_x, concat_with], dim=1)
        return self._net(f)


class Conv2d_WS(nn.Conv2d):
    """ weight standardization
    """ 
    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 padding=0, dilation=1, groups=1, bias=True):
        super(Conv2d_WS, self).__init__(in_channels, out_channels, kernel_size, stride,
                 padding, dilation, groups, bias)

    def forward(self, x):
        weight = self.weight
        weight_mean = weight.mean(dim=1, keepdim=True).mean(dim=2,
                                  keepdim=True).mean(dim=3, keepdim=True)
        weight = weight - weight_mean
        std = weight.view(weight.size(0), -1).std(dim=1).view(-1, 1, 1, 1) + 1e-5
        weight = weight / std.expand_as(weight)
        return F.conv2d(x, weight, self.bias, self.stride, self.padding, self.dilation, self.groups)


class UpSampleGN(nn.Module):
    """ UpSample with GroupNorm
    """
    def __init__(self, skip_input, output_features, align_corners=True):
        super(UpSampleGN, self).__init__()
        self._net = nn.Sequential(Conv2d_WS(skip_input, output_features, kernel_size=3, stride=1, padding=1),
                                  nn.GroupNorm(8, output_features),
                                  nn.LeakyReLU(),
                                  Conv2d_WS(output_features, output_features, kernel_size=3, stride=1, padding=1),
                                  nn.GroupNorm(8, output_features),
                                  nn.LeakyReLU())
        self.align_corners = align_corners

    def forward(self, x, concat_with):
        up_x = F.interpolate(x, size=[concat_with.size(2), concat_with.size(3)], mode='bilinear', align_corners=self.align_corners)
        f = torch.cat([up_x, concat_with], dim=1)
        return self._net(f)


def upsample_via_bilinear(out, up_mask, downsample_ratio):
    """ bilinear upsampling (up_mask is a dummy variable)
    """
    return F.interpolate(out, scale_factor=downsample_ratio, mode='bilinear', align_corners=False)


def upsample_via_mask(out, up_mask, downsample_ratio, padding='zero'):
    """ convex upsampling
    """
    # out: low-resolution output (B, o_dim, H, W)
    # up_mask: (B, 9*k*k, H, W)
    k = downsample_ratio

    B, C, H, W = out.shape
    up_mask = up_mask.view(B, 1, 9, k, k, H, W)
    up_mask = torch.softmax(up_mask, dim=2)             # (B, 1, 9, k, k, H, W)

    if padding == 'zero':
        # with zero padding
        up_out = F.unfold(out, [3, 3], padding=1)       # (B, 2, H, W) -> (B, 2 X 3*3, H*W)
    elif padding == 'replicate':
        # with replicate padding
        out = F.pad(out, pad=(1,1,1,1), mode='replicate')
        up_out = F.unfold(out, [3, 3], padding=0)           # (B, C, H, W) -> (B, C X 3*3, H*W)
    else:
        raise Exception('invalid padding for convex upsampling')

    up_out = up_out.view(B, C, 9, 1, 1, H, W)           # (B, C, 9, 1, 1, H, W)

    up_out = torch.sum(up_mask * up_out, dim=2)         # (B, C, k, k, H, W)
    up_out = up_out.permute(0, 1, 4, 2, 5, 3)           # (B, C, H, k, W, k)
    return up_out.reshape(B, C, k*H, k*W)               # (B, C, kH, kW)


def dummy_activation(out):
    return out


def get_prediction_head(input_dim, hidden_dim, output_dim):
    return nn.Sequential(
        nn.Conv2d(input_dim, hidden_dim, 3, padding=1), 
        nn.ReLU(inplace=True),
        nn.Conv2d(hidden_dim, hidden_dim, 1), 
        nn.ReLU(inplace=True),
        nn.Conv2d(hidden_dim, output_dim, 1),
    )


class Decoder(nn.Module):
    def __init__(self, num_classes=2, 
                 B=5, NF=2048, BN=True,
                 down=2, learned_upsampling=True,
                 activation_fn=dummy_activation):
        super(Decoder, self).__init__()
        input_channels = INPUT_CHANNELS_DICT[B]

        # use BN or GN
        UpSample = UpSampleBN if BN else UpSampleGN

        features = NF
        self.conv2 = nn.Conv2d(input_channels[0], features, kernel_size=1, stride=1, padding=0)
        self.up1 = UpSample(skip_input=features // 1 + input_channels[1], output_features=features // 2)
        self.up2 = UpSample(skip_input=features // 2 + input_channels[2], output_features=features // 4)

        if down == 8:
            i_dim = features // 4
        elif down == 4:
            self.up3 = UpSample(skip_input=features // 4 + input_channels[3], output_features=features // 8)
            i_dim = features // 8
        elif down == 2:
            self.up3 = UpSample(skip_input=features // 4 + input_channels[3], output_features=features // 8)
            self.up4 = UpSample(skip_input=features // 8 + input_channels[4], output_features=features // 16)
            i_dim = features // 16
        else:
            raise Exception('invalid downsampling ratio')

        self.downsample_ratio = down
        self.output_dim = num_classes

        h_dim = 128
        self.pred_head = get_prediction_head(i_dim, h_dim, num_classes)

        if learned_upsampling:
            h_dim = 128
            self.mask_head = get_prediction_head(i_dim, h_dim, 9*self.downsample_ratio*self.downsample_ratio)
            self.upsample_fn = upsample_via_mask
        else:
            self.mask_head = lambda a: None
            self.upsample_fn = upsample_via_bilinear

        self.activation_fn = activation_fn


    def forward(self, features):
        x_block0, x_block1, x_block2, x_block3, x_block4 = features[4], features[5], features[6], features[8], features[11]

        x_d0 = self.conv2(x_block4)
        x_d1 = self.up1(x_d0, x_block3)

        if self.downsample_ratio == 8:
            x_feat = self.up2(x_d1, x_block2)
        elif self.downsample_ratio == 4:
            x_d2 = self.up2(x_d1, x_block2)
            x_feat = self.up3(x_d2, x_block1)
        elif self.downsample_ratio == 2:
            x_d2 = self.up2(x_d1, x_block2)
            x_d3 = self.up3(x_d2, x_block1)
            x_feat = self.up4(x_d3, x_block0)

        out = self.activation_fn(self.pred_head(x_feat))
        mask = self.mask_head(x_feat)
        up_out = self.upsample_fn(out, mask, self.downsample_ratio)
        return up_out