模型通过三个head得到响应输出，即Rf（RGB输出）；Re（event输出）；Rc（融合特征输出）。此时为了探索不同模态信息对跟踪任务的重要程度，设计了自适应模态响应决策机制，使得网络在学习的过程中判断利用RGB模态信息、Event模态信息还是融合模态信息，以提升特征信息利用的灵活性，最大化模型性能。那么基于"""
Basic ceutrack model.
"""
import math
import os
from typing import List

import torch
from torch import nn
from torch.nn.modules.transformer import _get_clones

from lib.models.layers.head import build_box_head
from lib.models.ceutrack.vit import vit_base_patch16_224
from lib.models.ceutrack.vit_ce import vit_large_patch16_224_ce, vit_base_patch16_224_ce
# from lib.models.ceutrack.vit_ce_BACKUPS import vit_large_patch16_224_ce, vit_base_patch16_224_ce
from lib.utils.box_ops import box_xyxy_to_cxcywh


class CEUTrack(nn.Module):
    """ This is the base class for ceutrack """

    def __init__(self, transformer, box_head, aux_loss=False, head_type="CORNER"):
        """ Initializes the model.
        Parameters:
            transformer: torch module of the transformer architecture.
            aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
        """
        super().__init__()
        self.backbone = transformer
        self.box_head_x = box_head
        # 新增event模态的head,以输出独立模态响应
        self.box_head_event_x = box_head

        self.aux_loss = aux_loss
        self.head_type = head_type
        if head_type == "CORNER" or head_type == "CENTER":
            self.feat_sz_s = int(box_head.feat_sz)
            self.feat_len_s = int(box_head.feat_sz ** 2)

        if self.aux_loss:
            self.box_head = _get_clones(self.box_head, 6)

    def forward(self, template: torch.Tensor,
                search: torch.Tensor,
                event_template: torch.Tensor,        # torch.Size([4, 1, 19, 10000])
                event_search: torch.Tensor,          # torch.Size([4, 1, 19, 10000])
                ce_template_mask=None,
                ce_keep_rate=None,
                return_last_attn=False,
                ):

        # before feeding into backbone, we need to concat four vectors, or two two concat;
        x, event_x, aux_dict = self.backbone(z=template, x=search, event_z=event_template, event_x=event_search,
                                    ce_template_mask=ce_template_mask,
                                    ce_keep_rate=ce_keep_rate,
                                    return_last_attn=return_last_attn, )

        # Forward head for x
        feat_last = x
        if isinstance(x, list):
            feat_last_x = x[-1]
        out_x = self.forward_head_x(feat_last_x, None)

        # Forward head for event_x
        feat_last_event_x = event_x
        if isinstance(event_x, list):
            feat_last_event_x = event_x[-1]
        out_event_x = self.forward_head_event_x(feat_last_event_x, None)

        # 新增 fusion of response
        fused_output = self.fuse_outputs(out_x,out_event_x)
        # update aux_dict 
        fused_output.update(fused_output)

        return fused_output

    # 单独构建forward_head
    def forward_head_x(self, cat_feature_x, gt_score_map=None):
            """
            Forward pass for the head of x modality.
            """
            # dual head
            enc_opt1_x = cat_feature_x[:, -self.feat_len_s:]
            enc_opt2_x = cat_feature_x[:, :self.feat_len_s]
            enc_opt_x = torch.cat([enc_opt1_x, enc_opt2_x], dim=-1)

            opt_x = (enc_opt_x.unsqueeze(-1)).permute((0, 3, 2, 1)).contiguous()
            bs, Nq, C, HW = opt_x.size()
            opt_feat_x = opt_x.view(-1, C, self.feat_sz_s, self.feat_sz_s)

            # Forward pass for x modality
            if self.head_type == "CORNER":
                pred_box_x, score_map_x = self.box_head_x(opt_feat_x, True)
                outputs_coord_x = box_xyxy_to_cxcywh(pred_box_x)
                outputs_coord_new_x = outputs_coord_x.view(bs, Nq, 4)

                out_x = {'pred_boxes_x': outputs_coord_new_x,
                        'score_map_x': score_map_x}
                return out_x

            elif self.head_type == "CENTER":
                score_map_ctr_x, bbox_x, size_map_x, offset_map_x = self.box_head_x(opt_feat_x, gt_score_map)
                outputs_coord_x = bbox_x
                outputs_coord_new_x = outputs_coord_x.view(bs, Nq, 4)

                out_x = {'pred_boxes_x': outputs_coord_new_x,
                        'score_map_x': score_map_ctr_x,
                        'size_map_x': size_map_x,
                        'offset_map_x': offset_map_x}
                return out_x
            else:
                raise NotImplementedError
    # Event模态响应forward head
    def forward_head_event_x(self, cat_feature_event_x, gt_score_map=None):
        """
        Forward pass for the head of event_x modality.
        """
        # dual head
        enc_opt1_event_x = cat_feature_event_x[:, -self.feat_len_s:]
        enc_opt2_event_x = cat_feature_event_x[:, :self.feat_len_s]
        enc_opt_event_x = torch.cat([enc_opt1_event_x, enc_opt2_event_x], dim=-1)

        opt_event_x = (enc_opt_event_x.unsqueeze(-1)).permute((0, 3, 2, 1)).contiguous()
        bs, Nq, C, HW = opt_event_x.size()
        opt_feat_event_x = opt_event_x.view(-1, C, self.feat_sz_s, self.feat_sz_s)

        # Forward pass for event_x modality
        if self.head_type == "CORNER":
            pred_box_event_x, score_map_event_x = self.box_head_event_x(opt_feat_event_x, True)
            outputs_coord_event_x = box_xyxy_to_cxcywh(pred_box_event_x)
            outputs_coord_new_event_x = outputs_coord_event_x.view(bs, Nq, 4)

            out_event_x = {'pred_boxes_event_x': outputs_coord_new_event_x,
                           'score_map_event_x': score_map_event_x}
            return out_event_x

        elif self.head_type == "CENTER":
            score_map_ctr_event_x, bbox_event_x, size_map_event_x, offset_map_event_x = self.box_head_event_x(
                opt_feat_event_x, gt_score_map)
            outputs_coord_event_x = bbox_event_x
            outputs_coord_new_event_x = outputs_coord_event_x.view(bs, Nq, 4)

            out_event_x = {'pred_boxes_event_x': outputs_coord_new_event_x,
                           'score_map_event_x': score_map_ctr_event_x,
                           'size_map_event_x': size_map_event_x,
                           'offset_map_event_x': offset_map_event_x}
            return out_event_x
        else:
            raise NotImplementedError
        
    def fuse_outputs(self, out_x, out_event_x):
        fused_pred_boxes = out_x['pred_boxes'] + out_event_x['pred_boxes']
        fused_score_map = out_x['score_map'] + out_event_x['score_map']
        fused_output = {
            'pred_boxes': fused_pred_boxes,
            'score_map': fused_score_map,
        }
        return fused_output

def build_ceutrack(cfg, training=True):
    current_dir = os.path.dirname(os.path.abspath(__file__))  # This is your Project Root
    pretrained_path = os.path.join(current_dir, 'pretrained_models')
    if cfg.MODEL.PRETRAIN_FILE and ('CEUTrack' not in cfg.MODEL.PRETRAIN_FILE) and training:
        pretrained = os.path.join(pretrained_path, cfg.MODEL.PRETRAIN_FILE)
    else:
        pretrained = ''

    if cfg.MODEL.BACKBONE.TYPE == 'vit_base_patch16_224':
        backbone = vit_base_patch16_224(pretrained, drop_path_rate=cfg.TRAIN.DROP_PATH_RATE)
        hidden_dim = backbone.embed_dim
        patch_start_index = 1

    elif cfg.MODEL.BACKBONE.TYPE == 'vit_base_patch16_224_ce':
        backbone = vit_base_patch16_224_ce(pretrained, drop_path_rate=cfg.TRAIN.DROP_PATH_RATE,
                                           ce_loc=cfg.MODEL.BACKBONE.CE_LOC,
                                           ce_keep_ratio=cfg.MODEL.BACKBONE.CE_KEEP_RATIO,
                                           )
        # hidden_dim = backbone.embed_dim
        hidden_dim = backbone.embed_dim * 2
        patch_start_index = 1

    elif cfg.MODEL.BACKBONE.TYPE == 'vit_large_patch16_224_ce':
        backbone = vit_large_patch16_224_ce(pretrained, drop_path_rate=cfg.TRAIN.DROP_PATH_RATE,
                                            ce_loc=cfg.MODEL.BACKBONE.CE_LOC,
                                            ce_keep_ratio=cfg.MODEL.BACKBONE.CE_KEEP_RATIO,
                                            )

        hidden_dim = backbone.embed_dim
        patch_start_index = 1

    else:
        raise NotImplementedError

    backbone.finetune_track(cfg=cfg, patch_start_index=patch_start_index)

    box_head_x = build_box_head(cfg, hidden_dim)
    box_head_event_x = build_box_head(cfg, hidden_dim)

    model = CEUTrack(
        backbone,
        box_head_x,
        box_head_event_x,
        aux_loss=False,
        head_type=cfg.MODEL.HEAD.TYPE,
    )

    if 'CEUTrack' in cfg.MODEL.PRETRAIN_FILE and training:
        checkpoint = torch.load(cfg.MODEL.PRETRAIN_FILE, map_location="cpu")
        missing_keys, unexpected_keys = model.load_state_dict(checkpoint["net"], strict=False)
        print('Load pretrained model from: ' + cfg.MODEL.PRETRAIN_FILE)

    return model
按照需求，写一段代码