GP_UNet.py

# Adapted from https://github.com/soumickmj/FTSuperResDynMRI/blob/main/models/unet2d.py

import torch
from torch import nn
import torch.nn.functional as F
import torchcomplex.nn.functional as cF



__author__ = "Soumick Chatterjee"
__copyright__ = "Copyright 2020, Faculty of Computer Science, Otto von Guericke University Magdeburg, Germany"
__credits__ = ["Soumick Chatterjee"]
__license__ = "GPL"
__version__ = "1.0.0"
__maintainer__ = "Soumick Chatterjee"
__email__ = "soumick.chatterjee@ovgu.de"
__status__ = "Production"


class GP_UNet(nn.Module):
    """
    Implementation of
    U-Net: Convolutional Networks for Biomedical Image Segmentation
    (Ronneberger et al., 2015)
    https://arxiv.org/abs/1505.04597

    Using the default arguments will yield the exact version used
    in the original paper

    Args:
        in_channels (int): number of input channels
        n_classes (int): number of output channels
        depth (int): depth of the network
        wf (int): number of filters in the first layer is 2**wf
        padding (bool): if True, apply padding such that the input shape
                        is the same as the output.
                        This may introduce artifacts
        batch_norm (bool): Use BatchNorm after layers with an
                            activation function
        up_mode (str): one of 'upconv' or 'upsample'.
                        'upconv' will use transposed convolutions for
                        learned upsampling.
                        'upsample_Bi' will use bilinear upsampling.
                        'upsample_Sinc' will use sinc upsampling.
    """
    def __init__(self, in_channels=1, n_classes=1, depth=3, wf=6, padding=True,
                 batch_norm=False, up_mode='upconv', dropout=False, Relu = "Relu", out_act="None"): #dropout=False
        super(GP_UNet, self).__init__()
        assert up_mode in ('upconv', 'bilinear', 'sinc', "upsample_Sinc")
        assert out_act in ("softmax", "None", "sigmoid", "relu")
        self.padding = padding
        self.depth = depth
        self.Relu = Relu
        self.dropout = nn.Dropout2d() if dropout else nn.Sequential()
        prev_channels = in_channels
        self.down_path = nn.ModuleList()
        for i in range(depth):
            self.down_path.append(UNetConvBlock(prev_channels, 2**(wf+i),
                                                padding, batch_norm, Relu))
            prev_channels = 2**(wf+i)

        self.up_path = nn.ModuleList()
        for i in reversed(range(depth - 1)):
            self.up_path.append(UNetUpBlock(prev_channels, 2**(wf+i), up_mode,
                                            padding, batch_norm, Relu))
            prev_channels = 2**(wf+i)

        if out_act == "softmax":
            self.last = nn.Sequential(
                      nn.Conv2d(prev_channels, n_classes, kernel_size=1),
                      nn.Softmax2d()
                    )
        
        elif out_act == "sigmoid":
            self.last = nn.Sequential(
                      nn.Conv2d(prev_channels, n_classes, kernel_size=1),
                      nn.Sigmoid()
                    )
        
        elif out_act == "relu":
            self.last = nn.Sequential(
                      nn.Conv2d(prev_channels, n_classes, kernel_size=1),
                      nn.ReLU()
                    )

        else:
            self.last = nn.Conv2d(prev_channels, n_classes, kernel_size=1)
        

        ### For Classification, following Florian's GP-UNet
        self.GMP = nn.AdaptiveMaxPool2d((1, 1))

    def forward(self, x):
        blocks = []
        for i, down in enumerate(self.down_path):
            x = down(x)
            if i != len(self.down_path)-1:
                blocks.append(x)
                #x = nn.AvgPool2d(x, 2)
                x = F.avg_pool2d(x, 2)
        x = self.dropout(x)

        for i, up in enumerate(self.up_path):
            x = up(x, blocks[-i-1])

        if self.training:
            x = self.GMP(x)
            return self.last(x).view(x.shape[0],-1)
        else:
            mask = self.last(x)
            x = self.GMP(x)
            pred = self.last(x).view(x.shape[0],-1)
            return pred, mask

class UNetConvBlock(nn.Module):
    def __init__(self, in_size, out_size, padding, batch_norm, Relu):
        super(UNetConvBlock, self).__init__()
        block = []

        block.append(nn.Conv2d(in_size, out_size, kernel_size=3,
                               padding=int(padding)))
        if Relu == "Relu":
            block.append(nn.ReLU())
        else:
            block.append(nn.PReLU())

        if batch_norm:
            block.append(nn.BatchNorm2d(out_size))

        block.append(nn.Conv2d(out_size, out_size, kernel_size=3,
                               padding=int(padding)))
        
        if Relu == "Relu":
            block.append(nn.ReLU())
        else:
            block.append(nn.PReLU())

        if batch_norm:
            block.append(nn.BatchNorm2d(out_size))

        self.block = nn.Sequential(*block)

    def forward(self, x):
        out = self.block(x)
        return out


class UNetUpBlock(nn.Module):
    def __init__(self, in_size, out_size, up_mode, padding, batch_norm, Relu):
        super(UNetUpBlock, self).__init__()

        self.up_mode = up_mode

        if up_mode == 'upconv':
            self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2,
                                         stride=2)
        elif up_mode == 'bilinear':
            self.up = nn.Sequential(nn.Upsample(mode='bilinear', scale_factor=2),      #'trilinear'
                                    nn.Conv2d(in_size, out_size, kernel_size=1))
        elif 'inc' in up_mode:   
            self.up = nn.Conv2d(in_size, out_size, kernel_size=1)


        self.conv_block = UNetConvBlock(in_size, out_size, padding, batch_norm, Relu)

    def forward(self, x, bridge):
        if self.up_mode == 'upconv': #  'upconv'
            up = self.up(x)
        elif self.up_mode == 'bilinear':
            up = self.up(x)
        elif 'inc' in self.up_mode: 
            x = cF._sinc_interpolate(x, size=[int(x.shape[2]*2), int(x.shape[3]*2)]) #'sinc' ###sth wrong
            up = self.up(x)

        # bridge = self.center_crop(bridge, up.shape[2:]) #sending shape ignoring 2 digit, so target size start with 0,1,2
        up = F.interpolate(up, size=bridge.shape[2:], mode='bilinear')
        out = torch.cat([up, bridge], 1)
        out = self.conv_block(out)

        return out

#to run it here from this script, uncomment the following

if __name__ == "__main__":                #to run it
    image = torch.rand(2, 4, 240, 240)    #specify your image: batch size, Channel, height, width
    model = GP_UNet(in_channels=4, n_classes=3, depth=4, wf=6, up_mode="upsample_Sinc", Relu = "Relu")          #Initialize the model, up_mode = "upconv" or "upsample1" == interpolate mode Bilinear or "upsample" == interpolate mode sinc
    model.eval()
    out = model(image)
    print(model(image))