text: # Copyright (c) SenseTime. All Rights Reserved.
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from __future__ import unicode_literals
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import numpy as np
  
  from pysot.core.config import cfg
  from pysot.models.loss import select_cross_entropy_loss, weight_l1_loss
  from pysot.models.backbone import get_backbone
  from pysot.models.head import get_rpn_head, get_mask_head, get_refine_head
  from pysot.models.temporal import get_temporal_head # Changed
  from pysot.models.neck import get_neck
  
  
  class ModelBuilder(nn.Module):
      def __init__(self):
          super(ModelBuilder, self).__init__()
  
          # build backbone
          self.backbone = get_backbone(cfg.BACKBONE.TYPE,
                                       **cfg.BACKBONE.KWARGS)
  
          # build adjust layer
          if cfg.ADJUST.ADJUST:
              self.neck = get_neck(cfg.ADJUST.TYPE,
                                   **cfg.ADJUST.KWARGS) # AdjustAllLayer
  
          # build rpn head
          self.rpn_head = get_rpn_head(cfg.RPN.TYPE,
                                       **cfg.RPN.KWARGS) # MultiRPN
  
          self.temporal_head = get_temporal_head(cfg.TEMPORAL.TYPE) # Temporal Network
  
          # build mask head
          if cfg.MASK.MASK:
              self.mask_head = get_mask_head(cfg.MASK.TYPE,
                                             **cfg.MASK.KWARGS)
  
              if cfg.REFINE.REFINE:
                  self.refine_head = get_refine_head(cfg.REFINE.TYPE)
  
      def template(self, z): # template
          zf = self.backbone(z)
          if cfg.MASK.MASK:
              zf = zf[-1]
          if cfg.ADJUST.ADJUST:
              zf = self.neck(zf)
          self.zf = zf
  
      def track(self, x): # Seach Image
          xf = self.backbone(x)
          if cfg.MASK.MASK: # False
              self.xf = xf[:-1]
              xf = xf[-1]
          if cfg.ADJUST.ADJUST: # True and TYPE: "AdjustAllLayer"
              xf = self.neck(xf)
          cls, loc = self.rpn_head(self.zf, xf)
          if cfg.MASK.MASK:
              mask, self.mask_corr_feature = self.mask_head(self.zf, xf)
          return {
                  'cls': cls,
                  'loc': loc,
                  'mask': mask if cfg.MASK.MASK else None
                 }
  
      def mask_refine(self, pos):
          return self.refine_head(self.xf, self.mask_corr_feature, pos)
  
      def log_softmax(self, cls):
          b, a2, h, w = cls.size()
          cls = cls.view(b, 2, a2//2, h, w)
          cls = cls.permute(0, 2, 3, 4, 1).contiguous()
          cls = F.log_softmax(cls, dim=4)
          return cls
  
      def forward(self, data):
          """ only used in training
          """
          template = data['template'].cuda()
          search = data['search'].cuda()
          #data['temporal'] = np.asarray(data['temporal'],dtype=np.float32)
          temporal_data = data['temporal'].cuda()
          #print(temporal.shape())# Temporral Image data : shape [1, 3, 255, 255, 50]
          label_cls = data['label_cls'].cuda()
          print(label_cls.shape)
          label_loc = data['label_loc'].cuda()
          label_loc_weight = data['label_loc_weight'].cuda()
  
          # get feature
          zf = self.backbone(template)
          xf = self.backbone(search)
          temp_f = self.temporal_head(temporal_data) # temp_f = [out2, out3, out4]
  
          # Converting 3D Tensor into 2D Tensor
          for i in range(len(temp_f)):
              b, c, h, w, _ = temp_f[i].shape
              temp_f[i] = torch.reshape(temp_f[i], (b, c, h, w))
  
          # Adding Temporal Features into ResNet Used Layers
  #        print(xf[0].shape)
  #        print(temp_f[0].shape)
          for i in range(len(xf)):
              xf[i] = xf[i].add(temp_f[i])
          print(xf[0].shape)
  
          if cfg.MASK.MASK:
              zf = zf[-1]
              self.xf_refine = xf[:-1]
              xf = xf[-1]
          if cfg.ADJUST.ADJUST:
              zf = self.neck(zf)
              xf = self.neck(xf)
          cls, loc = self.rpn_head(zf, xf)
  
          # get loss
          cls = self.log_softmax(cls)
          print(cls.shape)
          cls_loss = select_cross_entropy_loss(cls, label_cls)
          loc_loss = weight_l1_loss(loc, label_loc, label_loc_weight)
  
          outputs = {}
          outputs['total_loss'] = cfg.TRAIN.CLS_WEIGHT * cls_loss + \
              cfg.TRAIN.LOC_WEIGHT * loc_loss
          outputs['cls_loss'] = cls_loss
          outputs['loc_loss'] = loc_loss
  
          if cfg.MASK.MASK:
              # TODO
              mask, self.mask_corr_feature = self.mask_head(zf, xf)
              mask_loss = None
              outputs['total_loss'] += cfg.TRAIN.MASK_WEIGHT * mask_loss
              outputs['mask_loss'] = mask_loss
             
          return outputs
  