text: # Copyright (c) 2020 Mitsubishi Electric Research Laboratories (MERL). All rights reserved. The software, documentation and/or data in this file is provided on an "as is" basis, and MERL has no obligations to provide maintenance, support, updates, enhancements or modifications. MERL specifically disclaims any warranties, including, but not limited to, the implied warranties of merchantability and fitness for any particular purpose. In no event shall MERL be liable to any party for direct, indirect, special, incidental, or consequential damages, including lost profits, arising out of the use of this software and its documentation, even if MERL has been advised of the possibility of such damages. As more fully described in the license agreement that was required in order to download this software, documentation and/or data, permission to use, copy and modify this software without fee is granted, but only for educational, research and non-commercial purposes.
  from numpy.core import shape_base
  import torch.nn.functional as F
  import torch.nn as nn
  import torch
  import ipdb
  
  class forecast_lstm(nn.Module):
      def __init__(self):
          super(forecast_lstm, self).__init__()
          self.embedding_dim = 4096
          self.hidden_size = 4096
          # self.proj_size = 32768
          self.lstm_layer = nn.LSTMCell(input_size=self.embedding_dim,hidden_size=self.hidden_size)
          # self.conv3d_1 = nn.Conv3D(64, 64, kernel_size=(3, 3, 3), stride=1, padding=(0, 0, 0))
          # self.conv3d_2 = nn.Conv3D(128, 128, kernel_size=(3, 3, 3), stride=1, padding=(0, 0, 0))
          self.linear_1 = nn.Linear(256, 32)
          self.linear_2 = nn.Linear(32,4)
          self.linear_3 = nn.Linear(4, 32)
          self.linear_4 = nn.Linear(32,256)
          # self.linear_5 = nn.Linear(4,16)
          # self.linear_6 = nn.Linear(16,4)
          
      def forward(self, x, delta_t):
          x = x.permute(0, 2, 3, 1)
          x = self.linear_1(x)
          x = F.relu(x)
          x = self.linear_2(x)
          x = F.relu(x)
          # x = self.linear_6(x)
          # x = F.relu(x)
          shape_a, shape_b, shape_c, shape_d = x.shape
          x = x.reshape(shape_a, 1, shape_b * shape_c * shape_d)
          for i in range(shape_a):
              x_temp = x[i]
              if i != 0:
                  (h_temp, c_temp) = self.lstm_layer(x_temp, (h_temp, c_temp))
              else:
                  (h_temp, c_temp) = self.lstm_layer(x_temp)
  
          # if delta_t < self.num_layers:
          #     x = x[1][int(delta_t)]
          # else:
          #     x = x[1][self.num_layers - 1]
  
          for j in range(int(delta_t)):
              x_temp = c_temp
              (h_temp, c_temp) = self.lstm_layer(x_temp, (h_temp, c_temp))
          x = c_temp
          x = x.reshape(32,32,4)
          # x = self.linear_5(x)
          # x = F.relu(x)
          x = self.linear_3(x)
          x = F.relu(x)
          x = self.linear_4(x)
          x = F.relu(x)
          x = x.permute(2, 0, 1)
          x = torch.unsqueeze(x, 0)
          return x
  
  class CellClassification(nn.Module):
      def __init__(self, category_num=5):
          super(CellClassification, self).__init__()
          self.conv1 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
          self.conv2 = nn.Conv2d(32, category_num, kernel_size=1, stride=1, padding=0)
  
          self.bn1 = nn.BatchNorm2d(32)
  
      def forward(self, x):
          x = F.relu(self.bn1(self.conv1(x)))
          x = self.conv2(x)
          return x
  
  
  class StateEstimation(nn.Module):
      def __init__(self, motion_category_num=2):
          super(StateEstimation, self).__init__()
          self.conv1 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
          self.conv2 = nn.Conv2d(32, motion_category_num, kernel_size=1, stride=1, padding=0)
  
          self.bn1 = nn.BatchNorm2d(32)
  
      def forward(self, x):
          x = F.relu(self.bn1(self.conv1(x)))
          x = self.conv2(x)
  
          return x
  
  
  class MotionPrediction(nn.Module):
      def __init__(self, seq_len):
          super(MotionPrediction, self).__init__()
          self.conv1 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
          self.conv2 = nn.Conv2d(32, 2 * seq_len, kernel_size=1, stride=1, padding=0)
  
          self.bn1 = nn.BatchNorm2d(32)
  
      def forward(self, x):
          x = F.relu(self.bn1(self.conv1(x)))
          x = self.conv2(x)
  
          return x
  
  
  class Conv3D(nn.Module):
      def __init__(self, in_channel, out_channel, kernel_size, stride, padding):
          super(Conv3D, self).__init__()
          self.conv3d = nn.Conv3d(in_channel, out_channel, kernel_size=kernel_size, stride=stride, padding=padding)
          self.bn3d = nn.BatchNorm3d(out_channel)
  
      def forward(self, x):
          # input x: (batch, seq, c, h, w)
          x = x.permute(0, 2, 1, 3, 4).contiguous()  # (batch, c, seq_len, h, w)
          x = F.relu(self.bn3d(self.conv3d(x)))
          x = x.permute(0, 2, 1, 3, 4).contiguous()  # (batch, seq_len, c, h, w)
  
          return x
  
  class MapExtractor(nn.Module):
      def __init__(self, map_channel=8):
          super(MapExtractor, self).__init__()
          self.conv_pre_1 = nn.Conv2d(map_channel, 16, kernel_size=3, stride=1, padding=1)
          self.conv_pre_2 = nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=1)
          self.bn_pre_1 = nn.BatchNorm2d(16)
          self.bn_pre_2 = nn.BatchNorm2d(16)
  
          self.conv1_1 = nn.Conv2d(16, 32, kernel_size=3, stride=2, padding=1)
          self.conv1_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
  
          self.conv2_1 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
          self.conv2_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
  
          # self.conv3_1 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
          # self.conv3_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
  
          # self.conv4_1 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
          # self.conv4_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
  
          # self.conv5_1 = nn.Conv2d(256 + 128, 128, kernel_size=3, stride=1, padding=1)
          # self.conv5_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
  
          # self.conv6_1 = nn.Conv2d(128 + 64, 64, kernel_size=3, stride=1, padding=1)
          # self.conv6_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
  
          self.conv7_1 = nn.Conv2d(64 + 32, 32, kernel_size=3, stride=1, padding=1)
          self.conv7_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
  
          self.conv8_1 = nn.Conv2d(32 + 16, 16, kernel_size=3, stride=1, padding=1)
          self.conv8_2 = nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=1)
  
          self.bn1_1 = nn.BatchNorm2d(32)
          self.bn1_2 = nn.BatchNorm2d(32)
  
          self.bn2_1 = nn.BatchNorm2d(64)
          self.bn2_2 = nn.BatchNorm2d(64)
  
          self.bn3_1 = nn.BatchNorm2d(128)
          self.bn3_2 = nn.BatchNorm2d(128)
  
          self.bn4_1 = nn.BatchNorm2d(256)
          self.bn4_2 = nn.BatchNorm2d(256)
  
          self.bn5_1 = nn.BatchNorm2d(128)
          self.bn5_2 = nn.BatchNorm2d(128)
  
          self.bn6_1 = nn.BatchNorm2d(64)
          self.bn6_2 = nn.BatchNorm2d(64)
  
          self.bn7_1 = nn.BatchNorm2d(32)
          self.bn7_2 = nn.BatchNorm2d(32)
  
          self.bn8_1 = nn.BatchNorm2d(16)
          self.bn8_2 = nn.BatchNorm2d(16)
  
      def forward(self, x):
  
          x = x.view(-1, x.size(-3), x.size(-2), x.size(-1))
          x = F.relu(self.bn_pre_1(self.conv_pre_1(x)))
          x = F.relu(self.bn_pre_2(self.conv_pre_2(x)))
  
  
          # -------------------------------- Encoder Path --------------------------------
          # -- STC block 1
          x_1 = F.relu(self.bn1_1(self.conv1_1(x)))
          x_1 = F.relu(self.bn1_2(self.conv1_2(x_1)))
  
          # -- STC block 2
          x_2 = F.relu(self.bn2_1(self.conv2_1(x_1)))
          x_2 = F.relu(self.bn2_2(self.conv2_2(x_2)))
  
          # # -- STC block 3
          # x_3 = F.relu(self.bn3_1(self.conv3_1(x_2)))
          # x_3 = F.relu(self.bn3_2(self.conv3_2(x_3)))
  
          # # -- STC block 4
          # x_4 = F.relu(self.bn4_1(self.conv4_1(x_3)))
          # x_4 = F.relu(self.bn4_2(self.conv4_2(x_4)))
  
  
          # -------------------------------- Decoder Path --------------------------------
          # x_5 = F.relu(self.bn5_1(self.conv5_1(torch.cat((F.interpolate(x_4, scale_factor=(2, 2)), x_3), dim=1))))
          # x_5 = F.relu(self.bn5_2(self.conv5_2(x_5)))
  
          # x_6 = F.relu(self.bn6_1(self.conv6_1(torch.cat((F.interpolate(x_5, scale_factor=(2, 2)), x_2), dim=1))))
          # x_6 = F.relu(self.bn6_2(self.conv6_2(x_6)))
  
          x_7 = F.relu(self.bn7_1(self.conv7_1(torch.cat((F.interpolate(x_2, scale_factor=(2, 2)), x_1), dim=1))))
          x_7 = F.relu(self.bn7_2(self.conv7_2(x_7)))
  
          x_8 = F.relu(self.bn8_1(self.conv8_1(torch.cat((F.interpolate(x_1, scale_factor=(2, 2)), x), dim=1))))
          res_x = F.relu(self.bn8_2(self.conv8_2(x_8)))
  
          return res_x
  
  class STPN(nn.Module):
      def __init__(self, height_feat_size=13):
          super(STPN, self).__init__()
          self.conv_pre_1 = nn.Conv2d(height_feat_size, 32, kernel_size=3, stride=1, padding=1)
          self.conv_pre_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
          self.bn_pre_1 = nn.BatchNorm2d(32)
          self.bn_pre_2 = nn.BatchNorm2d(32)
  
  #        self.conv3d_1 = Conv3D(64, 64, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))
  #        self.conv3d_2 = Conv3D(128, 128, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))
          self.conv3d_1 = Conv3D(64, 64, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
          self.conv3d_2 = Conv3D(128, 128, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
  
          self.conv1_1 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
          self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
  
          self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
          self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
  
          self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
          self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
  
          self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)
          self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
  
          self.conv5_1 = nn.Conv2d(512 + 256, 256, kernel_size=3, stride=1, padding=1)
          self.conv5_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
  
          self.conv6_1 = nn.Conv2d(256 + 128, 128, kernel_size=3, stride=1, padding=1)
          self.conv6_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
  
          self.conv7_1 = nn.Conv2d(128 + 64, 64, kernel_size=3, stride=1, padding=1)
          self.conv7_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
  
          self.conv8_1 = nn.Conv2d(64 + 32, 32, kernel_size=3, stride=1, padding=1)
          self.conv8_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
  
          self.bn1_1 = nn.BatchNorm2d(64)
          self.bn1_2 = nn.BatchNorm2d(64)
  
          self.bn2_1 = nn.BatchNorm2d(128)
          self.bn2_2 = nn.BatchNorm2d(128)
  
          self.bn3_1 = nn.BatchNorm2d(256)
          self.bn3_2 = nn.BatchNorm2d(256)
  
          self.bn4_1 = nn.BatchNorm2d(512)
          self.bn4_2 = nn.BatchNorm2d(512)
  
          self.bn5_1 = nn.BatchNorm2d(256)
          self.bn5_2 = nn.BatchNorm2d(256)
  
          self.bn6_1 = nn.BatchNorm2d(128)
          self.bn6_2 = nn.BatchNorm2d(128)
  
          self.bn7_1 = nn.BatchNorm2d(64)
          self.bn7_2 = nn.BatchNorm2d(64)
  
          self.bn8_1 = nn.BatchNorm2d(32)
          self.bn8_2 = nn.BatchNorm2d(32)
  
      def forward(self, x):
          batch, seq, z, h, w = x.size()
  
          x = x.view(-1, x.size(-3), x.size(-2), x.size(-1))
          x = F.relu(self.bn_pre_1(self.conv_pre_1(x)))
          x = F.relu(self.bn_pre_2(self.conv_pre_2(x)))
  
          # -------------------------------- Encoder Path --------------------------------
          # -- STC block 1
          x_1 = F.relu(self.bn1_1(self.conv1_1(x)))
          x_1 = F.relu(self.bn1_2(self.conv1_2(x_1)))
  
          x_1 = x_1.view(batch, -1, x_1.size(1), x_1.size(2), x_1.size(3)).contiguous()  # (batch, seq, c, h, w)
          x_1 = self.conv3d_1(x_1)
          x_1 = x_1.view(-1, x_1.size(2), x_1.size(3), x_1.size(4)).contiguous()  # (batch * seq, c, h, w)
  
          # -- STC block 2
          x_2 = F.relu(self.bn2_1(self.conv2_1(x_1)))
          x_2 = F.relu(self.bn2_2(self.conv2_2(x_2)))
  
          x_2 = x_2.view(batch, -1, x_2.size(1), x_2.size(2), x_2.size(3)).contiguous()  # (batch, seq, c, h, w)
          x_2 = self.conv3d_2(x_2)
          x_2 = x_2.view(-1, x_2.size(2), x_2.size(3), x_2.size(4)).contiguous()  # (batch * seq, c, h, w), seq = 1
  
          # -- STC block 3
          x_3 = F.relu(self.bn3_1(self.conv3_1(x_2)))
          x_3 = F.relu(self.bn3_2(self.conv3_2(x_3)))
  
          # -- STC block 4
          x_4 = F.relu(self.bn4_1(self.conv4_1(x_3)))
          x_4 = F.relu(self.bn4_2(self.conv4_2(x_4)))
  
  
          # -------------------------------- Decoder Path --------------------------------
          x_5 = F.relu(self.bn5_1(self.conv5_1(torch.cat((F.interpolate(x_4, scale_factor=(2, 2)), x_3), dim=1))))
          x_5 = F.relu(self.bn5_2(self.conv5_2(x_5)))
  
          x_2 = x_2.view(batch, -1, x_2.size(1), x_2.size(2), x_2.size(3))
          x_2 = x_2.permute(0, 2, 1, 3, 4).contiguous()
          x_2 = F.adaptive_max_pool3d(x_2, (1, None, None))
          x_2 = x_2.permute(0, 2, 1, 3, 4).contiguous()
          x_2 = x_2.view(-1, x_2.size(2), x_2.size(3), x_2.size(4)).contiguous()
  
          x_6 = F.relu(self.bn6_1(self.conv6_1(torch.cat((F.interpolate(x_5, scale_factor=(2, 2)), x_2), dim=1))))
          x_6 = F.relu(self.bn6_2(self.conv6_2(x_6)))
  
          x_1 = x_1.view(batch, -1, x_1.size(1), x_1.size(2), x_1.size(3))
          x_1 = x_1.permute(0, 2, 1, 3, 4).contiguous()
          x_1 = F.adaptive_max_pool3d(x_1, (1, None, None))
          x_1 = x_1.permute(0, 2, 1, 3, 4).contiguous()
          x_1 = x_1.view(-1, x_1.size(2), x_1.size(3), x_1.size(4)).contiguous()
  
          x_7 = F.relu(self.bn7_1(self.conv7_1(torch.cat((F.interpolate(x_6, scale_factor=(2, 2)), x_1), dim=1))))
          x_7 = F.relu(self.bn7_2(self.conv7_2(x_7)))
  
          x = x.view(batch, -1, x.size(1), x.size(2), x.size(3))
          x = x.permute(0, 2, 1, 3, 4).contiguous()
          x = F.adaptive_max_pool3d(x, (1, None, None))
          x = x.permute(0, 2, 1, 3, 4).contiguous()
          x = x.view(-1, x.size(2), x.size(3), x.size(4)).contiguous()
  
          x_8 = F.relu(self.bn8_1(self.conv8_1(torch.cat((F.interpolate(x_7, scale_factor=(2, 2)), x), dim=1))))
          res_x = F.relu(self.bn8_2(self.conv8_2(x_8)))
  
          return res_x
  
  class STPN_KD(nn.Module):
      def __init__(self, height_feat_size=13):
          super(STPN_KD, self).__init__()
          self.conv_pre_1 = nn.Conv2d(height_feat_size, 32, kernel_size=3, stride=1, padding=1)
          self.conv_pre_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
          self.bn_pre_1 = nn.BatchNorm2d(32)
          self.bn_pre_2 = nn.BatchNorm2d(32)
  
  #        self.conv3d_1 = Conv3D(64, 64, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))
  #        self.conv3d_2 = Conv3D(128, 128, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))
          self.conv3d_1 = Conv3D(64, 64, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
          self.conv3d_2 = Conv3D(128, 128, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
  
          self.conv1_1 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
          self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
  
          self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
          self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
  
          self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
          self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
  
          self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)
          self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
  
          self.conv5_1 = nn.Conv2d(512 + 256, 256, kernel_size=3, stride=1, padding=1)
          self.conv5_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
  
          self.conv6_1 = nn.Conv2d(256 + 128, 128, kernel_size=3, stride=1, padding=1)
          self.conv6_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
  
          self.conv7_1 = nn.Conv2d(128 + 64, 64, kernel_size=3, stride=1, padding=1)
          self.conv7_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
  
          self.conv8_1 = nn.Conv2d(64 + 32, 32, kernel_size=3, stride=1, padding=1)
          self.conv8_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
  
          self.bn1_1 = nn.BatchNorm2d(64)
          self.bn1_2 = nn.BatchNorm2d(64)
  
          self.bn2_1 = nn.BatchNorm2d(128)
          self.bn2_2 = nn.BatchNorm2d(128)
  
          self.bn3_1 = nn.BatchNorm2d(256)
          self.bn3_2 = nn.BatchNorm2d(256)
  
          self.bn4_1 = nn.BatchNorm2d(512)
          self.bn4_2 = nn.BatchNorm2d(512)
  
          self.bn5_1 = nn.BatchNorm2d(256)
          self.bn5_2 = nn.BatchNorm2d(256)
  
          self.bn6_1 = nn.BatchNorm2d(128)
          self.bn6_2 = nn.BatchNorm2d(128)
  
          self.bn7_1 = nn.BatchNorm2d(64)
          self.bn7_2 = nn.BatchNorm2d(64)
  
          self.bn8_1 = nn.BatchNorm2d(32)
          self.bn8_2 = nn.BatchNorm2d(32)
  
      def forward(self, x):
          batch, seq, z, h, w = x.size()
  
          x = x.view(-1, x.size(-3), x.size(-2), x.size(-1))
          x = F.relu(self.bn_pre_1(self.conv_pre_1(x)))
          x = F.relu(self.bn_pre_2(self.conv_pre_2(x)))
  
          # -------------------------------- Encoder Path --------------------------------
          # -- STC block 1
          x_1 = F.relu(self.bn1_1(self.conv1_1(x)))
          x_1 = F.relu(self.bn1_2(self.conv1_2(x_1)))
  
          x_1 = x_1.view(batch, -1, x_1.size(1), x_1.size(2), x_1.size(3)).contiguous()  # (batch, seq, c, h, w)
          x_1 = self.conv3d_1(x_1)
          x_1 = x_1.view(-1, x_1.size(2), x_1.size(3), x_1.size(4)).contiguous()  # (batch * seq, c, h, w)
  
          # -- STC block 2
          x_2 = F.relu(self.bn2_1(self.conv2_1(x_1)))
          x_2 = F.relu(self.bn2_2(self.conv2_2(x_2)))
  
          x_2 = x_2.view(batch, -1, x_2.size(1), x_2.size(2), x_2.size(3)).contiguous()  # (batch, seq, c, h, w)
          x_2 = self.conv3d_2(x_2)
          x_2 = x_2.view(-1, x_2.size(2), x_2.size(3), x_2.size(4)).contiguous()  # (batch * seq, c, h, w), seq = 1
  
          # -- STC block 3
          x_3 = F.relu(self.bn3_1(self.conv3_1(x_2)))
          x_3 = F.relu(self.bn3_2(self.conv3_2(x_3)))
  
          # -- STC block 4
          x_4 = F.relu(self.bn4_1(self.conv4_1(x_3)))
          x_4 = F.relu(self.bn4_2(self.conv4_2(x_4)))
  
  
          # -------------------------------- Decoder Path --------------------------------
          x_5 = F.relu(self.bn5_1(self.conv5_1(torch.cat((F.interpolate(x_4, scale_factor=(2, 2)), x_3), dim=1))))
          x_5 = F.relu(self.bn5_2(self.conv5_2(x_5)))
  
          x_2 = x_2.view(batch, -1, x_2.size(1), x_2.size(2), x_2.size(3))
          x_2 = x_2.permute(0, 2, 1, 3, 4).contiguous()
          x_2 = F.adaptive_max_pool3d(x_2, (1, None, None))
          x_2 = x_2.permute(0, 2, 1, 3, 4).contiguous()
          x_2 = x_2.view(-1, x_2.size(2), x_2.size(3), x_2.size(4)).contiguous()
  
          x_6 = F.relu(self.bn6_1(self.conv6_1(torch.cat((F.interpolate(x_5, scale_factor=(2, 2)), x_2), dim=1))))
          x_6 = F.relu(self.bn6_2(self.conv6_2(x_6)))
  
          x_1 = x_1.view(batch, -1, x_1.size(1), x_1.size(2), x_1.size(3))
          x_1 = x_1.permute(0, 2, 1, 3, 4).contiguous()
          x_1 = F.adaptive_max_pool3d(x_1, (1, None, None))
          x_1 = x_1.permute(0, 2, 1, 3, 4).contiguous()
          x_1 = x_1.view(-1, x_1.size(2), x_1.size(3), x_1.size(4)).contiguous()
  
          x_7 = F.relu(self.bn7_1(self.conv7_1(torch.cat((F.interpolate(x_6, scale_factor=(2, 2)), x_1), dim=1))))
          x_7 = F.relu(self.bn7_2(self.conv7_2(x_7)))
  
          x = x.view(batch, -1, x.size(1), x.size(2), x.size(3))
          x = x.permute(0, 2, 1, 3, 4).contiguous()
          x = F.adaptive_max_pool3d(x, (1, None, None))
          x = x.permute(0, 2, 1, 3, 4).contiguous()
          x = x.view(-1, x.size(2), x.size(3), x.size(4)).contiguous()
  
          x_8 = F.relu(self.bn8_1(self.conv8_1(torch.cat((F.interpolate(x_7, scale_factor=(2, 2)), x), dim=1))))
          res_x = F.relu(self.bn8_2(self.conv8_2(x_8)))
  
          return res_x, x_7, x_6, x_5, x_3
  
  
  class MotionNet(nn.Module):
      def __init__(self, out_seq_len=20, motion_category_num=2, height_feat_size=13):
          super(MotionNet, self).__init__()
          self.out_seq_len = out_seq_len
  
          self.cell_classify = CellClassification()
          self.motion_pred = MotionPrediction(seq_len=self.out_seq_len)
          self.state_classify = StateEstimation(motion_category_num=motion_category_num)
          self.stpn = STPN(height_feat_size=height_feat_size)
  
  
      def forward(self, bevs):
          bevs = bevs.permute(0, 1, 4, 2, 3)  # (Batch, seq, z, h, w)
  
          # Backbone network
          x = self.stpn(bevs)
  
          # Cell Classification head
          cell_class_pred = self.cell_classify(x)
  
          # Motion State Classification head
          state_class_pred = self.state_classify(x)
  
          # Motion Displacement prediction
          disp = self.motion_pred(x)
          disp = disp.view(-1, 2, x.size(-2), x.size(-1))
  
          return disp, cell_class_pred, state_class_pred
  
  
  # For MGDA loss computation
  class FeatEncoder(nn.Module):
      def __init__(self, height_feat_size=13):
          super(FeatEncoder, self).__init__()
          self.stpn = STPN(height_feat_size=height_feat_size)
  
      def forward(self, bevs):
          bevs = bevs.permute(0, 1, 4, 2, 3)  # (Batch, seq, z, h, w)
          x = self.stpn(bevs)
  
          return x
  
  
  class MotionNetMGDA(nn.Module):
      def __init__(self, out_seq_len=20, motion_category_num=2):
          super(MotionNetMGDA, self).__init__()
          self.out_seq_len = out_seq_len
  
          self.cell_classify = CellClassification()
          self.motion_pred = MotionPrediction(seq_len=self.out_seq_len)
          self.state_classify = StateEstimation(motion_category_num=motion_category_num)
  
      def forward(self, stpn_out):
          # Cell Classification head
          cell_class_pred = self.cell_classify(stpn_out)
  
          # Motion State Classification head
          state_class_pred = self.state_classify(stpn_out)
  
          # Motion Displacement prediction
          disp = self.motion_pred(stpn_out)
          disp = disp.view(-1, 2, stpn_out.size(-2), stpn_out.size(-1))
  
          return disp, cell_class_pred, state_class_pred
  
  '''''''''''''''''''''
  Added by Yiming
  
  '''''''''''''''''''''
  class conv2DBatchNormRelu(nn.Module):
      def __init__(
          self,
          in_channels,
          n_filters,
          k_size,
          stride,
          padding,
          bias=True,
          dilation=1,
          is_batchnorm=True,
      ):
          super(conv2DBatchNormRelu, self).__init__()
  
          conv_mod = nn.Conv2d(
              int(in_channels),
              int(n_filters),
              kernel_size=k_size,
              padding=padding,
              stride=stride,
              bias=bias,
              dilation=dilation,
          )
  
          if is_batchnorm:
              self.cbr_unit = nn.Sequential(
                  conv_mod, nn.BatchNorm2d(int(n_filters)), nn.ReLU(inplace=True)
              )
          else:
              self.cbr_unit = nn.Sequential(conv_mod, nn.ReLU(inplace=True))
  
      def forward(self, inputs):
          outputs = self.cbr_unit(inputs)
          return outputs
  
  
  class Sparsemax(nn.Module):
      """Sparsemax function."""
  
      def __init__(self, dim=None):
          """Initialize sparsemax activation
          
          Args:
              dim (int, optional): The dimension over which to apply the sparsemax function.
          """
          super(Sparsemax, self).__init__()
  
          self.dim = -1 if dim is None else dim
  
      def forward(self, input):
          """Forward function.
          Args:
              input (torch.Tensor): Input tensor. First dimension should be the batch size
          Returns:
              torch.Tensor: [batch_size x number_of_logits] Output tensor
          """
          # Sparsemax currently only handles 2-dim tensors,
          # so we reshape and reshape back after sparsemax
          original_size = input.size()
          input = input.view(-1, input.size(self.dim))
          
          dim = 1
          number_of_logits = input.size(dim)
  
          # Translate input by max for numerical stability
          input = input - torch.max(input, dim=dim, keepdim=True)[0].expand_as(input)
  
          # Sort input in descending order.
          # (NOTE: Can be replaced with linear time selection method described here:
          # http://stanford.edu/~jduchi/projects/DuchiShSiCh08.html)
          zs = torch.sort(input=input, dim=dim, descending=True)[0]
          range = torch.range(start=1, end=number_of_logits, device=input.device).view(1, -1)
          range = range.expand_as(zs)
  
          # Determine sparsity of projection
          bound = 1 + range * zs
          cumulative_sum_zs = torch.cumsum(zs, dim)
          is_gt = torch.gt(bound, cumulative_sum_zs).type(input.type())
          k = torch.max(is_gt * range, dim, keepdim=True)[0]
  
          # Compute threshold function
          zs_sparse = is_gt * zs
  
          # Compute taus
          taus = (torch.sum(zs_sparse, dim, keepdim=True) - 1) / k
          taus = taus.expand_as(input)
  
          # Sparsemax
          self.output = torch.max(torch.zeros_like(input), input - taus)
  
          output = self.output.view(original_size)
  
          return output
  
      def backward(self, grad_output):
          """Backward function."""
          dim = 1
  
          nonzeros = torch.ne(self.output, 0)
          sum = torch.sum(grad_output * nonzeros, dim=dim) / torch.sum(nonzeros, dim=dim)
          self.grad_input = nonzeros * (grad_output - sum.expand_as(grad_output))
  
          return self.grad_input
  
  
  class lidar_encoder(nn.Module):
      def __init__(self, height_feat_size=13):
          super(lidar_encoder, self).__init__()
          self.conv_pre_1 = nn.Conv2d(height_feat_size, 32, kernel_size=3, stride=1, padding=1)
          self.conv_pre_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
          self.bn_pre_1 = nn.BatchNorm2d(32)
          self.bn_pre_2 = nn.BatchNorm2d(32)
  
  #        self.conv3d_1 = Conv3D(64, 64, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))
  #        self.conv3d_2 = Conv3D(128, 128, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))
          self.conv3d_1 = Conv3D(64, 64, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
          self.conv3d_2 = Conv3D(128, 128, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
  
          self.conv1_1 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
          self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
  
          self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
          self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
  
          self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
          self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
  
          self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)
          self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
  
          self.conv5_1 = nn.Conv2d(512 + 256, 256, kernel_size=3, stride=1, padding=1)
          self.conv5_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
  
          self.conv6_1 = nn.Conv2d(256 + 128, 128, kernel_size=3, stride=1, padding=1)
          self.conv6_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
  
          self.conv7_1 = nn.Conv2d(128 + 64, 64, kernel_size=3, stride=1, padding=1)
          self.conv7_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
  
          self.conv8_1 = nn.Conv2d(64 + 32, 32, kernel_size=3, stride=1, padding=1)
          self.conv8_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
  
          self.bn1_1 = nn.BatchNorm2d(64)
          self.bn1_2 = nn.BatchNorm2d(64)
  
          self.bn2_1 = nn.BatchNorm2d(128)
          self.bn2_2 = nn.BatchNorm2d(128)
  
          self.bn3_1 = nn.BatchNorm2d(256)
          self.bn3_2 = nn.BatchNorm2d(256)
  
          self.bn4_1 = nn.BatchNorm2d(512)
          self.bn4_2 = nn.BatchNorm2d(512)
  
          self.bn5_1 = nn.BatchNorm2d(256)
          self.bn5_2 = nn.BatchNorm2d(256)
  
          self.bn6_1 = nn.BatchNorm2d(128)
          self.bn6_2 = nn.BatchNorm2d(128)
  
          self.bn7_1 = nn.BatchNorm2d(64)
          self.bn7_2 = nn.BatchNorm2d(64)
  
          self.bn8_1 = nn.BatchNorm2d(32)
          self.bn8_2 = nn.BatchNorm2d(32)
  
      def forward(self, x):
          batch, seq, z, h, w = x.size()
  
          x = x.view(-1, x.size(-3), x.size(-2), x.size(-1))
          x = F.relu(self.bn_pre_1(self.conv_pre_1(x)))
          x = F.relu(self.bn_pre_2(self.conv_pre_2(x)))
  
          # -------------------------------- Encoder Path --------------------------------
          # -- STC block 1
          x_1 = F.relu(self.bn1_1(self.conv1_1(x)))
          x_1 = F.relu(self.bn1_2(self.conv1_2(x_1)))
  
          x_1 = x_1.view(batch, -1, x_1.size(1), x_1.size(2), x_1.size(3)).contiguous()  # (batch, seq, c, h, w)
          x_1 = self.conv3d_1(x_1)
          x_1 = x_1.view(-1, x_1.size(2), x_1.size(3), x_1.size(4)).contiguous()  # (batch * seq, c, h, w)
  
          # -- STC block 2
          x_2 = F.relu(self.bn2_1(self.conv2_1(x_1)))
          x_2 = F.relu(self.bn2_2(self.conv2_2(x_2)))
  
          x_2 = x_2.view(batch, -1, x_2.size(1), x_2.size(2), x_2.size(3)).contiguous()  # (batch, seq, c, h, w)
          x_2 = self.conv3d_2(x_2)
          x_2 = x_2.view(-1, x_2.size(2), x_2.size(3), x_2.size(4)).contiguous()  # (batch * seq, c, h, w), seq = 1
  
          # -- STC block 3
          x_3 = F.relu(self.bn3_1(self.conv3_1(x_2)))
          x_3 = F.relu(self.bn3_2(self.conv3_2(x_3)))
  
          # -- STC block 4
          x_4 = F.relu(self.bn4_1(self.conv4_1(x_3)))
          x_4 = F.relu(self.bn4_2(self.conv4_2(x_4)))
  
          return x, x_1, x_2, x_3, x_4
  
  class lidar_decoder(nn.Module):
      def __init__(self, height_feat_size=13):
          super(lidar_decoder, self).__init__()
          self.conv_pre_1 = nn.Conv2d(height_feat_size, 32, kernel_size=3, stride=1, padding=1)
          self.conv_pre_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
          self.bn_pre_1 = nn.BatchNorm2d(32)
          self.bn_pre_2 = nn.BatchNorm2d(32)
  
  #        self.conv3d_1 = Conv3D(64, 64, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))
  #        self.conv3d_2 = Conv3D(128, 128, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))
          self.conv3d_1 = Conv3D(64, 64, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
          self.conv3d_2 = Conv3D(128, 128, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
  
          self.conv1_1 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
          self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
  
          self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
          self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
  
          self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
          self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
  
          self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)
          self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
  
          self.conv5_1 = nn.Conv2d(512 + 256, 256, kernel_size=3, stride=1, padding=1)
          self.conv5_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
  
          self.conv6_1 = nn.Conv2d(256 + 128, 128, kernel_size=3, stride=1, padding=1)
          self.conv6_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
  
          self.conv7_1 = nn.Conv2d(128 + 64, 64, kernel_size=3, stride=1, padding=1)
          self.conv7_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
  
          self.conv8_1 = nn.Conv2d(64 + 32, 32, kernel_size=3, stride=1, padding=1)
          self.conv8_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
  
          self.bn1_1 = nn.BatchNorm2d(64)
          self.bn1_2 = nn.BatchNorm2d(64)
  
          self.bn2_1 = nn.BatchNorm2d(128)
          self.bn2_2 = nn.BatchNorm2d(128)
  
          self.bn3_1 = nn.BatchNorm2d(256)
          self.bn3_2 = nn.BatchNorm2d(256)
  
          self.bn4_1 = nn.BatchNorm2d(512)
          self.bn4_2 = nn.BatchNorm2d(512)
  
          self.bn5_1 = nn.BatchNorm2d(256)
          self.bn5_2 = nn.BatchNorm2d(256)
  
          self.bn6_1 = nn.BatchNorm2d(128)
          self.bn6_2 = nn.BatchNorm2d(128)
  
          self.bn7_1 = nn.BatchNorm2d(64)
          self.bn7_2 = nn.BatchNorm2d(64)
  
          self.bn8_1 = nn.BatchNorm2d(32)
          self.bn8_2 = nn.BatchNorm2d(32)
  
      def forward(self, x, x_1, x_2, x_3, x_4, batch):
          # -------------------------------- Decoder Path --------------------------------
          x_5 = F.relu(self.bn5_1(self.conv5_1(torch.cat((F.interpolate(x_4, scale_factor=(2, 2)), x_3), dim=1))))
          x_5 = F.relu(self.bn5_2(self.conv5_2(x_5)))
  
          x_2 = x_2.view(batch, -1, x_2.size(1), x_2.size(2), x_2.size(3))
          x_2 = x_2.permute(0, 2, 1, 3, 4).contiguous()
          #x_2 = F.adaptive_max_pool3d(x_2, (1, None, None))
          x_2 = x_2.permute(0, 2, 1, 3, 4).contiguous()
          x_2 = x_2.view(-1, x_2.size(2), x_2.size(3), x_2.size(4)).contiguous()
  
          x_6 = F.relu(self.bn6_1(self.conv6_1(torch.cat((F.interpolate(x_5, scale_factor=(2, 2)), x_2), dim=1))))
          x_6 = F.relu(self.bn6_2(self.conv6_2(x_6)))
  
          x_1 = x_1.view(batch, -1, x_1.size(1), x_1.size(2), x_1.size(3))
          x_1 = x_1.permute(0, 2, 1, 3, 4).contiguous()
          #x_1 = F.adaptive_max_pool3d(x_1, (1, None, None))
          x_1 = x_1.permute(0, 2, 1, 3, 4).contiguous()
          x_1 = x_1.view(-1, x_1.size(2), x_1.size(3), x_1.size(4)).contiguous()
  
          x_7 = F.relu(self.bn7_1(self.conv7_1(torch.cat((F.interpolate(x_6, scale_factor=(2, 2)), x_1), dim=1))))
          x_7 = F.relu(self.bn7_2(self.conv7_2(x_7)))
  
          x = x.view(batch, -1, x.size(1), x.size(2), x.size(3))
          x = x.permute(0, 2, 1, 3, 4).contiguous()
          #x = F.adaptive_max_pool3d(x, (1, None, None))
          x = x.permute(0, 2, 1, 3, 4).contiguous()
          x = x.view(-1, x.size(2), x.size(3), x.size(4)).contiguous()
  
          x_8 = F.relu(self.bn8_1(self.conv8_1(torch.cat((F.interpolate(x_7, scale_factor=(2, 2)), x), dim=1))))
          res_x = F.relu(self.bn8_2(self.conv8_2(x_8)))
  
          return res_x
  
  class lidar_decoder_kd(nn.Module):
      def __init__(self, height_feat_size=13):
          super(lidar_decoder_kd, self).__init__()
          self.conv_pre_1 = nn.Conv2d(height_feat_size, 32, kernel_size=3, stride=1, padding=1)
          self.conv_pre_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
          self.bn_pre_1 = nn.BatchNorm2d(32)
          self.bn_pre_2 = nn.BatchNorm2d(32)
  
  #        self.conv3d_1 = Conv3D(64, 64, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))
  #        self.conv3d_2 = Conv3D(128, 128, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))
          self.conv3d_1 = Conv3D(64, 64, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
          self.conv3d_2 = Conv3D(128, 128, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
  
          self.conv1_1 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
          self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
  
          self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
          self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
  
          self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
          self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
  
          self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)
          self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
  
          self.conv5_1 = nn.Conv2d(512 + 256, 256, kernel_size=3, stride=1, padding=1)
          self.conv5_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
  
          self.conv6_1 = nn.Conv2d(256 + 128, 128, kernel_size=3, stride=1, padding=1)
          self.conv6_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
  
          self.conv7_1 = nn.Conv2d(128 + 64, 64, kernel_size=3, stride=1, padding=1)
          self.conv7_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
  
          self.conv8_1 = nn.Conv2d(64 + 32, 32, kernel_size=3, stride=1, padding=1)
          self.conv8_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
  
          self.bn1_1 = nn.BatchNorm2d(64)
          self.bn1_2 = nn.BatchNorm2d(64)
  
          self.bn2_1 = nn.BatchNorm2d(128)
          self.bn2_2 = nn.BatchNorm2d(128)
  
          self.bn3_1 = nn.BatchNorm2d(256)
          self.bn3_2 = nn.BatchNorm2d(256)
  
          self.bn4_1 = nn.BatchNorm2d(512)
          self.bn4_2 = nn.BatchNorm2d(512)
  
          self.bn5_1 = nn.BatchNorm2d(256)
          self.bn5_2 = nn.BatchNorm2d(256)
  
          self.bn6_1 = nn.BatchNorm2d(128)
          self.bn6_2 = nn.BatchNorm2d(128)
  
          self.bn7_1 = nn.BatchNorm2d(64)
          self.bn7_2 = nn.BatchNorm2d(64)
  
          self.bn8_1 = nn.BatchNorm2d(32)
          self.bn8_2 = nn.BatchNorm2d(32)
  
      def forward(self, x, x_1, x_2, x_3, x_4, batch):
          # -------------------------------- Decoder Path --------------------------------
          x_5 = F.relu(self.bn5_1(self.conv5_1(torch.cat((F.interpolate(x_4, scale_factor=(2, 2)), x_3), dim=1))))
          x_5 = F.relu(self.bn5_2(self.conv5_2(x_5)))
  
          x_2 = x_2.view(batch, -1, x_2.size(1), x_2.size(2), x_2.size(3))
          x_2 = x_2.permute(0, 2, 1, 3, 4).contiguous()
          #x_2 = F.adaptive_max_pool3d(x_2, (1, None, None))
          x_2 = x_2.permute(0, 2, 1, 3, 4).contiguous()
          x_2 = x_2.view(-1, x_2.size(2), x_2.size(3), x_2.size(4)).contiguous()
  
          x_6 = F.relu(self.bn6_1(self.conv6_1(torch.cat((F.interpolate(x_5, scale_factor=(2, 2)), x_2), dim=1))))
          x_6 = F.relu(self.bn6_2(self.conv6_2(x_6)))
  
          x_1 = x_1.view(batch, -1, x_1.size(1), x_1.size(2), x_1.size(3))
          x_1 = x_1.permute(0, 2, 1, 3, 4).contiguous()
          #x_1 = F.adaptive_max_pool3d(x_1, (1, None, None))
          x_1 = x_1.permute(0, 2, 1, 3, 4).contiguous()
          x_1 = x_1.view(-1, x_1.size(2), x_1.size(3), x_1.size(4)).contiguous()
  
          x_7 = F.relu(self.bn7_1(self.conv7_1(torch.cat((F.interpolate(x_6, scale_factor=(2, 2)), x_1), dim=1))))
          x_7 = F.relu(self.bn7_2(self.conv7_2(x_7)))
  
          x = x.view(batch, -1, x.size(1), x.size(2), x.size(3))
          x = x.permute(0, 2, 1, 3, 4).contiguous()
          #x = F.adaptive_max_pool3d(x, (1, None, None))
          x = x.permute(0, 2, 1, 3, 4).contiguous()
          x = x.view(-1, x.size(2), x.size(3), x.size(4)).contiguous()
  
          x_8 = F.relu(self.bn8_1(self.conv8_1(torch.cat((F.interpolate(x_7, scale_factor=(2, 2)), x), dim=1))))
          res_x = F.relu(self.bn8_2(self.conv8_2(x_8)))
  
          return res_x, x_7, x_6, x_5
  
  
  
  class adafusionlayer(nn.Module):
      def __init__(self,input_channel=128):
          super(adafusionlayer, self).__init__()
          self.attn = nn.Conv2d(input_channel, 1, kernel_size=1, stride=1)
          self.bn = nn.BatchNorm2d(1)
  
      def forward(self,x):
  
          _,c,h,w = x.size()
          num_agent = x.size()[0]
  
          fusion_weight = F.relu(self.bn(self.attn(x)))
          # fusion_weight = F.relu(self.attn(x))
  
          fusion_weight = F.softmax(fusion_weight,dim=0).cuda()
  
  
          feat = torch.zeros(x[0].size()).cuda()
          for j in range(num_agent):
  
              feat = feat + (x[j]*(fusion_weight[j].repeat(c,1,1)))
          
          return feat
  
  class multifusionlayer(nn.Module):
      def __init__(self,input_channel=128):
          super(multifusionlayer,self).__init__()
          c1 = 64
          c2 = 32 
          self.attn1 = nn.Conv2d(input_channel,c1,kernel_size=1, stride=1)
          self.attn2 = nn.Conv2d(c1,c2,kernel_size=1, stride=1)
          self.attn3 = nn.Conv2d(c2,1,kernel_size=1, stride=1)
          
          self.bn_1=nn.BatchNorm2d(c1)
          self.bn_2=nn.BatchNorm2d(c2)
          self.bn_3=nn.BatchNorm2d(1)
  
      
      def forward(self,x):
  
          _,c,h,w = x.size()
          num_agent = 5
  
          x = F.relu(self.bn_1(self.attn1(x)))
          x = F.relu(self.bn_2(self.attn2(x)))
          x = F.relu(self.bn_3(self.attn3(x)))
          fusion_weight = F.softmax(x,dim=0).cuda()
  
          feat = torch.zeros(x[0].size()).cuda()
          for j in range(num_agent):
  
              feat = feat + (x[j]*(fusion_weight[j].repeat(c,1,1)))
          
          return feat
  
  
  
  class sigmoidfusionlayer(nn.Module):
      def __init__(self,input_channel=128):
          super(sigmoidfusionlayer, self).__init__()
          self.attn = nn.Conv2d(input_channel, 1, kernel_size=1, stride=1)
          self.bn = nn.BatchNorm2d(1)
  
      def forward(self,x):
  
          _,c,h,w = x.size()
          num_agent = 5
  
          fusion_weight = F.sigmoid(F.relu(self.bn(self.attn(x)))).cuda()
          # fusion_weight = fusion_weight.sum(dim=0).cuda()
  
  
          feat = torch.zeros(x[0].size()).cuda()
          for j in range(num_agent):
  
              feat = feat + (x[j]*(fusion_weight[j]))
          
          return feat
  
  
  # class MLPfusionlayer(nn.Module):
  #     def __init__(self,input_channel=128):
  #         super(MLPfusionlayer,self).__init__()
  
  #         self.MLP = 
  
  
  class pairfusionlayer(nn.Module):
      def __init__(self,input_channel=512):
          super(pairfusionlayer, self).__init__()
          self.attn = nn.Conv2d(512, 1, kernel_size=1, stride=1)
          self.bn = nn.BatchNorm2d(1)
  
      def forward(self,x):
  
          _,c,h,w = x.size()
          num_agent = x.size()[0]
          
          cat_list=[]
          for i in range(num_agent):
              
              cat_list.append(torch.cat((x[0],x[i])))
          
          
          feat_list=torch.stack(cat_list)
  
          # fusion_weight = F.relu(self.bn(self.attn(x)))
          fusion_weight = F.relu(self.attn(feat_list))
  
          fusion_weight = F.softmax(fusion_weight,dim=0).cuda()
  
          feat = torch.zeros(x[0].size()).cuda()
          for j in range(num_agent):
  
              feat = feat + (x[j]*(fusion_weight[j].repeat(c,1,1)))
          
          return feat
  
  
  class pairfusionlayer_1(nn.Module):
      def __init__(self,input_channel=512):
          super(pairfusionlayer_1, self).__init__()
          self.conv1_1 = nn.Conv2d(512, 128, kernel_size=1, stride=1, padding=0)
          self.bn1_1 = nn.BatchNorm2d(128)
  
          self.conv1_2 = nn.Conv2d(128, 32, kernel_size=1, stride=1, padding=0)
          self.bn1_2 = nn.BatchNorm2d(32)
  
          self.conv1_3 = nn.Conv2d(32, 8, kernel_size=1, stride=1, padding=0)
          self.bn1_3 = nn.BatchNorm2d(8)
  
          self.conv1_4 = nn.Conv2d(8, 1, kernel_size=1, stride=1, padding=0)
  
      def forward(self,x):
  
          _,c,h,w = x.size()
          num_agent = x.size()[0]
          
          cat_list=[]
          for i in range(num_agent):
              
              cat_list.append(torch.cat((x[0],x[i])))
          
          
          feat_list=torch.stack(cat_list)
  
          # # fusion_weight = F.relu(self.bn(self.attn(x)))
          # fusion_weight = F.relu(self.attn(feat_list))
          x_1 = F.relu(self.bn1_1(self.conv1_1(feat_list)))
          x_1 = F.relu(self.bn1_2(self.conv1_2(x_1)))
          x_1 = F.relu(self.bn1_3(self.conv1_3(x_1)))
          fusion_weight = F.relu(self.conv1_4(x_1))
  
  
          fusion_weight = F.softmax(fusion_weight,dim=0).cuda()
  
          feat = torch.zeros(x[0].size()).cuda()
          for j in range(num_agent):
  
              feat = feat + (x[j]*(fusion_weight[j].repeat(c,1,1)))
          
          return feat
  
  
  class pairfusionlayer_2(nn.Module):
      def __init__(self,input_channel=512):
          super(pairfusionlayer_2, self).__init__()
          self.conv1_1 = nn.Conv2d(512, 128, kernel_size=1, stride=1, padding=0)
          self.bn1_1 = nn.BatchNorm2d(128)
  
          self.conv1_2 = nn.Conv2d(128, 32, kernel_size=1, stride=1, padding=0)
          self.bn1_2 = nn.BatchNorm2d(32)
  
          self.conv1_3 = nn.Conv2d(32, 8, kernel_size=1, stride=1, padding=0)
          self.bn1_3 = nn.BatchNorm2d(8)
  
          self.conv1_4 = nn.Conv2d(8, 1, kernel_size=1, stride=1, padding=0)
          self.bn1_4 = nn.BatchNorm2d(1)
  
  
      def forward(self,x):
  
          _,c,h,w = x.size()
          num_agent = x.size()[0]
          
          cat_list=[]
          for i in range(num_agent):
              
              cat_list.append(torch.cat((x[0],x[i])))
          
          
          feat_list=torch.stack(cat_list)
  
          # # fusion_weight = F.relu(self.bn(self.attn(x)))
          # fusion_weight = F.relu(self.attn(feat_list))
          x_1 = F.relu(self.bn1_1(self.conv1_1(feat_list)))
          x_1 = F.relu(self.bn1_2(self.conv1_2(x_1)))
          x_1 = F.relu(self.bn1_3(self.conv1_3(x_1)))
          fusion_weight = F.relu(self.bn1_4(self.conv1_4(x_1)))
  
  
          fusion_weight = F.softmax(fusion_weight,dim=0).cuda()
  
          feat = torch.zeros(x[0].size()).cuda()
          for j in range(num_agent):
  
              feat = feat + (x[j]*(fusion_weight[j].repeat(c,1,1)))
          
          return feat
  
  
  class pairfusionlayer_3(nn.Module):
      def __init__(self,input_channel=512):
          super(pairfusionlayer_3, self).__init__()
          self.conv1_1 = nn.Conv2d(512, 128, kernel_size=1, stride=1, padding=0)
          self.bn1_1 = nn.BatchNorm2d(128)
  
          self.conv1_2 = nn.Conv2d(128, 8, kernel_size=1, stride=1, padding=0)
          self.bn1_2 = nn.BatchNorm2d(8)
  
          # self.conv1_3 = nn.Conv2d(32, 8, kernel_size=1, stride=1, padding=0)
          # self.bn1_3 = nn.BatchNorm2d(8)
  
          self.conv1_4 = nn.Conv2d(8, 1, kernel_size=1, stride=1, padding=0)
          self.bn1_4 = nn.BatchNorm2d(1)
  
  
      def forward(self,x):
  
          _,c,h,w = x.size()
          num_agent = x.size()[0]
          
          cat_list=[]
          for i in range(num_agent):
              
              cat_list.append(torch.cat((x[0],x[i])))
          
          
          feat_list=torch.stack(cat_list)
  
          # # fusion_weight = F.relu(self.bn(self.attn(x)))
          # fusion_weight = F.relu(self.attn(feat_list))
          x_1 = F.relu(self.bn1_1(self.conv1_1(feat_list)))
          x_1 = F.relu(self.bn1_2(self.conv1_2(x_1)))
          # x_1 = F.relu(self.bn1_3(self.conv1_3(x_1)))
          fusion_weight = F.relu(self.bn1_4(self.conv1_4(x_1)))
  
  
          fusion_weight = F.softmax(fusion_weight,dim=0).cuda()
  
          feat = torch.zeros(x[0].size()).cuda()
          for j in range(num_agent):
  
              feat = feat + (x[j]*(fusion_weight[j].repeat(c,1,1)))
          
          return feat
  
  
  class pairfusionlayer_4(nn.Module):
      def __init__(self,input_channel=512):
          super(pairfusionlayer_4, self).__init__()
          self.conv1_1 = nn.Conv2d(512, 256, kernel_size=1, stride=1, padding=0)
          self.bn1_1 = nn.BatchNorm2d(256)
  
          # self.conv1_2 = nn.Conv2d(128, 8, kernel_size=1, stride=1, padding=0)
          # self.bn1_2 = nn.BatchNorm2d(8)
  
          # self.conv1_3 = nn.Conv2d(32, 8, kernel_size=1, stride=1, padding=0)
          # self.bn1_3 = nn.BatchNorm2d(8)
  
          self.conv1_4 = nn.Conv2d(256, 1, kernel_size=1, stride=1, padding=0)
          self.bn1_4 = nn.BatchNorm2d(1)
  
  
      def forward(self,x):
  
          _,c,h,w = x.size()
          num_agent = x.size()[0]
          
          cat_list=[]
          for i in range(num_agent):
              
              cat_list.append(torch.cat((x[0],x[i])))
          
          
          feat_list=torch.stack(cat_list)
  
          # # fusion_weight = F.relu(self.bn(self.attn(x)))
          # fusion_weight = F.relu(self.attn(feat_list))
          x_1 = F.relu(self.bn1_1(self.conv1_1(feat_list)))
          # x_1 = F.relu(self.bn1_2(self.conv1_2(x_1)))
          # x_1 = F.relu(self.bn1_3(self.conv1_3(x_1)))
          fusion_weight = F.relu(self.bn1_4(self.conv1_4(x_1)))
  
          # print(fusion_weight.size())
  
  
          fusion_weight = F.softmax(fusion_weight,dim=0).cuda()
          # print(fusion_weight.size())
          # ipdb.set_trace()
  
  
          feat = torch.zeros(x[0].size()).cuda()
          for j in range(num_agent):
  
              feat = feat + (x[j]*(fusion_weight[j].repeat(c,1,1)))
          
          return feat
  
  
  
  
  