text: from functools import partial
  
  import torch
  from torch import nn as nn
  from torch.nn import functional as F
  
  
  def conv3d(in_channels, out_channels, kernel_size, bias, padding):
      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):
      """
      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
          kernel_size(int or tuple): size of the convolving kernel
          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 or tuple): add zero-padding added to all three sides of 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 batchnorm/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, f'Expected number of channels in input to be divisible by num_groups. num_channels={num_channels}, num_groups={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)))
          # vivian added
          elif char == 'd':
              modules.append(('dropout', nn.Dropout3d(p=0.1))) # vivian added 
          else:
              raise ValueError(f"Unsupported layer type '{char}'. MUST be one of ['b', 'g', 'r', 'l', 'e', 'c']")
  
      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 or tuple): 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
          padding (int or tuple):
      """
  
      def __init__(self, in_channels, out_channels, kernel_size=3, order='gcr', 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 or tuple): 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
          padding (int or tuple): add zero-padding added to all three sides of the input
      """
  
      def __init__(self, in_channels, out_channels, encoder, kernel_size=3, order='gcr', num_groups=8, padding=1):
          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,
                                     padding=padding))
          # conv2
          self.add_module('SingleConv2',
                          SingleConv(conv2_in_channels, conv2_out_channels, kernel_size, order, num_groups,
                                     padding=padding))
  
  
  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 or tuple): size of the convolving kernel
          apply_pooling (bool): if True use MaxPool3d before DoubleConv
          pool_kernel_size (int or tuple): the size of the window
          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
          padding (int or tuple): add zero-padding added to all three sides of the input
      """
  
      def __init__(self, in_channels, out_channels, conv_kernel_size=3, apply_pooling=True,
                   pool_kernel_size=2, pool_type='max', basic_module=DoubleConv, conv_layer_order='gcr',
                   num_groups=8, padding=1):
          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,
                                           padding=padding)
  
      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 upsampling layer
      (either learned ConvTranspose3d or nearest neighbor interpolation) followed by a basic module (DoubleConv or ExtResNetBlock).
      Args:
          in_channels (int): number of input channels
          out_channels (int): number of output channels
          conv_kernel_size (int or tuple): 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
          padding (int or tuple): add zero-padding added to all three sides of the input
      """
  
      def __init__(self, in_channels, out_channels, conv_kernel_size=3, scale_factor=(2, 2, 2), basic_module=DoubleConv,
                   conv_layer_order='gcr', num_groups=8, mode='nearest', padding=1):
          super(Decoder, self).__init__()
          if basic_module == DoubleConv:
              # if DoubleConv is the basic_module use interpolation for upsampling and concatenation joining
              self.upsampling = Upsampling(transposed_conv=False, in_channels=in_channels, out_channels=out_channels,
                                           kernel_size=conv_kernel_size, scale_factor=scale_factor, mode=mode)
              # concat joining
              self.joining = partial(self._joining, concat=True)
          else:
              # if basic_module=ExtResNetBlock use transposed convolution upsampling and summation joining
              self.upsampling = Upsampling(transposed_conv=True, in_channels=in_channels, out_channels=out_channels,
                                           kernel_size=conv_kernel_size, scale_factor=scale_factor, mode=mode)
              # sum joining
              self.joining = partial(self._joining, concat=False)
              # 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=conv_kernel_size,
                                           order=conv_layer_order,
                                           num_groups=num_groups,
                                           padding=padding)
  
      def forward(self, encoder_features, x):
          x = self.upsampling(encoder_features=encoder_features, x=x)
          x = self.joining(encoder_features, x)
          x = self.basic_module(x)
          return x
  
      @staticmethod
      def _joining(encoder_features, x, concat):
          if concat:
              return torch.cat((encoder_features, x), dim=1)
          else:
              return encoder_features + x
  
  
  class Upsampling(nn.Module):
      """
      Upsamples a given multi-channel 3D data using either interpolation or learned transposed convolution.
  
      Args:
          transposed_conv (bool): if True uses ConvTranspose3d for upsampling, otherwise uses interpolation
          in_channels (int): number of input channels for transposed conv
              used only if transposed_conv is True
          out_channels (int): number of output channels for transpose conv
              used only if transposed_conv is True
          kernel_size (int or tuple): size of the convolving kernel
              used only if transposed_conv is True
          scale_factor (int or tuple): stride of the convolution
              used only if transposed_conv is True
          mode (str): algorithm used for upsampling:
              'nearest' | 'linear' | 'bilinear' | 'trilinear' | 'area'. Default: 'nearest'
              used only if transposed_conv is False
      """
  
      def __init__(self, transposed_conv, in_channels=None, out_channels=None, kernel_size=3,
                   scale_factor=(2, 2, 2), mode='nearest'):
          super(Upsampling, self).__init__()
  
          if transposed_conv:
              # 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])
              self.upsample = nn.ConvTranspose3d(in_channels, out_channels, kernel_size=kernel_size, stride=scale_factor,
                                                 padding=1)
          else:
              self.upsample = partial(self._interpolate, mode=mode)
  
      def forward(self, encoder_features, x):
          output_size = encoder_features.size()[2:]
          return self.upsample(x, output_size)
  
  <<<<<<< HEAD:unet3d/buildingblocks.py
          # 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)
          
  =======
      @staticmethod
      def _interpolate(x, size, mode):
          return F.interpolate(x, size=size, mode=mode)
  >>>>>>> 62e10674c5cf9f44e252297203213b8c7f23c5f7:pytorch3dunet/unet3d/buildingblocks.py
  