text: from models.siampose_RCNN_corr import SiamMask
  from models.features import MultiStageFeature
  from models.rpn import RPN, DepthCorr
  from models.mask import Mask
  # from models.DCNv2.dcn_v2 import DCN
  import torch
  import torch.nn as nn
  import math
  from utils.load_helper import load_pretrain
  from resnet import resnet50
  
  BN_MOMENTUM = 0.1
  
  
  class ResDownS(nn.Module):
      def __init__(self, inplane, outplane):
          super(ResDownS, self).__init__()
          self.downsample = nn.Sequential(
                  nn.Conv2d(inplane, outplane, kernel_size=1, bias=False),
                  nn.BatchNorm2d(outplane))
  
      def forward(self, x):
          x = self.downsample(x)
          if x.size(3) < 20:
              l = 4
              r = -4
              x = x[:, :, l:r, l:r]
          return x
  
  
  class ResDown(MultiStageFeature):
      def __init__(self, pretrain=False):
          super(ResDown, self).__init__()
          self.features = resnet50(layer3=True, layer4=True)
          if pretrain:
              load_pretrain(self.features, '../resnet.model')
  
          self.downsample = ResDownS(1024, 256)
          self.downsample_p4 = ResDownS(2048, 1024)
  
          self.layers = [self.downsample, self.features.layer2, self.features.layer3, self.downsample_p4, self.features.layer4]
          self.train_nums = [3, 3]
          self.change_point = [0, 0.5]
  
          self.unfix(0.0)
  
      def param_groups(self, start_lr, feature_mult=1):
          lr = start_lr * feature_mult
  
          def _params(module, mult=1):
              params = list(filter(lambda x:x.requires_grad, module.parameters()))
              if len(params):
                  return [{'params': params, 'lr': lr * mult}]
              else:
                  return []
  
          groups = []
          groups += _params(self.downsample)
          groups += _params(self.features, 0.1)
          return groups
  
      def forward(self, x):
          output = self.features(x)
          p3 = self.downsample(output[-2])
          return p3
  
      def forward_all(self, x):
          output = self.features(x)
          p3 = self.downsample(output[-2])
          p4 = self.downsample_p4(output[-1])
          return p3, p4
  
  
  class UP(RPN):
      def __init__(self, anchor_num=5, feature_in=256, feature_out=256):
          super(UP, self).__init__()
  
          self.anchor_num = anchor_num
          self.feature_in = feature_in
          self.feature_out = feature_out
  
          self.cls_output = 2 * self.anchor_num
          self.loc_output = 4 * self.anchor_num
  
          self.cls = DepthCorr(feature_in, feature_out, self.cls_output)
          self.loc = DepthCorr(feature_in, feature_out, self.loc_output)
  
      def forward(self, z_f, x_f):
          cls = self.cls(z_f, x_f)
          loc = self.loc(z_f, x_f)
          return cls, loc
  
  
  class Center_pose_head(nn.Module):
      def __init__(self, head_conv=256):
          super(Center_pose_head, self).__init__()
  
          self.inplanes = 512
          self.deconv_with_bias = False
          # self.deconv_layers = self._make_deconv_layer(
          #     3,
          #     [256, 128, 64],
          #     [4, 4, 4],
          # )
          self.deconv_layers_new = self._make_deconv_layer(
              3,
              [512, 256, 64],
              [4, 4, 4],
          )
          self.heads = {'hps': 34, 'hm_hp': 17, 'hp_offset': 2}
  
          for head in self.heads:
              classes = self.heads[head]
              if head_conv > 0:
                fc = nn.Sequential(
                    nn.Conv2d(self.inplanes, head_conv,
                      kernel_size=3, padding=1, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Conv2d(head_conv, classes,
                      kernel_size=1, stride=1,
                      padding=0, bias=True))
                if 'hm' in head:
                  fc[-1].bias.data.fill_(-2.19)
                else:
                  fill_fc_weights(fc)
              else:
                fc = nn.Conv2d(self.inplanes, classes,
                    kernel_size=1, stride=1,
                    padding=0, bias=True)
                if 'hm' in head:
                  fc.bias.data.fill_(-2.19)
                else:
                  fill_fc_weights(fc)
              self.__setattr__(head, fc)
  
      def _get_deconv_cfg(self, deconv_kernel, index):
          if deconv_kernel == 4:
              padding = 1
              output_padding = 0
          elif deconv_kernel == 3:
              padding = 1
              output_padding = 1
          elif deconv_kernel == 2:
              padding = 0
              output_padding = 0
  
          return deconv_kernel, padding, output_padding
  
      def _make_deconv_layer(self, num_layers, num_filters, num_kernels):
          assert num_layers == len(num_filters), \
              'ERROR: num_deconv_layers is different len(num_deconv_filters)'
          assert num_layers == len(num_kernels), \
              'ERROR: num_deconv_layers is different len(num_deconv_filters)'
  
          layers = []
          for i in range(num_layers):
              kernel, padding, output_padding = \
                  self._get_deconv_cfg(num_kernels[i], i)
  
              planes = num_filters[i]
              # fc = DCN(self.inplanes, planes,
              #         kernel_size=(3,3), stride=1,
              #         padding=1, dilation=1, deformable_groups=1)
              fc = nn.Conv2d(self.inplanes, planes,
                      kernel_size=3, stride=1,
                      padding=1, dilation=1, bias=False)
              fill_fc_weights(fc)
              up = nn.ConvTranspose2d(
                      in_channels=planes,
                      out_channels=planes,
                      kernel_size=kernel,
                      stride=2,
                      padding=padding,
                      output_padding=output_padding,
                      bias=self.deconv_with_bias)
              fill_up_weights(up)
  
              layers.append(fc)
              layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM))
              layers.append(nn.ReLU(inplace=True))
              layers.append(up)
              layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM))
              layers.append(nn.ReLU(inplace=True))
              self.inplanes = planes
  
          return nn.Sequential(*layers)
  
      def forward(self, x):
          x = self.deconv_layers_new(x)
          ret = {}
          for head in self.heads:
              ret[head] = self.__getattr__(head)(x)
          return [ret]
  
      def param_groups(self, start_lr, feature_mult=1, key=None):
          if key is None:
              params = filter(lambda x:x.requires_grad, self.parameters())
          else:
              params = [v for k, v in self.named_parameters() if (key in k) and v.requires_grad]
          params = [{'params': params, 'lr': start_lr * feature_mult}]
          return params
  
  
  class PoseCorr(Mask):
      def __init__(self):
          super(PoseCorr, self).__init__()
          self.kp = DepthCorr(256, 512, 512)
  
      def forward(self, z, x):
          return self.kp(z, x)
  
  
  class Custom(SiamMask):
      def __init__(self, pretrain=False, **kwargs):
          super(Custom, self).__init__(**kwargs)
          # self.opt = opt
          self.features = ResDown(pretrain=pretrain)
          self.rpn_model = UP(anchor_num=self.anchor_num, feature_in=256, feature_out=256)
          self.pose_corr = PoseCorr()
          self.kp_model = Center_pose_head()
  
      def template(self, template):
          self.zf = self.features(template)
  
      def track(self, search):
          search = self.features(search)
          rpn_pred_cls, rpn_pred_loc = self.rpn(self.zf, search)
          return rpn_pred_cls, rpn_pred_loc
  
      def track_mask(self, search):
          search = self.features(search)
          rpn_pred_cls, rpn_pred_loc = self.rpn(self.zf, search)
          pred_mask = self.mask(self.zf, search)
          return rpn_pred_cls, rpn_pred_loc, pred_mask
  
  
  def fill_fc_weights(layers):
      for m in layers.modules():
          if isinstance(m, nn.Conv2d):
              if m.bias is not None:
                  nn.init.constant_(m.bias, 0)
  
  def fill_up_weights(up):
      w = up.weight.data
      f = math.ceil(w.size(2) / 2)
      c = (2 * f - 1 - f % 2) / (2. * f)
      for i in range(w.size(2)):
          for j in range(w.size(3)):
              w[0, 0, i, j] = \
                  (1 - math.fabs(i / f - c)) * (1 - math.fabs(j / f - c))
      for c in range(1, w.size(0)):
          w[c, 0, :, :] = w[0, 0, :, :]
  