<gh_stars>1-10
import torch

import torch.nn as nn

from hdc_duc import HDC, DUC
import torch.nn.functional as F


# 3d unet model borrow from : https://github.com/wolny/pytorch-3dunet

def create_feature_maps(init_channel_number, number_of_fmaps):
    return [init_channel_number * 2 ** k for k in range(number_of_fmaps)]


def conv3d(in_channels, out_channels, kernel_size, bias, padding=1):
    return nn.Conv3d(in_channels,
                     out_channels,
                     kernel_size,
                     padding=padding,
                     bias=bias)


def create_conv(in_channels,
                out_channels,
                kernel_size,
                order,
                num_groups,
                padding=1):
    """
    Create a list of modules with together constitute a single conv layer with non-linearity
    and optional batchnorm/groupnorm.

    Args:
        in_channels (int): number of input channels
        out_channels (int): number of output channels
        order (string): order of things, e.g.
            'cr' -> conv + ReLU
            'gcr' -> groupnorm + conv + ReLU
            'cl' -> conv + LeakyReLU
            'ce' -> conv + ELU
            'bcr' -> batchnorm + conv + ReLU
        num_groups (int): number of groups for the GroupNorm
        padding (int): add zero-padding to the input

    Return:
        list of tuple (name, module)
    """
    assert 'c' in order, "Conv layer MUST be present"
    assert order[
               0] not in 'rle', 'Non-linearity cannot be the first operation in the layer'

    modules = []
    for i, char in enumerate(order):
        if char == 'r':
            modules.append(('ReLU', nn.ReLU(inplace=True)))
        elif char == 'l':
            modules.append(
                ('LeakyReLU', nn.LeakyReLU(negative_slope=0.1, inplace=True)))
        elif char == 'e':
            modules.append(('ELU', nn.ELU(inplace=True)))
        elif char == 'c':
            # add learnable bias only in the absence of gatchnorm/groupnorm
            bias = not ('g' in order or 'b' in order)
            modules.append(('conv',
                            conv3d(in_channels,
                                   out_channels,
                                   kernel_size,
                                   bias,
                                   padding=padding)))
        elif char == 'g':
            is_before_conv = i < order.index('c')
            if is_before_conv:
                num_channels = in_channels
            else:
                num_channels = out_channels

            # use only one group if the given number of groups is greater than the number of channels
            if num_channels < num_groups:
                num_groups = 1

            assert num_channels % num_groups == 0, 'Expected number of channels in input to be divisible by num_groups. num_channels={}, num_groups={}'.format(
                num_channels, num_groups)
            modules.append(('groupnorm',
                            nn.GroupNorm(num_groups=num_groups,
                                         num_channels=num_channels)))
        elif char == 'b':
            is_before_conv = i < order.index('c')
            if is_before_conv:
                modules.append(('batchnorm', nn.BatchNorm3d(in_channels)))
            else:
                modules.append(('batchnorm', nn.BatchNorm3d(out_channels)))
        else:
            raise ValueError(
                "Unsupported layer type '{}'. MUST be one of ['b', 'g', 'r', 'l', 'e', 'c']"
                    .format(char))

    return modules


class SingleConv(nn.Sequential):
    """
    Basic convolutional module consisting of a Conv3d, non-linearity and optional batchnorm/groupnorm. The order
    of operations can be specified via the `order` parameter

    Args:
        in_channels (int): number of input channels
        out_channels (int): number of output channels
        kernel_size (int): size of the convolving kernel
        order (string): determines the order of layers, e.g.
            'cr' -> conv + ReLU
            'crg' -> conv + ReLU + groupnorm
            'cl' -> conv + LeakyReLU
            'ce' -> conv + ELU
        num_groups (int): number of groups for the GroupNorm
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=3,
                 order='crg',
                 num_groups=8,
                 padding=1):
        super(SingleConv, self).__init__()

        for name, module in create_conv(in_channels,
                                        out_channels,
                                        kernel_size,
                                        order,
                                        num_groups,
                                        padding=padding):
            self.add_module(name, module)


class DoubleConv(nn.Sequential):
    """
    A module consisting of two consecutive convolution layers (e.g. BatchNorm3d+ReLU+Conv3d).
    We use (Conv3d+ReLU+GroupNorm3d) by default.
    This can be changed however by providing the 'order' argument, e.g. in order
    to change to Conv3d+BatchNorm3d+ELU use order='cbe'.
    Use padded convolutions to make sure that the output (H_out, W_out) is the same
    as (H_in, W_in), so that you don't have to crop in the decoder path.

    Args:
        in_channels (int): number of input channels
        out_channels (int): number of output channels
        encoder (bool): if True we're in the encoder path, otherwise we're in the decoder
        kernel_size (int): size of the convolving kernel
        order (string): determines the order of layers, e.g.
            'cr' -> conv + ReLU
            'crg' -> conv + ReLU + groupnorm
            'cl' -> conv + LeakyReLU
            'ce' -> conv + ELU
        num_groups (int): number of groups for the GroupNorm
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 encoder,
                 kernel_size=3,
                 order='crg',
                 num_groups=8):
        super(DoubleConv, self).__init__()
        if encoder:
            # we're in the encoder path
            conv1_in_channels = in_channels
            conv1_out_channels = out_channels // 2
            if conv1_out_channels < in_channels:
                conv1_out_channels = in_channels
            conv2_in_channels, conv2_out_channels = conv1_out_channels, out_channels
        else:
            # we're in the decoder path, decrease the number of channels in the 1st convolution
            conv1_in_channels, conv1_out_channels = in_channels, out_channels
            conv2_in_channels, conv2_out_channels = out_channels, out_channels

        # conv1
        self.add_module(
            'SingleConv1',
            SingleConv(conv1_in_channels, conv1_out_channels, kernel_size,
                       order, num_groups))
        # conv2
        self.add_module(
            'SingleConv2',
            SingleConv(conv2_in_channels, conv2_out_channels, kernel_size,
                       order, num_groups))


class ExtResNetBlock(nn.Module):
    """
    Basic UNet block consisting of a SingleConv followed by the residual block.
    The SingleConv takes care of increasing/decreasing the number of channels and also ensures that the number
    of output channels is compatible with the residual block that follows.
    This block can be used instead of standard DoubleConv in the Encoder module.
    Motivated by: https://arxiv.org/pdf/1706.00120.pdf

    Notice we use ELU instead of ReLU (order='cge') and put non-linearity after the groupnorm.
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=3,
                 order='cge',
                 num_groups=8,
                 **kwargs):
        super(ExtResNetBlock, self).__init__()

        # first convolution
        self.conv1 = SingleConv(in_channels,
                                out_channels,
                                kernel_size=kernel_size,
                                order=order,
                                num_groups=num_groups)
        # residual block
        self.conv2 = SingleConv(out_channels,
                                out_channels,
                                kernel_size=kernel_size,
                                order=order,
                                num_groups=num_groups)
        # remove non-linearity from the 3rd convolution since it's going to be applied after adding the residual
        n_order = order
        for c in 'rel':
            n_order = n_order.replace(c, '')
        self.conv3 = SingleConv(out_channels,
                                out_channels,
                                kernel_size=kernel_size,
                                order=n_order,
                                num_groups=num_groups)

        # create non-linearity separately
        if 'l' in order:
            self.non_linearity = nn.LeakyReLU(negative_slope=0.1, inplace=True)
        elif 'e' in order:
            self.non_linearity = nn.ELU(inplace=True)
        else:
            self.non_linearity = nn.ReLU(inplace=True)

    def forward(self, x):
        # apply first convolution and save the output as a residual
        out = self.conv1(x)
        residual = out

        # residual block
        out = self.conv2(out)
        out = self.conv3(out)

        out += residual
        out = self.non_linearity(out)

        return out


class Encoder(nn.Module):
    """
    A single module from the encoder path consisting of the optional max
    pooling layer (one may specify the MaxPool kernel_size to be different
    than the standard (2,2,2), e.g. if the volumetric data is anisotropic
    (make sure to use complementary scale_factor in the decoder path) followed by
    a DoubleConv module.
    Args:
        in_channels (int): number of input channels
        out_channels (int): number of output channels
        conv_kernel_size (int): size of the convolving kernel
        apply_pooling (bool): if True use MaxPool3d before DoubleConv
        pool_kernel_size (tuple): the size of the window to take a max over
        pool_type (str): pooling layer: 'max' or 'avg'
        basic_module(nn.Module): either ResNetBlock or DoubleConv
        conv_layer_order (string): determines the order of layers
            in `DoubleConv` module. See `DoubleConv` for more info.
        num_groups (int): number of groups for the GroupNorm
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 conv_kernel_size=3,
                 apply_pooling=True,
                 pool_kernel_size=(2, 2, 2),
                 pool_type='max',
                 basic_module=DoubleConv,
                 conv_layer_order='crg',
                 num_groups=8):
        super(Encoder, self).__init__()
        assert pool_type in ['max', 'avg']
        if apply_pooling:
            if pool_type == 'max':
                self.pooling = nn.MaxPool3d(kernel_size=pool_kernel_size)
            else:
                self.pooling = nn.AvgPool3d(kernel_size=pool_kernel_size)
        else:
            self.pooling = None

        self.basic_module = basic_module(in_channels,
                                         out_channels,
                                         encoder=True,
                                         kernel_size=conv_kernel_size,
                                         order=conv_layer_order,
                                         num_groups=num_groups)

    def forward(self, x):
        if self.pooling is not None:
            x = self.pooling(x)
        x = self.basic_module(x)
        return x


class Decoder(nn.Module):
    """
    A single module for decoder path consisting of the upsample layer
    (either learned ConvTranspose3d or interpolation) followed by a DoubleConv
    module.
    Args:
        in_channels (int): number of input channels
        out_channels (int): number of output channels
        kernel_size (int): size of the convolving kernel
        scale_factor (tuple): used as the multiplier for the image H/W/D in
            case of nn.Upsample or as stride in case of ConvTranspose3d, must reverse the MaxPool3d operation
            from the corresponding encoder
        basic_module(nn.Module): either ResNetBlock or DoubleConv
        conv_layer_order (string): determines the order of layers
            in `DoubleConv` module. See `DoubleConv` for more info.
        num_groups (int): number of groups for the GroupNorm
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=3,
                 scale_factor=(2, 2, 2),
                 basic_module=DoubleConv,
                 conv_layer_order='crg',
                 num_groups=8):
        super(Decoder, self).__init__()
        if basic_module == DoubleConv:
            # if DoubleConv is the basic_module use nearest neighbor interpolation for upsampling
            self.upsample = None
        else:
            # otherwise use ConvTranspose3d (bear in mind your GPU memory)
            # make sure that the output size reverses the MaxPool3d from the corresponding encoder
            # (D_out = (D_in − 1) ×  stride[0] − 2 ×  padding[0] +  kernel_size[0] +  output_padding[0])
            # also scale the number of channels from in_channels to out_channels so that summation joining
            # works correctly
            self.upsample = nn.ConvTranspose3d(in_channels,
                                               out_channels,
                                               kernel_size=kernel_size,
                                               stride=scale_factor,
                                               padding=1,
                                               output_padding=1)
            # adapt the number of in_channels for the ExtResNetBlock
            in_channels = out_channels

        self.basic_module = basic_module(in_channels,
                                         out_channels,
                                         encoder=False,
                                         kernel_size=kernel_size,
                                         order=conv_layer_order,
                                         num_groups=num_groups)

    def forward(self, encoder_features, x):
        if self.upsample is None:
            # use nearest neighbor interpolation and concatenation joining
            output_size = encoder_features.size()[2:]
            x = F.interpolate(x, size=output_size, mode='nearest')  # interpolation in decoder
            # concatenate encoder_features (encoder path) with the upsampled input across channel dimension
            x = torch.cat((encoder_features, x), dim=1)
        else:
            # use ConvTranspose3d and summation joining
            x = self.upsample(x)
            x += encoder_features

        x = self.basic_module(x)
        return x


class FinalConv(nn.Sequential):
    """
    A module consisting of a convolution layer (e.g. Conv3d+ReLU+GroupNorm3d) and the final 1x1 convolution
    which reduces the number of channels to 'out_channels'.
    with the number of output channels 'out_channels // 2' and 'out_channels' respectively.
    We use (Conv3d+ReLU+GroupNorm3d) by default.
    This can be change however by providing the 'order' argument, e.g. in order
    to change to Conv3d+BatchNorm3d+ReLU use order='cbr'.
    Args:
        in_channels (int): number of input channels
        out_channels (int): number of output channels
        kernel_size (int): size of the convolving kernel
        order (string): determines the order of layers, e.g.
            'cr' -> conv + ReLU
            'crg' -> conv + ReLU + groupnorm
        num_groups (int): number of groups for the GroupNorm
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=3,
                 order='crg',
                 num_groups=8):
        super(FinalConv, self).__init__()

        # conv1
        self.add_module(
            'SingleConv',
            SingleConv(in_channels, in_channels, kernel_size, order,
                       num_groups))

        # in the last layer a 1×1 convolution reduces the number of output channels to out_channels
        final_conv = nn.Conv3d(in_channels, out_channels, 1)
        self.add_module('final_conv', final_conv)


class UNet3D_HDC(nn.Module):
    """
    3DUnet model from
    "3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation"
        <https://arxiv.org/pdf/1606.0660.pdf>.
    HDC from:
    "Holistic decomposition convolution for effective semantic segmentation of medical volume images"
         <https://doi.org/10.1016/j.media.2019.07.003>.
    DUC from:
    "Real-time single image and video super-resolution using an efficient sub-pixel convolutional neural network"
        <https://doi.org/10.1109/CVPR.2016.207>.
    Args:
        in_channels (int): number of input channels
        out_channels (int): number of output segmentation masks;
            Note that that the of out_channels might correspond to either
            different semantic classes or to different binary segmentation mask.
            It's up to the user of the class to interpret the out_channels and
            use the proper loss criterion during training (i.e. CrossEntropyLoss (multi-class)
            or BCEWithLogitsLoss (two-class) respectively)
        f_maps (int, tuple): number of feature maps at each level of the encoder; if it's an integer the number
            of feature maps is given by the geometric progression: f_maps ^ k, k=1,2,3,4
        final_sigmoid (bool): if True apply element-wise nn.Sigmoid after the
            final 1x1 convolution, otherwise apply nn.Softmax. MUST be True if nn.BCELoss (two-class) is used
            to train the model. MUST be False if nn.CrossEntropyLoss (multi-class) is used to train the model.
        layer_order (string): determines the order of layers
            in `SingleConv` module. e.g. 'crg' stands for Conv3d+ReLU+GroupNorm3d.
            See `SingleConv` for more info
        init_channel_number (int): number of feature maps in the first conv layer of the encoder; default: 64
        num_groups (int): number of groups for the GroupNorm
        number_of_fmas (int): depth of the model
        scale_factor (int or tuple or list): Scale up and scale down factors.
    Notes:
        input data shape : (batchsize,channels,height,width,depth) or (batch,channels,x,y,z)
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 final_sigmoid,
                 f_maps=64,
                 layer_order='crg',
                 num_groups=8,
                 number_of_fmaps=4,
                 scale_factor=2,
                 **kwargs):
        super(UNet3D_HDC, self).__init__()

        self.in_hdc = HDC(scale_factor)

        if isinstance(f_maps, int):
            # use number_of_fmaps levels in the encoder path as suggested in the paper
            f_maps = create_feature_maps(f_maps,
                                         number_of_fmaps=number_of_fmaps)

        # create encoder path consisting of Encoder modules. The length of the encoder is equal to `len(f_maps)`
        # uses DoubleConv as a basic_module for the Encoder
        encoders = []
        for i, out_feature_num in enumerate(f_maps):
            if i == 0:
                encoder = Encoder(in_channels,
                                  out_feature_num,
                                  apply_pooling=False,  # no pooling in the first layer
                                  basic_module=DoubleConv,
                                  conv_layer_order=layer_order,
                                  num_groups=num_groups)
            else:
                encoder = Encoder(f_maps[i - 1],
                                  out_feature_num,
                                  basic_module=DoubleConv,
                                  conv_layer_order=layer_order,
                                  num_groups=num_groups)
            encoders.append(encoder)

        self.encoders = nn.ModuleList(encoders)

        # create decoder path consisting of the Decoder modules. The length of the decoder is equal to `len(f_maps) - 1`
        # uses DoubleConv as a basic_module for the Decoder
        decoders = []
        reversed_f_maps = list(reversed(f_maps))
        for i in range(len(reversed_f_maps) - 1):
            in_feature_num = reversed_f_maps[i] + reversed_f_maps[i + 1]
            out_feature_num = reversed_f_maps[i + 1]
            decoder = Decoder(in_feature_num,
                              out_feature_num,
                              basic_module=DoubleConv,
                              conv_layer_order=layer_order,
                              num_groups=num_groups)
            decoders.append(decoder)

        self.decoders = nn.ModuleList(decoders)

        # in the last layer a 1×1 convolution reduces the number of output
        # channels to the number of labels
        # self.final_conv = nn.Conv3d(f_maps[0], out_channels, 1)
        self.out_duc = DUC(scale_factor, out_channels, f_maps[0])
        if final_sigmoid:
            self.final_activation = nn.Sigmoid()
        else:
            self.final_activation = nn.Softmax(dim=1)

    def forward(self, x):
        # encoder part
        encoders_features = []
        x = self.in_hdc(x)
        for encoder in self.encoders:
            x = encoder(x)
            # reverse the encoder outputs to be aligned with the decoder
            encoders_features.insert(
                0, x)  # insert(index,object)

        # remove the last encoder's output from the list
        # !!remember: it's the 1st in the list
        encoders_features = encoders_features[1:]

        # decoder part
        for decoder, encoder_features in zip(self.decoders, encoders_features):
            # pass the output from the corresponding encoder and the output
            # of the previous decoder
            x = decoder(encoder_features, x)

        # x = self.final_conv(x)
        x = self.out_duc(x)
        # apply final_activation (i.e. Sigmoid or Softmax) only for prediction. During training the network outputs
        # logits and it's up to the user to normalize it before visualising with tensorboard or computing validation metric
        if not self.training:
            x = self.final_activation(x)  # training also output sigmoid value

        return x

import torch

import torch.nn as nn

from hdc_duc import HDC, DUC
import torch.nn.functional as F


# 3d unet model borrow from : https://github.com/wolny/pytorch-3dunet

def create_feature_maps(init_channel_number, number_of_fmaps):
    return [init_channel_number * 2 ** k for k in range(number_of_fmaps)]


def conv3d(in_channels, out_channels, kernel_size, bias, padding=1):
    return nn.Conv3d(in_channels,
                     out_channels,
                     kernel_size,
                     padding=padding,
                     bias=bias)


def create_conv(in_channels,
                out_channels,
                kernel_size,
                order,
                num_groups,
                padding=1):
    """
    Create a list of modules with together constitute a single conv layer with non-linearity
    and optional batchnorm/groupnorm.

    Args:
        in_channels (int): number of input channels
        out_channels (int): number of output channels
        order (string): order of things, e.g.
            'cr' -> conv + ReLU
            'gcr' -> groupnorm + conv + ReLU
            'cl' -> conv + LeakyReLU
            'ce' -> conv + ELU
            'bcr' -> batchnorm + conv + ReLU
        num_groups (int): number of groups for the GroupNorm
        padding (int): add zero-padding to the input

    Return:
        list of tuple (name, module)
    """
    assert 'c' in order, "Conv layer MUST be present"
    assert order[
               0] not in 'rle', 'Non-linearity cannot be the first operation in the layer'

    modules = []
    for i, char in enumerate(order):
        if char == 'r':
            modules.append(('ReLU', nn.ReLU(inplace=True)))
        elif char == 'l':
            modules.append(
                ('LeakyReLU', nn.LeakyReLU(negative_slope=0.1, inplace=True)))
        elif char == 'e':
            modules.append(('ELU', nn.ELU(inplace=True)))
        elif char == 'c':
            # add learnable bias only in the absence of gatchnorm/groupnorm
            bias = not ('g' in order or 'b' in order)
            modules.append(('conv',
                            conv3d(in_channels,
                                   out_channels,
                                   kernel_size,
                                   bias,
                                   padding=padding)))
        elif char == 'g':
            is_before_conv = i < order.index('c')
            if is_before_conv:
                num_channels = in_channels
            else:
                num_channels = out_channels

            # use only one group if the given number of groups is greater than the number of channels
            if num_channels < num_groups:
                num_groups = 1

            assert num_channels % num_groups == 0, 'Expected number of channels in input to be divisible by num_groups. num_channels={}, num_groups={}'.format(
                num_channels, num_groups)
            modules.append(('groupnorm',
                            nn.GroupNorm(num_groups=num_groups,
                                         num_channels=num_channels)))
        elif char == 'b':
            is_before_conv = i < order.index('c')
            if is_before_conv:
                modules.append(('batchnorm', nn.BatchNorm3d(in_channels)))
            else:
                modules.append(('batchnorm', nn.BatchNorm3d(out_channels)))
        else:
            raise ValueError(
                "Unsupported layer type '{}'. MUST be one of ['b', 'g', 'r', 'l', 'e', 'c']"
                    .format(char))

    return modules


class SingleConv(nn.Sequential):
    """
    Basic convolutional module consisting of a Conv3d, non-linearity and optional batchnorm/groupnorm. The order
    of operations can be specified via the `order` parameter

    Args:
        in_channels (int): number of input channels
        out_channels (int): number of output channels
        kernel_size (int): size of the convolving kernel
        order (string): determines the order of layers, e.g.
            'cr' -> conv + ReLU
            'crg' -> conv + ReLU + groupnorm
            'cl' -> conv + LeakyReLU
            'ce' -> conv + ELU
        num_groups (int): number of groups for the GroupNorm
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=3,
                 order='crg',
                 num_groups=8,
                 padding=1):
        super(SingleConv, self).__init__()

        for name, module in create_conv(in_channels,
                                        out_channels,
                                        kernel_size,
                                        order,
                                        num_groups,
                                        padding=padding):
            self.add_module(name, module)


class DoubleConv(nn.Sequential):
    """
    A module consisting of two consecutive convolution layers (e.g. BatchNorm3d+ReLU+Conv3d).
    We use (Conv3d+ReLU+GroupNorm3d) by default.
    This can be changed however by providing the 'order' argument, e.g. in order
    to change to Conv3d+BatchNorm3d+ELU use order='cbe'.
    Use padded convolutions to make sure that the output (H_out, W_out) is the same
    as (H_in, W_in), so that you don't have to crop in the decoder path.

    Args:
        in_channels (int): number of input channels
        out_channels (int): number of output channels
        encoder (bool): if True we're in the encoder path, otherwise we're in the decoder
        kernel_size (int): size of the convolving kernel
        order (string): determines the order of layers, e.g.
            'cr' -> conv + ReLU
            'crg' -> conv + ReLU + groupnorm
            'cl' -> conv + LeakyReLU
            'ce' -> conv + ELU
        num_groups (int): number of groups for the GroupNorm
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 encoder,
                 kernel_size=3,
                 order='crg',
                 num_groups=8):
        super(DoubleConv, self).__init__()
        if encoder:
            # we're in the encoder path
            conv1_in_channels = in_channels
            conv1_out_channels = out_channels // 2
            if conv1_out_channels < in_channels:
                conv1_out_channels = in_channels
            conv2_in_channels, conv2_out_channels = conv1_out_channels, out_channels
        else:
            # we're in the decoder path, decrease the number of channels in the 1st convolution
            conv1_in_channels, conv1_out_channels = in_channels, out_channels
            conv2_in_channels, conv2_out_channels = out_channels, out_channels

        # conv1
        self.add_module(
            'SingleConv1',
            SingleConv(conv1_in_channels, conv1_out_channels, kernel_size,
                       order, num_groups))
        # conv2
        self.add_module(
            'SingleConv2',
            SingleConv(conv2_in_channels, conv2_out_channels, kernel_size,
                       order, num_groups))


class ExtResNetBlock(nn.Module):
    """
    Basic UNet block consisting of a SingleConv followed by the residual block.
    The SingleConv takes care of increasing/decreasing the number of channels and also ensures that the number
    of output channels is compatible with the residual block that follows.
    This block can be used instead of standard DoubleConv in the Encoder module.
    Motivated by: https://arxiv.org/pdf/1706.00120.pdf

    Notice we use ELU instead of ReLU (order='cge') and put non-linearity after the groupnorm.
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=3,
                 order='cge',
                 num_groups=8,
                 **kwargs):
        super(ExtResNetBlock, self).__init__()

        # first convolution
        self.conv1 = SingleConv(in_channels,
                                out_channels,
                                kernel_size=kernel_size,
                                order=order,
                                num_groups=num_groups)
        # residual block
        self.conv2 = SingleConv(out_channels,
                                out_channels,
                                kernel_size=kernel_size,
                                order=order,
                                num_groups=num_groups)
        # remove non-linearity from the 3rd convolution since it's going to be applied after adding the residual
        n_order = order
        for c in 'rel':
            n_order = n_order.replace(c, '')
        self.conv3 = SingleConv(out_channels,
                                out_channels,
                                kernel_size=kernel_size,
                                order=n_order,
                                num_groups=num_groups)

        # create non-linearity separately
        if 'l' in order:
            self.non_linearity = nn.LeakyReLU(negative_slope=0.1, inplace=True)
        elif 'e' in order:
            self.non_linearity = nn.ELU(inplace=True)
        else:
            self.non_linearity = nn.ReLU(inplace=True)

    def forward(self, x):
        # apply first convolution and save the output as a residual
        out = self.conv1(x)
        residual = out

        # residual block
        out = self.conv2(out)
        out = self.conv3(out)

        out += residual
        out = self.non_linearity(out)

        return out


class Encoder(nn.Module):
    """
    A single module from the encoder path consisting of the optional max
    pooling layer (one may specify the MaxPool kernel_size to be different
    than the standard (2,2,2), e.g. if the volumetric data is anisotropic
    (make sure to use complementary scale_factor in the decoder path) followed by
    a DoubleConv module.
    Args:
        in_channels (int): number of input channels
        out_channels (int): number of output channels
        conv_kernel_size (int): size of the convolving kernel
        apply_pooling (bool): if True use MaxPool3d before DoubleConv
        pool_kernel_size (tuple): the size of the window to take a max over
        pool_type (str): pooling layer: 'max' or 'avg'
        basic_module(nn.Module): either ResNetBlock or DoubleConv
        conv_layer_order (string): determines the order of layers
            in `DoubleConv` module. See `DoubleConv` for more info.
        num_groups (int): number of groups for the GroupNorm
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 conv_kernel_size=3,
                 apply_pooling=True,
                 pool_kernel_size=(2, 2, 2),
                 pool_type='max',
                 basic_module=DoubleConv,
                 conv_layer_order='crg',
                 num_groups=8):
        super(Encoder, self).__init__()
        assert pool_type in ['max', 'avg']
        if apply_pooling:
            if pool_type == 'max':
                self.pooling = nn.MaxPool3d(kernel_size=pool_kernel_size)
            else:
                self.pooling = nn.AvgPool3d(kernel_size=pool_kernel_size)
        else:
            self.pooling = None

        self.basic_module = basic_module(in_channels,
                                         out_channels,
                                         encoder=True,
                                         kernel_size=conv_kernel_size,
                                         order=conv_layer_order,
                                         num_groups=num_groups)

    def forward(self, x):
        if self.pooling is not None:
            x = self.pooling(x)
        x = self.basic_module(x)
        return x


class Decoder(nn.Module):
    """
    A single module for decoder path consisting of the upsample layer
    (either learned ConvTranspose3d or interpolation) followed by a DoubleConv
    module.
    Args:
        in_channels (int): number of input channels
        out_channels (int): number of output channels
        kernel_size (int): size of the convolving kernel
        scale_factor (tuple): used as the multiplier for the image H/W/D in
            case of nn.Upsample or as stride in case of ConvTranspose3d, must reverse the MaxPool3d operation
            from the corresponding encoder
        basic_module(nn.Module): either ResNetBlock or DoubleConv
        conv_layer_order (string): determines the order of layers
            in `DoubleConv` module. See `DoubleConv` for more info.
        num_groups (int): number of groups for the GroupNorm
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=3,
                 scale_factor=(2, 2, 2),
                 basic_module=DoubleConv,
                 conv_layer_order='crg',
                 num_groups=8):
        super(Decoder, self).__init__()
        if basic_module == DoubleConv:
            # if DoubleConv is the basic_module use nearest neighbor interpolation for upsampling
            self.upsample = None
        else:
            # otherwise use ConvTranspose3d (bear in mind your GPU memory)
            # make sure that the output size reverses the MaxPool3d from the corresponding encoder
            # (D_out = (D_in − 1) ×  stride[0] − 2 ×  padding[0] +  kernel_size[0] +  output_padding[0])
            # also scale the number of channels from in_channels to out_channels so that summation joining
            # works correctly
            self.upsample = nn.ConvTranspose3d(in_channels,
                                               out_channels,
                                               kernel_size=kernel_size,
                                               stride=scale_factor,
                                               padding=1,
                                               output_padding=1)
            # adapt the number of in_channels for the ExtResNetBlock
            in_channels = out_channels

        self.basic_module = basic_module(in_channels,
                                         out_channels,
                                         encoder=False,
                                         kernel_size=kernel_size,
                                         order=conv_layer_order,
                                         num_groups=num_groups)

    def forward(self, encoder_features, x):
        if self.upsample is None:
            # use nearest neighbor interpolation and concatenation joining
            output_size = encoder_features.size()[2:]
            x = F.interpolate(x, size=output_size, mode='nearest')  # interpolation in decoder
            # concatenate encoder_features (encoder path) with the upsampled input across channel dimension
            x = torch.cat((encoder_features, x), dim=1)
        else:
            # use ConvTranspose3d and summation joining
            x = self.upsample(x)
            x += encoder_features

        x = self.basic_module(x)
        return x


class FinalConv(nn.Sequential):
    """
    A module consisting of a convolution layer (e.g. Conv3d+ReLU+GroupNorm3d) and the final 1x1 convolution
    which reduces the number of channels to 'out_channels'.
    with the number of output channels 'out_channels // 2' and 'out_channels' respectively.
    We use (Conv3d+ReLU+GroupNorm3d) by default.
    This can be change however by providing the 'order' argument, e.g. in order
    to change to Conv3d+BatchNorm3d+ReLU use order='cbr'.
    Args:
        in_channels (int): number of input channels
        out_channels (int): number of output channels
        kernel_size (int): size of the convolving kernel
        order (string): determines the order of layers, e.g.
            'cr' -> conv + ReLU
            'crg' -> conv + ReLU + groupnorm
        num_groups (int): number of groups for the GroupNorm
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=3,
                 order='crg',
                 num_groups=8):
        super(FinalConv, self).__init__()

        # conv1
        self.add_module(
            'SingleConv',
            SingleConv(in_channels, in_channels, kernel_size, order,
                       num_groups))

        # in the last layer a 1×1 convolution reduces the number of output channels to out_channels
        final_conv = nn.Conv3d(in_channels, out_channels, 1)
        self.add_module('final_conv', final_conv)


class UNet3D_HDC(nn.Module):
    """
    3DUnet model from
    "3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation"
        <https://arxiv.org/pdf/1606.0660.pdf>.
    HDC from:
    "Holistic decomposition convolution for effective semantic segmentation of medical volume images"
         <https://doi.org/10.1016/j.media.2019.07.003>.
    DUC from:
    "Real-time single image and video super-resolution using an efficient sub-pixel convolutional neural network"
        <https://doi.org/10.1109/CVPR.2016.207>.
    Args:
        in_channels (int): number of input channels
        out_channels (int): number of output segmentation masks;
            Note that that the of out_channels might correspond to either
            different semantic classes or to different binary segmentation mask.
            It's up to the user of the class to interpret the out_channels and
            use the proper loss criterion during training (i.e. CrossEntropyLoss (multi-class)
            or BCEWithLogitsLoss (two-class) respectively)
        f_maps (int, tuple): number of feature maps at each level of the encoder; if it's an integer the number
            of feature maps is given by the geometric progression: f_maps ^ k, k=1,2,3,4
        final_sigmoid (bool): if True apply element-wise nn.Sigmoid after the
            final 1x1 convolution, otherwise apply nn.Softmax. MUST be True if nn.BCELoss (two-class) is used
            to train the model. MUST be False if nn.CrossEntropyLoss (multi-class) is used to train the model.
        layer_order (string): determines the order of layers
            in `SingleConv` module. e.g. 'crg' stands for Conv3d+ReLU+GroupNorm3d.
            See `SingleConv` for more info
        init_channel_number (int): number of feature maps in the first conv layer of the encoder; default: 64
        num_groups (int): number of groups for the GroupNorm
        number_of_fmas (int): depth of the model
        scale_factor (int or tuple or list): Scale up and scale down factors.
    Notes:
        input data shape : (batchsize,channels,height,width,depth) or (batch,channels,x,y,z)
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 final_sigmoid,
                 f_maps=64,
                 layer_order='crg',
                 num_groups=8,
                 number_of_fmaps=4,
                 scale_factor=2,
                 **kwargs):
        super(UNet3D_HDC, self).__init__()

        self.in_hdc = HDC(scale_factor)

        if isinstance(f_maps, int):
            # use number_of_fmaps levels in the encoder path as suggested in the paper
            f_maps = create_feature_maps(f_maps,
                                         number_of_fmaps=number_of_fmaps)

        # create encoder path consisting of Encoder modules. The length of the encoder is equal to `len(f_maps)`
        # uses DoubleConv as a basic_module for the Encoder
        encoders = []
        for i, out_feature_num in enumerate(f_maps):
            if i == 0:
                encoder = Encoder(in_channels,
                                  out_feature_num,
                                  apply_pooling=False,  # no pooling in the first layer
                                  basic_module=DoubleConv,
                                  conv_layer_order=layer_order,
                                  num_groups=num_groups)
            else:
                encoder = Encoder(f_maps[i - 1],
                                  out_feature_num,
                                  basic_module=DoubleConv,
                                  conv_layer_order=layer_order,
                                  num_groups=num_groups)
            encoders.append(encoder)

        self.encoders = nn.ModuleList(encoders)

        # create decoder path consisting of the Decoder modules. The length of the decoder is equal to `len(f_maps) - 1`
        # uses DoubleConv as a basic_module for the Decoder
        decoders = []
        reversed_f_maps = list(reversed(f_maps))
        for i in range(len(reversed_f_maps) - 1):
            in_feature_num = reversed_f_maps[i] + reversed_f_maps[i + 1]
            out_feature_num = reversed_f_maps[i + 1]
            decoder = Decoder(in_feature_num,
                              out_feature_num,
                              basic_module=DoubleConv,
                              conv_layer_order=layer_order,
                              num_groups=num_groups)
            decoders.append(decoder)

        self.decoders = nn.ModuleList(decoders)

        # in the last layer a 1×1 convolution reduces the number of output
        # channels to the number of labels
        # self.final_conv = nn.Conv3d(f_maps[0], out_channels, 1)
        self.out_duc = DUC(scale_factor, out_channels, f_maps[0])
        if final_sigmoid:
            self.final_activation = nn.Sigmoid()
        else:
            self.final_activation = nn.Softmax(dim=1)

    def forward(self, x):
        # encoder part
        encoders_features = []
        x = self.in_hdc(x)
        for encoder in self.encoders:
            x = encoder(x)
            # reverse the encoder outputs to be aligned with the decoder
            encoders_features.insert(
                0, x)  # insert(index,object)

        # remove the last encoder's output from the list
        # !!remember: it's the 1st in the list
        encoders_features = encoders_features[1:]

        # decoder part
        for decoder, encoder_features in zip(self.decoders, encoders_features):
            # pass the output from the corresponding encoder and the output
            # of the previous decoder
            x = decoder(encoder_features, x)

        # x = self.final_conv(x)
        x = self.out_duc(x)
        # apply final_activation (i.e. Sigmoid or Softmax) only for prediction. During training the network outputs
        # logits and it's up to the user to normalize it before visualising with tensorboard or computing validation metric
        if not self.training:
            x = self.final_activation(x)  # training also output sigmoid value

        return x
# 3d unet model borrow from : https://github.com/wolny/pytorch-3dunet Create a list of modules with together constitute a single conv layer with non-linearity
     and optional batchnorm/groupnorm.
 
     Args:
         in_channels (int): number of input channels
         out_channels (int): number of output channels
         order (string): order of things, e.g.
             'cr' -> conv + ReLU
             'gcr' -> groupnorm + conv + ReLU
             'cl' -> conv + LeakyReLU
             'ce' -> conv + ELU
             'bcr' -> batchnorm + conv + ReLU
         num_groups (int): number of groups for the GroupNorm
         padding (int): add zero-padding to the input
 
     Return:
         list of tuple (name, module) # add learnable bias only in the absence of gatchnorm/groupnorm # use only one group if the given number of groups is greater than the number of channels Basic convolutional module consisting of a Conv3d, non-linearity and optional batchnorm/groupnorm. The order
     of operations can be specified via the `order` parameter
 
     Args:
         in_channels (int): number of input channels
         out_channels (int): number of output channels
         kernel_size (int): size of the convolving kernel
         order (string): determines the order of layers, e.g.
             'cr' -> conv + ReLU
             'crg' -> conv + ReLU + groupnorm
             'cl' -> conv + LeakyReLU
             'ce' -> conv + ELU
         num_groups (int): number of groups for the GroupNorm A module consisting of two consecutive convolution layers (e.g. BatchNorm3d+ReLU+Conv3d).
     We use (Conv3d+ReLU+GroupNorm3d) by default.
     This can be changed however by providing the 'order' argument, e.g. in order
     to change to Conv3d+BatchNorm3d+ELU use order='cbe'.
     Use padded convolutions to make sure that the output (H_out, W_out) is the same
     as (H_in, W_in), so that you don't have to crop in the decoder path.
 
     Args:
         in_channels (int): number of input channels
         out_channels (int): number of output channels
         encoder (bool): if True we're in the encoder path, otherwise we're in the decoder
         kernel_size (int): size of the convolving kernel
         order (string): determines the order of layers, e.g.
             'cr' -> conv + ReLU
             'crg' -> conv + ReLU + groupnorm
             'cl' -> conv + LeakyReLU
             'ce' -> conv + ELU
         num_groups (int): number of groups for the GroupNorm # we're in the encoder path # we're in the decoder path, decrease the number of channels in the 1st convolution # conv1 # conv2 Basic UNet block consisting of a SingleConv followed by the residual block.
     The SingleConv takes care of increasing/decreasing the number of channels and also ensures that the number
     of output channels is compatible with the residual block that follows.
     This block can be used instead of standard DoubleConv in the Encoder module.
     Motivated by: https://arxiv.org/pdf/1706.00120.pdf
 
     Notice we use ELU instead of ReLU (order='cge') and put non-linearity after the groupnorm. # first convolution # residual block # remove non-linearity from the 3rd convolution since it's going to be applied after adding the residual # create non-linearity separately # apply first convolution and save the output as a residual # residual block A single module from the encoder path consisting of the optional max
     pooling layer (one may specify the MaxPool kernel_size to be different
     than the standard (2,2,2), e.g. if the volumetric data is anisotropic
     (make sure to use complementary scale_factor in the decoder path) followed by
     a DoubleConv module.
     Args:
         in_channels (int): number of input channels
         out_channels (int): number of output channels
         conv_kernel_size (int): size of the convolving kernel
         apply_pooling (bool): if True use MaxPool3d before DoubleConv
         pool_kernel_size (tuple): the size of the window to take a max over
         pool_type (str): pooling layer: 'max' or 'avg'
         basic_module(nn.Module): either ResNetBlock or DoubleConv
         conv_layer_order (string): determines the order of layers
             in `DoubleConv` module. See `DoubleConv` for more info.
         num_groups (int): number of groups for the GroupNorm A single module for decoder path consisting of the upsample layer
     (either learned ConvTranspose3d or interpolation) followed by a DoubleConv
     module.
     Args:
         in_channels (int): number of input channels
         out_channels (int): number of output channels
         kernel_size (int): size of the convolving kernel
         scale_factor (tuple): used as the multiplier for the image H/W/D in
             case of nn.Upsample or as stride in case of ConvTranspose3d, must reverse the MaxPool3d operation
             from the corresponding encoder
         basic_module(nn.Module): either ResNetBlock or DoubleConv
         conv_layer_order (string): determines the order of layers
             in `DoubleConv` module. See `DoubleConv` for more info.
         num_groups (int): number of groups for the GroupNorm # if DoubleConv is the basic_module use nearest neighbor interpolation for upsampling # otherwise use ConvTranspose3d (bear in mind your GPU memory) # make sure that the output size reverses the MaxPool3d from the corresponding encoder # (D_out = (D_in − 1) ×  stride[0] − 2 ×  padding[0] +  kernel_size[0] +  output_padding[0]) # also scale the number of channels from in_channels to out_channels so that summation joining # works correctly # adapt the number of in_channels for the ExtResNetBlock # use nearest neighbor interpolation and concatenation joining # interpolation in decoder # concatenate encoder_features (encoder path) with the upsampled input across channel dimension # use ConvTranspose3d and summation joining A module consisting of a convolution layer (e.g. Conv3d+ReLU+GroupNorm3d) and the final 1x1 convolution
     which reduces the number of channels to 'out_channels'.
     with the number of output channels 'out_channels // 2' and 'out_channels' respectively.
     We use (Conv3d+ReLU+GroupNorm3d) by default.
     This can be change however by providing the 'order' argument, e.g. in order
     to change to Conv3d+BatchNorm3d+ReLU use order='cbr'.
     Args:
         in_channels (int): number of input channels
         out_channels (int): number of output channels
         kernel_size (int): size of the convolving kernel
         order (string): determines the order of layers, e.g.
             'cr' -> conv + ReLU
             'crg' -> conv + ReLU + groupnorm
         num_groups (int): number of groups for the GroupNorm # conv1 # in the last layer a 1×1 convolution reduces the number of output channels to out_channels 3DUnet model from
     "3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation"
         <https://arxiv.org/pdf/1606.0660.pdf>.
     HDC from:
     "Holistic decomposition convolution for effective semantic segmentation of medical volume images"
          <https://doi.org/10.1016/j.media.2019.07.003>.
     DUC from:
     "Real-time single image and video super-resolution using an efficient sub-pixel convolutional neural network"
         <https://doi.org/10.1109/CVPR.2016.207>.
     Args:
         in_channels (int): number of input channels
         out_channels (int): number of output segmentation masks;
             Note that that the of out_channels might correspond to either
             different semantic classes or to different binary segmentation mask.
             It's up to the user of the class to interpret the out_channels and
             use the proper loss criterion during training (i.e. CrossEntropyLoss (multi-class)
             or BCEWithLogitsLoss (two-class) respectively)
         f_maps (int, tuple): number of feature maps at each level of the encoder; if it's an integer the number
             of feature maps is given by the geometric progression: f_maps ^ k, k=1,2,3,4
         final_sigmoid (bool): if True apply element-wise nn.Sigmoid after the
             final 1x1 convolution, otherwise apply nn.Softmax. MUST be True if nn.BCELoss (two-class) is used
             to train the model. MUST be False if nn.CrossEntropyLoss (multi-class) is used to train the model.
         layer_order (string): determines the order of layers
             in `SingleConv` module. e.g. 'crg' stands for Conv3d+ReLU+GroupNorm3d.
             See `SingleConv` for more info
         init_channel_number (int): number of feature maps in the first conv layer of the encoder; default: 64
         num_groups (int): number of groups for the GroupNorm
         number_of_fmas (int): depth of the model
         scale_factor (int or tuple or list): Scale up and scale down factors.
     Notes:
         input data shape : (batchsize,channels,height,width,depth) or (batch,channels,x,y,z) # use number_of_fmaps levels in the encoder path as suggested in the paper # create encoder path consisting of Encoder modules. The length of the encoder is equal to `len(f_maps)` # uses DoubleConv as a basic_module for the Encoder # no pooling in the first layer # create decoder path consisting of the Decoder modules. The length of the decoder is equal to `len(f_maps) - 1` # uses DoubleConv as a basic_module for the Decoder # in the last layer a 1×1 convolution reduces the number of output # channels to the number of labels # self.final_conv = nn.Conv3d(f_maps[0], out_channels, 1) # encoder part # reverse the encoder outputs to be aligned with the decoder # insert(index,object) # remove the last encoder's output from the list # !!remember: it's the 1st in the list # decoder part # pass the output from the corresponding encoder and the output # of the previous decoder # x = self.final_conv(x) # apply final_activation (i.e. Sigmoid or Softmax) only for prediction. During training the network outputs # logits and it's up to the user to normalize it before visualising with tensorboard or computing validation metric # training also output sigmoid value