# -*- coding: utf-8 -*-

# Author: Guotai Wang

# Date:   12 June, 2020

# Implementation of of COPLENet for COVID-19 pneumonia lesion segmentation from CT images.

# Reference: 

#     G. Wang et al. A Noise-robust Framework for Automatic Segmentation of COVID-19 Pneumonia Lesions 

#     from CT Images. IEEE Transactions on Medical Imaging, 2020. DOI:10.1109/TMI.2020.3000314.

from __future__ import print_function, division

import torch

import torch.nn as nn

class ConvLayer(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size = 1):

        super(ConvLayer, self).__init__()

        padding = int((kernel_size - 1) / 2)

        self.conv = nn.Sequential(

            nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=padding),

            nn.BatchNorm2d(out_channels),

            nn.LeakyReLU()

        )

    def forward(self, x):

        return self.conv(x)

class SEBlock(nn.Module):

    def __init__(self, in_channels, r):

        super(SEBlock, self).__init__()

        redu_chns = int(in_channels / r)

        self.se_layers = nn.Sequential(

            nn.AdaptiveAvgPool2d(1),

            nn.Conv2d(in_channels, redu_chns, kernel_size=1, padding=0),

            nn.LeakyReLU(),

            nn.Conv2d(redu_chns, in_channels, kernel_size=1, padding=0),

            nn.ReLU())

    def forward(self, x):

        f = self.se_layers(x)

        return f*x + x

class ASPPBlock(nn.Module):

    def __init__(self,in_channels, out_channels_list, kernel_size_list, dilation_list):

        super(ASPPBlock, self).__init__()

        self.conv_num = len(out_channels_list)

        assert(self.conv_num == 4)

        assert(self.conv_num == len(kernel_size_list) and self.conv_num == len(dilation_list))

        pad0 = int((kernel_size_list[0] - 1) / 2 * dilation_list[0])

        pad1 = int((kernel_size_list[1] - 1) / 2 * dilation_list[1])

        pad2 = int((kernel_size_list[2] - 1) / 2 * dilation_list[2])

        pad3 = int((kernel_size_list[3] - 1) / 2 * dilation_list[3])

        self.conv_1 = nn.Conv2d(in_channels, out_channels_list[0], kernel_size = kernel_size_list[0], 

                    dilation = dilation_list[0], padding = pad0 )

        self.conv_2 = nn.Conv2d(in_channels, out_channels_list[1], kernel_size = kernel_size_list[1], 

                    dilation = dilation_list[1], padding = pad1 )

        self.conv_3 = nn.Conv2d(in_channels, out_channels_list[2], kernel_size = kernel_size_list[2], 

                    dilation = dilation_list[2], padding = pad2 )

        self.conv_4 = nn.Conv2d(in_channels, out_channels_list[3], kernel_size = kernel_size_list[3], 

                    dilation = dilation_list[3], padding = pad3 )

        out_channels  = out_channels_list[0] + out_channels_list[1] + out_channels_list[2] + out_channels_list[3] 

        self.conv_1x1 = nn.Sequential(

            nn.Conv2d(out_channels, out_channels, kernel_size=1, padding=0),

            nn.BatchNorm2d(out_channels),

            nn.LeakyReLU())

    def forward(self, x):

        x1 = self.conv_1(x)

        x2 = self.conv_2(x)

        x3 = self.conv_3(x)

        x4 = self.conv_4(x)

        y  = torch.cat([x1, x2, x3, x4], dim=1)

        y  = self.conv_1x1(y)

        return y

class ConvBNActBlock(nn.Module):

    """Two convolution layers with batch norm, leaky relu, dropout and SE block"""

    def __init__(self,in_channels, out_channels, dropout_p):

        super(ConvBNActBlock, self).__init__()

        self.conv_conv = nn.Sequential(

            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),

            nn.BatchNorm2d(out_channels),

            nn.LeakyReLU(),

            nn.Dropout(dropout_p),

            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),

            nn.BatchNorm2d(out_channels),

            nn.LeakyReLU(),

            SEBlock(out_channels, 2)

        )

    def forward(self, x):

        return self.conv_conv(x)

class DownBlock(nn.Module):

    """Downsampling by a concantenation of max-pool and avg-pool, followed by ConvBNActBlock

    """

    def __init__(self, in_channels, out_channels, dropout_p):

        super(DownBlock, self).__init__()

        self.maxpool = nn.MaxPool2d(2)

        self.avgpool = nn.AvgPool2d(2)

        self.conv    = ConvBNActBlock(2 * in_channels, out_channels, dropout_p)

    def forward(self, x):

        x_max = self.maxpool(x)

        x_avg = self.avgpool(x)

        x_cat = torch.cat([x_max, x_avg], dim=1)

        y     = self.conv(x_cat)

        return y + x_cat

class UpBlock(nn.Module):

    """Upssampling followed by ConvBNActBlock"""

    def __init__(self, in_channels1, in_channels2, out_channels, 

                 bilinear=True, dropout_p = 0.5):

        super(UpBlock, self).__init__()

        self.bilinear = bilinear

        if bilinear:

            self.conv1x1 = nn.Conv2d(in_channels1, in_channels2, kernel_size = 1)

            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

        else:

            self.up = nn.ConvTranspose2d(in_channels1, in_channels2, kernel_size=2, stride=2)

        self.conv = ConvBNActBlock(in_channels2 * 2, out_channels, dropout_p)

    def forward(self, x1, x2):

        if self.bilinear:

            x1 = self.conv1x1(x1)

        x1    = self.up(x1)

        x_cat = torch.cat([x2, x1], dim=1)

        y     = self.conv(x_cat)

        return y + x_cat

class COPLENet(nn.Module):

    def __init__(self, params):

        super(COPLENet, self).__init__()

        self.params    = params

        self.in_chns   = self.params['in_chns']

        self.ft_chns   = self.params['feature_chns']

        self.n_class   = self.params['class_num']

        self.bilinear  = self.params['bilinear']

        self.dropout   = self.params['dropout']

        assert(len(self.ft_chns) == 5)

        f0_half = int(self.ft_chns[0] / 2)

        f1_half = int(self.ft_chns[1] / 2)

        f2_half = int(self.ft_chns[2] / 2)

        f3_half = int(self.ft_chns[3] / 2)

        self.in_conv= ConvBNActBlock(self.in_chns, self.ft_chns[0], self.dropout[0])

        self.down1  = DownBlock(self.ft_chns[0], self.ft_chns[1], self.dropout[1])

        self.down2  = DownBlock(self.ft_chns[1], self.ft_chns[2], self.dropout[2])

        self.down3  = DownBlock(self.ft_chns[2], self.ft_chns[3], self.dropout[3])

        self.down4  = DownBlock(self.ft_chns[3], self.ft_chns[4], self.dropout[4])

        self.bridge0= ConvLayer(self.ft_chns[0], f0_half)

        self.bridge1= ConvLayer(self.ft_chns[1], f1_half)

        self.bridge2= ConvLayer(self.ft_chns[2], f2_half)

        self.bridge3= ConvLayer(self.ft_chns[3], f3_half)

        self.up1    = UpBlock(self.ft_chns[4], f3_half, self.ft_chns[3], dropout_p = self.dropout[3])

        self.up2    = UpBlock(self.ft_chns[3], f2_half, self.ft_chns[2], dropout_p = self.dropout[2])

        self.up3    = UpBlock(self.ft_chns[2], f1_half, self.ft_chns[1], dropout_p = self.dropout[1])

        self.up4    = UpBlock(self.ft_chns[1], f0_half, self.ft_chns[0], dropout_p = self.dropout[0])

        f4 = self.ft_chns[4]

        aspp_chns = [int(f4 / 4), int(f4 / 4), int(f4 / 4), int(f4 / 4)]

        aspp_knls = [1, 3, 3, 3]

        aspp_dila = [1, 2, 4, 6]

        self.aspp = ASPPBlock(f4, aspp_chns, aspp_knls, aspp_dila)

        self.out_conv = nn.Conv2d(self.ft_chns[0], self.n_class,  

            kernel_size = 3, padding = 1)

    def forward(self, x):

        x_shape = list(x.shape)

        if(len(x_shape) == 5):

          [N, C, D, H, W] = x_shape

          new_shape = [N*D, C, H, W]

          x = torch.transpose(x, 1, 2)

          x = torch.reshape(x, new_shape)

        x0  = self.in_conv(x)

        x0b = self.bridge0(x0)

        x1  = self.down1(x0)

        x1b = self.bridge1(x1)

        x2  = self.down2(x1)

        x2b = self.bridge2(x2)

        x3  = self.down3(x2)

        x3b = self.bridge3(x3)

        x4  = self.down4(x3)

        x4  = self.aspp(x4) 

        x   = self.up1(x4, x3b)

        x   = self.up2(x, x2b)

        x   = self.up3(x, x1b)

        x   = self.up4(x, x0b)

        output = self.out_conv(x)

        if(len(x_shape) == 5):

            new_shape = [N, D] + list(output.shape)[1:]

            output    = torch.reshape(output, new_shape)

            output    = torch.transpose(output, 1, 2)

        return output