models.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchview import draw_graph
from typing import Type, Callable, List, Optional
from torch import Tensor
import numpy as np



def weights_init(m):
    if isinstance(m, nn.Conv2d):
        nn.init.kaiming_uniform_(m.weight)
    if isinstance(m, nn.ConvTranspose2d):
        nn.init.kaiming_uniform_(m.weight)


""" Entry Block """
class EntryConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=1, padding=0, partial_conv=False):
        super(EntryConv, self).__init__()
        self.entry = nn.Sequential(
            PartialConv2d(in_channels, out_channels, kernel_size=kernel_size, padding=padding, bias=False) if partial_conv else nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=padding),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.entry(x)



### Convolution Wrapper Classes

""" Partial Convolution """
class PartialConv2d(nn.Conv2d):
    """ Implementation of Partial-convolution (https://github.com/NVIDIA/partialconv)"""
    def __init__(self, *args, **kwargs):
        # whether the mask is multi-channel or not
        if 'multi_channel' in kwargs:
            self.multi_channel = kwargs['multi_channel']
            kwargs.pop('multi_channel')
        else:
            self.multi_channel = False
        if 'return_mask' in kwargs:
            self.return_mask = kwargs['return_mask']
            kwargs.pop('return_mask')
        else:
            self.return_mask = False
        super(PartialConv2d, self).__init__(*args, **kwargs)
        if self.multi_channel:
            self.weight_maskUpdater = torch.ones(self.out_channels, self.in_channels, self.kernel_size[0],
                                                 self.kernel_size[1])
        else:
            self.weight_maskUpdater = torch.ones(1, 1, self.kernel_size[0], self.kernel_size[1])
        self.slide_winsize = self.weight_maskUpdater.shape[1] * self.weight_maskUpdater.shape[2] * \
                             self.weight_maskUpdater.shape[3]
        self.last_size = (None, None, None, None)
        self.update_mask = None
        self.mask_ratio = None

    def forward(self, input, mask_in=None):
        assert len(input.shape) == 4
        if mask_in is not None or self.last_size != tuple(input.shape):
            self.last_size = tuple(input.shape)

            with torch.no_grad():
                if self.weight_maskUpdater.type() != input.type():
                    self.weight_maskUpdater = self.weight_maskUpdater.to(input)

                if mask_in is None:
                    # if mask is not provided, create a mask
                    if self.multi_channel:
                        mask = torch.ones(input.data.shape[0], input.data.shape[1], input.data.shape[2],
                                          input.data.shape[3]).to(input)
                    else:
                        mask = torch.ones(1, 1, input.data.shape[2], input.data.shape[3]).to(input)
                else:
                    mask = mask_in

                self.update_mask = F.conv2d(mask, self.weight_maskUpdater, bias=None, stride=self.stride,
                                            padding=self.padding, dilation=self.dilation, groups=1)

                self.mask_ratio = self.slide_winsize / (self.update_mask + 1e-8)
                # self.mask_ratio = torch.max(self.update_mask)/(self.update_mask + 1e-8)
                self.update_mask = torch.clamp(self.update_mask, 0, 1)
                self.mask_ratio = torch.mul(self.mask_ratio, self.update_mask)

        # if self.update_mask.type() != input.type() or self.mask_ratio.type() != input.type():
        #     self.update_mask.to(input)
        #     self.mask_ratio.to(input)

        raw_out = super(PartialConv2d, self).forward(torch.mul(input, mask) if mask_in is not None else input)

        if self.bias is not None:
            bias_view = self.bias.view(1, self.out_channels, 1, 1)
            output = torch.mul(raw_out - bias_view, self.mask_ratio) + bias_view
            output = torch.mul(output, self.update_mask)
        else:
            output = torch.mul(raw_out, self.mask_ratio)

        if self.return_mask:
            return output, self.update_mask
        else:
            return output

class SingleConv(nn.Module):
    def __init__(self, in_channels, out_channels, partial_conv=False):
        super(SingleConv, self).__init__()
        self.conv = nn.Sequential(
            PartialConv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False) if partial_conv else nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
    def forward(self, x):
        return self.conv(x)

class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels, partial_conv=False):
        super(DoubleConv, self).__init__()
        self.conv = nn.Sequential(
            PartialConv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False) if partial_conv else nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            PartialConv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False) if partial_conv else nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
            #nn.Dropout2d(p=0.3)
        )

    def forward(self, x):
        return self.conv(x)

class TransposeConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=2, stride=2, batchNorm=False, relu=False):
        super(TransposeConv, self).__init__()
        self.conv = nn.Sequential(
            nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride)
        )
        if batchNorm:
            self.conv.append(nn.BatchNorm2d(out_channels))
        if relu:
            self.conv.append(nn.ReLU(inplace=True))

    def forward(self, x):
        return self.conv(x)

class Upsample(nn.Module):
    def __init__(self, in_channels, out_channels, mode='bilinear', scale_factor=2, trainable=True, batchNorm=False, relu=False, align_corners=False):
        super(Upsample, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.mode = mode
        self.scale_factor = scale_factor
        self.align_corners = align_corners
        self.conv = nn.Sequential(
            nn.Upsample(scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners),
        )
        if trainable:
            self.conv.append(nn.Conv2d(self.in_channels, self.out_channels, kernel_size=1, stride=1))
        if batchNorm:
            self.conv.append(nn.BatchNorm2d(out_channels))
        if relu:
            self.conv.append(nn.ReLU(inplace=True))

    def forward(self, x):
        return self.conv(x)

""" Squeeze-and-Excitation block """
class SqueezeExcitation(nn.Module):
    def __init__(self, channels):
        super(SqueezeExcitation, self).__init__()
        self.conv = nn.Sequential(
            # Squeeze
            nn.AdaptiveAvgPool2d(1),
            # Excitation
            nn.Conv2d(channels, int(channels//2), 1, 1, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(int(channels//2), channels, 1, 1, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        calibrate = self.conv(x)
        return calibrate * x

# Sharp block


def get_kernel():
    """
    See https://setosa.io/ev/image-kernels/
    """

    k1 = np.array([[0.0625, 0.125, 0.0625],
                   [0.125,  0.25, 0.125],
                   [0.0625, 0.125, 0.0625]])
    
    # Sharpening Spatial Kernel, used in paper
    k2 = np.array([[-1, -1, -1],
                   [-1,  8, -1],
                   [-1, -1, -1]])
    
    k3 = np.array([[0, -1, 0],
                   [-1,  5, -1],
                   [0, -1, 0]])
    
    return k1, k2, k3


def build_sharp_blocks(in_channels):
    """
    Sharp Blocks
    """
    # Get kernel
    w, _, _ = get_kernel()    
    # Change dimension
    w = np.expand_dims(w, axis=0)
    # Expand dimension
    w = np.expand_dims(w, axis=0)
    # Repeat filter by in_channels times to get (H, W, in_channels)
    w = np.repeat(w, in_channels, axis=0)
    
    return w
class Sharping(nn.Module):
    def __init__(self, in_channels):
        super(Sharping, self).__init__()
        kernel_size = 3
        W = build_sharp_blocks(in_channels) 
        self.sb = nn.Conv2d(in_channels,in_channels, kernel_size, bias=False,groups=in_channels, padding='same')
        self.sb.weight =torch.nn.Parameter(torch.from_numpy(W).float())
        self.sb.requires_grad_=False
    
    def forward(self, x):
        sb =self.sb(x)
        return sb
        

## Attention blocks

class GatingSignal(nn.Module):
    def __init__(self, in_channels, out_channels, batch_norm=True):
        super(GatingSignal, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=(1, 1))
        self.batch_norm = batch_norm
        self.bn = nn.BatchNorm2d(out_channels)
        self.activation = nn.ReLU()

    def forward(self, x):
        x = self.conv(x)
        if self.batch_norm:
            x = self.bn(x)
        return self.activation(x)


class Attention_Gate(nn.Module):
    def __init__(self, in_channels):
        super(Attention_Gate, self).__init__()
        self.conv_theta_x = nn.Conv2d(
            in_channels, in_channels, kernel_size=(1, 1), stride=(2, 2)
        )
        self.conv_phi_g = nn.Conv2d(in_channels, in_channels, kernel_size=(1, 1))
        self.att = nn.Sequential(
            nn.ReLU(),
            nn.Conv2d(in_channels, 1, kernel_size=(1, 1)),
            nn.Sigmoid(),
            nn.Upsample(scale_factor=2),
        )

    def forward(self, x, gat):
        theta_x = self.conv_theta_x(x)
        phi_g = self.conv_phi_g(gat)
        res = torch.add(phi_g, theta_x)
        res = self.att(res)
        return torch.mul(res, x)


#### Model Architectures

""" UNet, Attention-UNet, Multitask Vegetation (uses UNet/Attention-Unet as backbone) """


class UNet(nn.Module):
    def __init__(
            self,
            in_channels=5,
            out_channels=4,
            upsample='transpose',
            features=[64, 128, 256, 512],
            use_se_block=False,
            partial_conv=False,
            entry_block=False,
            dropout=False,
            sar_channels=1,
            init_weights=False,
            sar_only=False,
            opt_only=False,
            seg_only=False,
            reg_only=False,
            backbone=False
        ):
        super(UNet, self).__init__()
        assert features[0] * 2**(len(features)-1) == features[-1]
        assert upsample in ['nearest','bilinear','transpose']
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.upsample = upsample
        self.features = features
        self.use_se_block = use_se_block
        self.partial_conv = partial_conv
        self.entry_block = entry_block
        self.dropout = dropout
        self.sar_channels = sar_channels
        self.init_weights = init_weights
        self.sar_only = sar_only
        self.opt_only = opt_only
        self.seg_only = seg_only
        self.reg_only = reg_only
        self.backbone = backbone

        features = features[:-1]
        # before unet
        if entry_block:
            self.opt_block = EntryConv(in_channels-sar_channels, features[0]//2, kernel_size=1, partial_conv=partial_conv) if not self.opt_only else EntryConv(in_channels, out_channels=features[0], kernel_size=1, partial_conv=partial_conv)
            self.sar_block = EntryConv(sar_channels, features[0]//2, kernel_size=5, padding=2, partial_conv=partial_conv) if not self.sar_only else EntryConv(in_channels, out_channels=features[0], kernel_size=5, padding=2, partial_conv=partial_conv)
            in_channels = features[0]
        if self.use_se_block:
            self.first_se = SqueezeExcitation(in_channels)
        # unet
        self.downwards = nn.ModuleList()
        self.skip = nn.ModuleList()
        self.upwards = nn.ModuleList()
        self.bottleneck = nn.ModuleList()
        self.pool = nn.MaxPool2d(2,2)
        self.projection = nn.Conv2d(features[0], out_channels, kernel_size=1)
        # contracting path
        for f in features:
            self.downwards.append(DoubleConv(in_channels, f, partial_conv=partial_conv))
            in_channels = f
            if self.use_se_block:
                self.skip.append(SqueezeExcitation(in_channels))
        # bottleneck
        self.bottleneck.append(DoubleConv(in_channels, features[-1]*2, partial_conv=partial_conv))
        in_channels = features[-1]*2
        if self.use_se_block:
            self.bottleneck.append(SqueezeExcitation(in_channels))
        if self.dropout:
            self.bottleneck.append(nn.Dropout2d(p=0.3))
        # expansive path
        for f in reversed(features):
            if upsample == 'transpose':
                self.upwards.append(nn.ConvTranspose2d(in_channels, f, 2, 2))
            else:
                self.upwards.append(Upsample(in_channels=in_channels, out_channels=f, mode=upsample, scale_factor=2, align_corners=False))
            self.upwards.append(DoubleConv(in_channels, f, partial_conv=partial_conv))
            in_channels = f

        if init_weights:
            self.apply(weights_init)

    def forward(self, x):
        skips = []
        #  ------------- Entry -------------  #
        if self.entry_block:
            if self.opt_only:
                x = self.opt_block(x)
            elif self.sar_only:
                x = self.sar_block(x)
            else:
                x_opt, x_sar = x[:,:-self.sar_channels], x[:, -self.sar_channels:]
                x_opt = self.opt_block(x_opt)
                x_sar = self.sar_block(x_sar)
                x = torch.cat([x_opt, x_sar],dim=1)
        if self.use_se_block:
            x = self.first_se(x)
        #  ------------- Contracting Path -------------  #
        for i, conv in enumerate(self.downwards):
            x = conv(x)
            skips.append(self.skip[i](x) if self.use_se_block else x)
            x = self.pool(x)
        skips.reverse()
        #  ------------- Bottleneck -------------  #
        for conv in self.bottleneck:
            x = conv(x)
        #  ------------- Expansive Path -------------  #
        for i in range(0,len(self.upwards),2):
            x = self.upwards[i](x)
            x = torch.cat((skips[i//2], x), dim=1)
            x = self.upwards[i+1](x)
        if self.backbone:
            pass
        else:
            x = self.projection(x)
        return x


    def to_out(self):
        return (
            f'{self.__class__.__name__}'
            f'(in_channels={self.in_channels}, '
            f'out_channels={self.out_channels}, '
            f'upsample={self.upsample}, '
            f'features={self.features}, '
            f'use_se_block={self.use_se_block}, '
            f'partial_conv={self.partial_conv}, '
            f'entry_block={self.entry_block}, '
            f'dropout={self.dropout}, '
            f'sar_channels={self.sar_channels}, '
            f'init_weights={self.init_weights}, '
            f'sar_only={self.sar_only}, '
            f'opt_only={self.opt_only})'
        )
    def visualize_graph(self, input_size):
        model_graph = draw_graph(self, input_size=input_size, expand_nested = True)
        return model_graph.visual_graph


class AttentionUNet(nn.Module):
    def __init__(
            self,
            in_channels=5,
            out_channels=1,
            upsample='transpose',
            features=[128, 256, 512],
            use_se_block=False,
            use_sharp_block=False,
            partial_conv=False,
            entry_block=False,
            dropout=False,
            sar_channels=1,
            init_weights=False,
            sar_only=False,
            opt_only=False,
            backbone=False,
            seg_only=False,
            reg_only=False

        ):
        super(AttentionUNet, self).__init__()
        assert features[0] * 2**(len(features)-1) == features[-1]
        assert upsample in ['nearest','bilinear','transpose']
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.upsample = upsample
        self.features = features
        self.use_se_block = use_se_block
        self.use_sharp_block = use_sharp_block
        self.partial_conv = partial_conv
        self.entry_block = entry_block
        self.dropout = dropout
        self.sar_channels = sar_channels
        self.init_weights = init_weights
        self.sar_only = sar_only
        self.opt_only = opt_only
        self.seg_only = seg_only
        self.reg_only = reg_only
        self.backbone = backbone

        features = features[:-1]
        # before unet
        if entry_block:
            self.opt_block = EntryConv(in_channels-sar_channels, features[0]//2, kernel_size=3, padding=1,partial_conv=partial_conv) if not self.opt_only else EntryConv(in_channels, out_channels=features[0], kernel_size=1, partial_conv=partial_conv)
            self.sar_block = EntryConv(sar_channels, features[0]//2, kernel_size=5, padding=2, partial_conv=partial_conv) if not self.sar_only else EntryConv(in_channels, out_channels=features[0], kernel_size=5, padding=2, partial_conv=partial_conv)
            in_channels = features[0]
        if self.use_se_block:
            self.first_se = SqueezeExcitation(in_channels)
        # unet
        self.downwards = nn.ModuleList()
        self.skip = nn.ModuleList()
        self.upwards = nn.ModuleList()
        self.bottleneck = nn.ModuleList()
        self.pool = nn.MaxPool2d(2,2)
        self.projection = nn.Conv2d(features[0], out_channels, kernel_size=1)
        # contracting path
        for f in features:
            self.downwards.append(DoubleConv(in_channels, f, partial_conv=partial_conv))
            in_channels = f
            if self.use_se_block: 
                self.skip.append(SqueezeExcitation(in_channels))
            if self.use_sharp_block:
                self.skip.append(Sharping(in_channels))
        # bottleneck
        self.bottleneck.append(DoubleConv(in_channels, features[-1]*2, partial_conv=partial_conv))
        in_channels = features[-1]*2
        if self.use_se_block:
            self.bottleneck.append(SqueezeExcitation(in_channels))
        if self.dropout:
            self.bottleneck.append(nn.Dropout2d(p=0.3))
        # expansive path
        for f in reversed(features):
            if upsample == 'transpose':
                self.upwards.append(nn.ConvTranspose2d(in_channels, f, 2, 2))
            else:
                self.upwards.append(Upsample(in_channels=in_channels, out_channels=f, mode=upsample, scale_factor=2, align_corners=False))
            self.upwards.append(GatingSignal(in_channels, f))
            self.upwards.append(Attention_Gate(f))
            self.upwards.append(DoubleConv(in_channels, f,  partial_conv=partial_conv))
            in_channels = f
        if init_weights:
            self.apply(weights_init)

    def forward(self, x):
        skips = []
        #  ------------- Entry -------------  #
        if self.entry_block:
            if self.opt_only:
                x = self.opt_block(x)
            elif self.sar_only:
                x = self.sar_block(x)
            else:
                x_opt, x_sar = x[:, :-self.sar_channels], x[:, -self.sar_channels:]
                x_opt = self.opt_block(x_opt)
                x_sar = self.sar_block(x_sar)
                x = torch.cat([x_opt, x_sar],dim=1)
        if self.use_se_block:
            x = self.first_se(x)
        #  ------------- Contracting Path -------------  #
        for i, conv in enumerate(self.downwards):
            c = conv(x)
            skips.append(self.skip[i](c) if self.use_sharp_block else c)
            #skips.append(c)
            x = self.pool(c)
        skips.reverse()
        #  ------------- Bottleneck -------------  #
        for conv in self.bottleneck:
            x = conv(x)
        #  ------------- Expansive Path -------------  #

        for i in range(0,len(self.upwards),4):
            x1 = self.upwards[i](x)
            gat = self.upwards[i+1](x)
            att = self.upwards[i+2](skips[i//4], gat) #//2
            x = torch.cat((x1, att), dim=1)
            x = self.upwards[i+3](x)
        if self.backbone:
            pass
        else:
            x = self.projection(x)
        return x

    def to_out(self):
        return (
            f'{self.__class__.__name__}'
            f'(in_channels={self.in_channels}, '
            f'out_channels={self.out_channels}, '
            f'upsample={self.upsample}, '
            f'features={self.features}, '
            f'use_se_block={self.use_se_block}, '
            f'partial_conv={self.partial_conv}, '
            f'entry_block={self.entry_block}, '
            f'dropout={self.dropout}, '
            f'sar_channels={self.sar_channels}, '
            f'init_weights={self.init_weights}, '
            f'sar_only={self.sar_only}, '
            f'opt_only={self.opt_only})'
        )
    def visualize_graph(self, input_size):
        model_graph = draw_graph(self, input_size=input_size, expand_nested = True)
        return model_graph.visual_graph


class Veg_MT(nn.Module):
    def __init__(self,
            in_channels=5,
            out_channels=1,
            upsample='bilinear',
            features=[64, 128, 256, 512],
            use_se_block=False,
            partial_conv=False,
            entry_block=False,
            dropout=False,
            sar_channels=1,
            init_weights=False,
            sar_only=False,
            opt_only=False,
            compute_seg=True,
            compute_reg=True,
            seg_only=False,
            reg_only=False,
            backbone=AttentionUNet
            #backbone=UNet
        ):
        super(Veg_MT, self).__init__()
        assert features[0] * 2**(len(features)-1) == features[-1]
        assert upsample in ['nearest','bilinear','transpose']
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.upsample = upsample
        self.features = features
        self.entry_block = entry_block
        self.use_se_block = use_se_block
        self.partial_conv = partial_conv
        self.dropout = dropout
        self.sar_channels = sar_channels
        self.init_weights = init_weights
        self.sar_only = sar_only
        self.opt_only = opt_only
        self.backbone = backbone(in_channels=self.in_channels, out_channels=self.out_channels, upsample=self.upsample, features=self.features,
                                 entry_block=self.entry_block, dropout=self.dropout, sar_channels=self.sar_channels, init_weights=self.init_weights,
                                 sar_only=self.sar_only, opt_only=self.opt_only, backbone=True)
        self.compute_seg = compute_seg
        self.compute_reg = compute_reg
        self.seg_only = seg_only
        self.reg_only = reg_only



        # backbone
        backbone_out_features = features[0]

        # Seg branch
        seg_channels = 0
        if self.compute_seg:
            seg_channels = 1
            self.seg_module = torch.nn.Sequential(
                torch.nn.Conv2d(backbone_out_features, backbone_out_features, 3, padding=1),
                torch.nn.BatchNorm2d(backbone_out_features),
                torch.nn.ReLU(),
                torch.nn.Conv2d(backbone_out_features, seg_channels, 1),
                #torch.nn.Sigmoid(),)
            )

        # Regression branch
        if self.compute_reg:
            reg_channels = 1
            self.reg_module = torch.nn.Sequential(
                torch.nn.Conv2d(backbone_out_features+ seg_channels, backbone_out_features, 3, padding=1), #
                torch.nn.BatchNorm2d(backbone_out_features),
                torch.nn.ReLU(),
                torch.nn.Conv2d(backbone_out_features, reg_channels, 1),
            )

    def forward(self, x):

        DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

        # --- Extract features for every pixel of the image with a U-Net --- #
        backbone_features = self.backbone(x)

        if self.compute_seg:
            # --- Output a segmentation of the image --- #
            seg = self.seg_module(backbone_features)

        if (self.compute_seg and self.compute_reg):
                m = torch.nn.Sigmoid()
                seg_to_cat = seg.clone().detach()
                seg_to_cat= m(seg_to_cat)
                backbone_features = torch.cat([backbone_features, seg_to_cat], dim=1)  # Add seg to image features

        if self.compute_reg:
            # --- Output a height estimation of the image --- #
            reg = self.reg_module(backbone_features)  # Outputs c_0, c_2 values in [-2, 2]

        if (self.compute_seg and self.compute_seg):
            output=torch.cat([seg, reg], dim=1)
        elif self.compute_reg:
            output = reg
        else:
            output = seg


        return output




    def to_out(self):
        return (
            f'{self.__class__.__name__}'
            f'(in_channels={self.in_channels}, '
            f'out_channels={self.out_channels}, '
            f'upsample={self.upsample}, '
            f'features={self.features}, '
            f'use_se_block={self.use_se_block}, '
            f'partial_conv={self.partial_conv}, '
            f'entry_block={self.entry_block}, '
            f'dropout={self.dropout}, '
            f'sar_channels={self.sar_channels}, '
            f'init_weights={self.init_weights}, '
            f'sar_only={self.sar_only}, '
            f'opt_only={self.opt_only})'
        )
    def visualize_graph(self, input_size):
        model_graph = draw_graph(self, input_size=input_size, expand_nested = True)
        return model_graph.visual_graph