GP_ShuffleUNet.py

##Dropout and out_act was added by hadya

import sys
import torch
import torch.nn as nn

from . import GP_ShuffleUNet_pixel_shuffle, GP_ShuffleUNet_pixel_unshuffle
# import pixel_shuffle, pixel_unshuffle

# -------------------------------------------------------------------------------------------------------------------------------------------------##

class _double_conv(nn.Module):
    """
    Double Convolution Block
    """

    def __init__(self, in_channels, out_channels, k_size, stride, bias=True, conv_layer=nn.Conv3d):
        super(_double_conv, self).__init__()
        self.conv_1 = conv_layer(in_channels=in_channels, out_channels=out_channels, kernel_size=k_size,
                      stride=stride, padding=k_size // 2, bias=bias)
        self.conv_2 = conv_layer(in_channels=out_channels, out_channels=out_channels, kernel_size=k_size,
                      stride=stride, padding=k_size // 2, bias=bias)

        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.conv_1(x)
        x = self.relu((x))
        x = self.conv_2(x)
        x = self.relu((x))

        return x


class _conv_decomp(nn.Module):
    """
    Convolutional Decomposition Block
    """

    def __init__(self, in_channels, out_channels, k_size, stride, bias=True, conv_layer=nn.Conv3d):
        super(_conv_decomp, self).__init__()
        self.conv1 = conv_layer(in_channels=in_channels, out_channels=out_channels, kernel_size=k_size,
                      stride=stride, padding=k_size // 2, bias=bias)
        self.conv2 = conv_layer(in_channels=in_channels, out_channels=out_channels, kernel_size=k_size,
                      stride=stride, padding=k_size // 2, bias=bias)
        self.conv3 = conv_layer(in_channels=in_channels, out_channels=out_channels, kernel_size=k_size,
                      stride=stride, padding=k_size // 2, bias=bias)
        self.conv4 = conv_layer(in_channels=in_channels, out_channels=out_channels, kernel_size=k_size,
                      stride=stride, padding=k_size // 2, bias=bias)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x1 = self.conv1(x)
        x1 = self.relu((x1))
        x2 = self.conv2(x)
        x2 = self.relu((x2))
        x3 = self.conv3(x)
        x3 = self.relu((x3))
        x4 = self.conv4(x)
        x4 = self.relu((x4))
        return x1, x2, x3, x4


class _concat(nn.Module):
    """
    Skip-Addition block
    """

    def __init__(self):
        super(_concat, self).__init__()

    def forward(self, e1, e2, e3, e4, d1, d2, d3, d4):
        self.X1 = e1 + d1
        self.X2 = e2 + d2
        self.X3 = e3 + d3
        self.X4 = e4 + d4
        x = torch.cat([self.X1, self.X2, self.X3, self.X4], dim=1)

        return x

# -------------------------------------------------------------------------------------------------------------------------------------------------##

class GP_ShuffleUNet(nn.Module):

    def __init__(self, d=3, in_ch=1, num_features=64, n_levels=3, out_ch=1, kernel_size=3, stride=1, dropout=False, out_act="softmax"):
        super(GP_ShuffleUNet, self).__init__()

        self.n_levels = n_levels
        self.dropout = nn.Dropout2d() if dropout else nn.Sequential()       #added by Hadya

        num_features = num_features
        filters = [num_features]
        for _ in range(n_levels):
            filters.append(filters[-1]*2)

        if d==3:
            conv_layer = nn.Conv3d
            ps_fact = (2 ** 2)
        elif d==2:
            conv_layer = nn.Conv2d
            ps_fact = 2
        else:
            sys.exit("Invalid d")

        # Input
        self.conv_inp = _double_conv(in_ch, filters[0], kernel_size, stride, conv_layer=conv_layer)     

        #Contraction path
        self.wave_down = nn.ModuleList()
        self.pix_unshuff = nn.ModuleList()
        self.conv_enc = nn.ModuleList()
        for i in range(0, n_levels):
            self.wave_down.append(_conv_decomp(filters[i], filters[i], kernel_size, stride, conv_layer=conv_layer))
            self.pix_unshuff.append(pixel_unshuffle.PixelUnshuffle(num_features * (2**i), num_features * (2**i), kernel_size, stride, d=d))       
            self.conv_enc.append(_double_conv(filters[i], filters[i+1], kernel_size, stride, conv_layer=conv_layer))     

        #Expansion path
        self.cat = _concat()
        self.pix_shuff = nn.ModuleList()
        self.wave_up = nn.ModuleList()
        self.convup = nn.ModuleList()
        for i in range(n_levels-1,-1,-1):
            self.pix_shuff.append(pixel_shuffle.PixelShuffle(num_features * (2**(i+1)), num_features * (2**(i+1)) * ps_fact, kernel_size, stride, d=d))
            self.wave_up.append(_conv_decomp(filters[i], filters[i], kernel_size, stride, conv_layer=conv_layer))
            self.convup.append(_double_conv(filters[i] * 5, filters[i], kernel_size, stride, conv_layer=conv_layer))

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

        #FC
        if out_act == "softmax":            #added by Hadya
            self.last = nn.Sequential(
                      conv_layer(filters[0], out_ch, kernel_size=1, stride=1, padding=0, bias=True),
                      nn.Softmax2d()
                    )
        else:
            self.out = conv_layer(filters[0], out_ch, kernel_size=1, stride=1, padding=0, bias=True)  #original

        #Weight init
        for m in self.modules():
            if isinstance(m, conv_layer):
                weight = nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
                m.weight.data.copy_(weight)
                if m.bias is not None:
                    m.bias.data.zero_()

    def forward(self, x):
        encs = [self.conv_inp(x)]

        waves = []
        for i in range(self.n_levels):
            waves.append(self.wave_down[i](encs[-1]))
            _tmp = self.pix_unshuff[i](waves[-1][-1])
            encs.append(self.conv_enc[i](_tmp))

        dec = encs.pop()

        dec = self.dropout(dec) #added by hadya

        for i in range(self.n_levels):
            _tmp = self.pix_shuff[i](dec)
            _tmp_waves = self.wave_up[i](_tmp) + waves.pop()
            _tmp_cat = self.cat(*_tmp_waves)
            dec = self.convup[i](torch.cat([encs.pop(), _tmp_cat], dim=1))

        ###added section to make it GP-UNet
        if self.training:
            x = self.GMP(dec)
            return self.out(x).view(x.shape[0],-1)
        else:
            mask = self.out(dec)
            x = self.GMP(dec)
            pred = self.out(x).view(x.shape[0],-1)
            return pred, mask


        # return self.out(dec) #####replace by line 154-161 to make it GP_ShuffleUNet


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

if __name__ == "__main__":                #to run it
    image = torch.rand(2, 1, 512, 512)    #specify your image: batch size, Channel, height, width
    model = GP_ShuffleUNet(d=2, in_ch=1, num_features=64, n_levels=3, out_ch=3, kernel_size=3, stride=1)    #Initialize the model, d=3 default is for dimensionality conv2d or 3d, default out channel = 1 but in gp we need 3
    model.eval()
    out = model(image)
    print(model(image))