模型是：import math
import logging
from functools import partial
from collections import OrderedDict
from copy import deepcopy

import torch
import torch.nn as nn
import torch.nn.functional as F

from timm.models.layers import to_2tuple

from lib.models.layers.patch_embed import PatchEmbed, PatchEmbed_event, xcorr_depthwise
from .utils import combine_tokens, recover_tokens
from .vit import VisionTransformer
from ..layers.attn_blocks import CEBlock
import random
import numpy as np

_logger = logging.getLogger(__name__)


class VisionTransformerCE(VisionTransformer):
    """ Vision Transformer with candidate elimination (CE) module

    A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`
        - https://arxiv.org/abs/2010.11929

    Includes distillation token & head support for `DeiT: Data-efficient Image Transformers`
        - https://arxiv.org/abs/2012.12877
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
                 num_heads=12, mlp_ratio=4., qkv_bias=True, representation_size=None, distilled=False,
                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0., embed_layer=PatchEmbed, norm_layer=None,
                 act_layer=None, weight_init='',
                 ce_loc=None, ce_keep_ratio=None):
        """
        Args:
            img_size (int, tuple): input image size
            patch_size (int, tuple): patch size
            in_chans (int): number of input channels
            num_classes (int): number of classes for classification head
            embed_dim (int): embedding dimension
            depth (int): depth of transformer
            num_heads (int): number of attention heads
            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
            qkv_bias (bool): enable bias for qkv if True
            representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
            distilled (bool): model includes a distillation token and head as in DeiT models
            drop_rate (float): dropout rate
            attn_drop_rate (float): attention dropout rate
            drop_path_rate (float): stochastic depth rate
            embed_layer (nn.Module): patch embedding layer
            norm_layer: (nn.Module): normalization layer
            weight_init: (str): weight init scheme
        """
        # super().__init__()
        super().__init__()
        if isinstance(img_size, tuple):
            self.img_size = img_size
        else:
            self.img_size = to_2tuple(img_size)
        self.patch_size = patch_size
        self.in_chans = in_chans

        self.num_classes = num_classes
        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
        self.num_tokens = 2 if distilled else 1
        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
        act_layer = act_layer or nn.GELU

        self.patch_embed = embed_layer(
            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
        self.pos_drop = nn.Dropout(p=drop_rate)
        self.pos_embed_event = PatchEmbed_event(in_chans=32, embed_dim=768, kernel_size=4, stride=4)
        # self.pos_embed_event = PatchEmbed_event(in_chans=32, embed_dim=768, kernel_size=4, stride=4)
        # self.pos_embed_event_z = PatchEmbed_event(in_chans=32, embed_dim=768, kernel_size=3, stride=1)
        # attn = CrossAttn(768, 4, 3072, 0.1, 'relu')
        # self.cross_attn = Iter_attn(attn, 2)

        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
        blocks = []
        ce_index = 0
        self.ce_loc = ce_loc
        for i in range(depth):
            ce_keep_ratio_i = 1.0
            if ce_loc is not None and i in ce_loc:
                ce_keep_ratio_i = ce_keep_ratio[ce_index]
                ce_index += 1

            blocks.append(
                CEBlock(
                    dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, drop=drop_rate,
                    attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, act_layer=act_layer,
                    keep_ratio_search=ce_keep_ratio_i)
            )

        self.blocks = nn.Sequential(*blocks)
        self.norm = norm_layer(embed_dim)

        self.init_weights(weight_init)
    def masking_fea(self,z, event_z, x, event_x, ratio=0.8 ):
        b,nz,c = z.shape
        b,nez,c = event_z.shape
        b,nx,c = x.shape
        b,nex,c = event_x.shape
        assert(nz == nez)
        assert(nx == nex)
        lenz_out = int(nz*ratio)
        lenx_out = int(nx*ratio)

        mask_nz = torch.rand(b,nz).float()
        mask_ez = torch.rand(b,nez).float()
        mask_nx = torch.rand(b,nx).float()
        mask_ex = torch.rand(b,nex).float()
        mask_nz = mask_nz>0.4
        mask_ez = mask_ez>0.4
        mask_ez = ~mask_nz + mask_ez
        mask_nz_idx = mask_nz.float().sort(1,descending=True)[-1].to(device = z.device)
        mask_ez_idx = mask_ez.float().sort(1,descending=True)[-1].to(device = z.device)
        mask_nx = mask_nx>0.4
        mask_ex = mask_ex>0.4
        mask_ex = ~mask_nx + mask_ex
        mask_nx_idx = mask_nx.float().sort(1,descending=True)[-1].to(device = z.device)
        mask_ex_idx = mask_ex.float().sort(1,descending=True)[-1].to(device = z.device)

        masked_z = torch.gather(z, 1, mask_nz_idx[:,:lenz_out,None].repeat([1,1,c]))
        masked_ez = torch.gather(event_z, 1, mask_ez_idx[:,:lenz_out,None].repeat([1,1,c]))
        masked_x = torch.gather(x, 1, mask_nx_idx[:,:lenx_out,None].repeat([1,1,c]))
        masked_ex = torch.gather(event_x, 1, mask_ex_idx[:,:lenx_out,None].repeat([1,1,c]))
        return masked_z, masked_ez, masked_x, masked_ex,{'x1':mask_nx_idx[:,:lenx_out],'x0':mask_nx_idx[:,lenx_out:],
                                                        'ex1':mask_ex_idx[:,:lenx_out],'ex0':mask_ex_idx[:,lenx_out:], }

    def forward_features(self, z, x, event_z, event_x,
                         mask_z=None, mask_x=None,
                         ce_template_mask=None, ce_keep_rate=None,
                         return_last_attn=False,Track=False
                         ):
        B, H, W = x.shape[0], x.shape[2], x.shape[3]
        # print('shape of event_z before projection:{}, event_x:{}'.format(event_z.shape, event_x.shape))
        event_z = self.pos_embed_event(event_z)     # [:,:,:,:1000]
        event_x = self.pos_embed_event(event_x)     # B 768 1024
        x = self.patch_embed(x)
        z = self.patch_embed(z)
        # print('shape of event_z:{}, event_x:{}, x:{}, z:{}'.format(event_z.shape,event_x.shape,x.shape,z.shape ))
        event_z += self.pos_embed_z
        event_x += self.pos_embed_x
        z += self.pos_embed_z
        x += self.pos_embed_x

        # attention mask handling   # B, H, W
        if mask_z is not None and mask_x is not None:
            mask_z = F.interpolate(mask_z[None].float(), scale_factor=1. / self.patch_size).to(torch.bool)[0]
            mask_z = mask_z.flatten(1).unsqueeze(-1)

            mask_x = F.interpolate(mask_x[None].float(), scale_factor=1. / self.patch_size).to(torch.bool)[0]
            mask_x = mask_x.flatten(1).unsqueeze(-1)

            mask_x = combine_tokens(mask_z, mask_x, mode=self.cat_mode)
            mask_x = mask_x.squeeze(-1)

        if self.add_cls_token:
            cls_tokens = self.cls_token.expand(B, -1, -1)
            cls_tokens = cls_tokens + self.cls_pos_embed

        if self.add_sep_seg:
            x += self.search_segment_pos_embed
            z += self.template_segment_pos_embed
        if Track == False:
            z, event_z, x, event_x, token_idx = self.masking_fea(z, event_z, x, event_x, ratio=0.9) 
        x = combine_tokens(z, event_z, x, event_x, mode=self.cat_mode)        # 64+64+256+256=640
        # x = combine_tokens(z, x, event_z, event_x, mode=self.cat_mode)        # 64+64+256+256=640
        if self.add_cls_token:
            x = torch.cat([cls_tokens, x], dim=1)

        x = self.pos_drop(x)
        # lens_z = self.pos_embed_z.shape[1]
        # lens_x = self.pos_embed_x.shape[1]
        lens_z = z.shape[1]
        lens_x = x.shape[1]
        global_index_t = torch.linspace(0, lens_z - 1, lens_z).to(x.device)
        global_index_t = global_index_t.repeat(B, 1)

        global_index_s = torch.linspace(0, lens_x - 1, lens_x).to(x.device)
        global_index_s = global_index_s.repeat(B, 1)
        removed_indexes_s = []
        out_attn = []
        out_attn5 = []
        fifth_layer_attn = None
        for i, blk in enumerate(self.blocks):
            # out_global_s.append(global_index_s)
            # out_global_t.append(global_index_t)
            x, global_index_t, global_index_s, removed_index_s, attn = \
                blk(x, global_index_t, global_index_s, mask_x, ce_template_mask, ce_keep_rate)
            # if i == 4:
            #     fifth_layer_attn = attn
            
            if self.ce_loc is not None and i in self.ce_loc:
                removed_indexes_s.append(removed_index_s)
            out_attn.append(attn)
            out_attn5.append(fifth_layer_attn)
                

        
        out_attn_idx = random.choice(np.arange(len(out_attn)))
        out_attn = out_attn[out_attn_idx]
        x = self.norm(x)
        lens_x_new = global_index_s.shape[1]
        lens_z_new = global_index_t.shape[1]

        z = x[:, :lens_z_new*2]
        x = x[:, lens_z_new*2:]
        if removed_indexes_s and removed_indexes_s[0] is not None:
            removed_indexes_cat = torch.cat(removed_indexes_s, dim=1)

            pruned_lens_x = lens_x - lens_x_new
            pad_x = torch.zeros([B, pruned_lens_x, x.shape[2]], device=x.device)
            x = torch.cat([x, pad_x], dim=1)
            index_all = torch.cat([global_index_s, removed_indexes_cat], dim=1)
            # recover original token order
            C = x.shape[-1]
            x = torch.zeros_like(x).scatter_(dim=1, index=index_all.unsqueeze(-1).expand(B, -1, C).to(torch.int64), src=x)

        x = recover_tokens(x, lens_z_new, lens_x, mode=self.cat_mode)
        # event_x = x[:, lens_x:]   # RGB head
        x = x[:, :lens_x]   # RGB head
        x = torch.cat([event_x, x], dim=1)
        
        aux_dict = {
            # "attn": attn,
            "attn": out_attn,
            # "fifth_layer_attn": out_attn5,
            "removed_indexes_s": removed_indexes_s,  # used for visualization
        }

        return x, aux_dict

    def forward(self, z, x, event_z, event_x,
                ce_template_mask=None, ce_keep_rate=None,
                tnc_keep_rate=None,
                return_last_attn=False,Track=False):

        x, aux_dict = self.forward_features(z, x, event_z, event_x, ce_template_mask=ce_template_mask, ce_keep_rate=ce_keep_rate,Track=Track)

        return x, aux_dict


def _create_vision_transformer(pretrained=False, **kwargs):
    model = VisionTransformerCE(**kwargs)

    if pretrained:
        if 'npz' in pretrained:
            model.load_pretrained(pretrained, prefix='')
        else:
            checkpoint = torch.load(pretrained, map_location="cpu")
            missing_keys, unexpected_keys = model.load_state_dict(checkpoint["model"], strict=False)
            print('Load pretrained model from: ' + pretrained)

    return model


def vit_base_patch16_224_ce(pretrained=False, **kwargs):
    """ ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
    """
    model_kwargs = dict(
        patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)
    model = _create_vision_transformer(pretrained=pretrained, **model_kwargs)
    return model


def vit_large_patch16_224_ce(pretrained=False, **kwargs):
    """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
    """
    model_kwargs = dict(
        patch_size=16, embed_dim=1024, depth=24, num_heads=16, **kwargs)
    model = _create_vision_transformer(pretrained=pretrained, **model_kwargs)
    return model 损失是：from . import BaseActor
from lib.utils.misc import NestedTensor
from lib.utils.box_ops import box_cxcywh_to_xyxy, box_xywh_to_xyxy
import torch
from lib.utils.merge import merge_template_search
from ...utils.heapmap_utils import generate_heatmap
from ...utils.ce_utils import generate_mask_cond, adjust_keep_rate
import matplotlib.pyplot as plt
import numpy as np
import torch.nn.functional as F
from skimage.measure import label, regionprops
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
class CEUTrackActor(BaseActor):
    """ Actor for training CEUTrack models """

    def __init__(self, net, objective, loss_weight, settings, cfg=None):
        super().__init__(net, objective)
        self.loss_weight = loss_weight
        self.settings = settings
        self.bs = self.settings.batchsize  # batch size
        self.cfg = cfg

    def __call__(self, data):
        """
        args:
            data - The input data, should contain the fields 'template', 'search', 'gt_bbox'.
            template_images: (N_t, batch, 3, H, W)
            search_images: (N_s, batch, 3, H, W)
        returns:
            loss    - the training loss
            status  -  dict containing detailed losses
        """
        # forward pass
        out_dict = self.forward_pass(data)

        # compute losses
        loss, status = self.compute_losses(out_dict, data)

        return loss, status

    def forward_pass(self, data):
        # currently only support 1 template and 1 search region
        assert len(data['template_images']) == 1
        assert len(data['search_images']) == 1
        assert len(data['template_event']) == 1
        assert len(data['search_event']) == 1

        template_list = []
        for i in range(self.settings.num_template):
            template_img_i = data['template_images'][i].view(-1,
                                                             *data['template_images'].shape[2:])  # (batch, 3, 128, 128)
            # template_att_i = data['template_att'][i].view(-1, *data['template_att'].shape[2:])  # (batch, 128, 128)
            template_list.append(template_img_i)

        search_img = data['search_images'][0].view(-1, *data['search_images'].shape[2:])  # (batch, 3, 320, 320)
        # search_att = data['search_att'][0].view(-1, *data['search_att'].shape[2:])  # (batch, 320, 320)

        template_event = data['template_event'][0].view(-1, *data['template_event'].shape[2:])
        search_event = data['search_event'][0].view(-1, *data['search_event'].shape[2:])

        box_mask_z = None
        ce_keep_rate = None
        if self.cfg.MODEL.BACKBONE.CE_LOC:
            box_mask_z = generate_mask_cond(self.cfg, template_list[0].shape[0], template_list[0].device,
                                            data['template_anno'][0])

            ce_start_epoch = self.cfg.TRAIN.CE_START_EPOCH
            ce_warm_epoch = self.cfg.TRAIN.CE_WARM_EPOCH
            ce_keep_rate = adjust_keep_rate(data['epoch'], warmup_epochs=ce_start_epoch,
                                                total_epochs=ce_start_epoch + ce_warm_epoch,
                                                ITERS_PER_EPOCH=1,
                                                base_keep_rate=self.cfg.MODEL.BACKBONE.CE_KEEP_RATIO[0])

        if len(template_list) == 1:
            template_list = template_list[0]

        out_dict = self.net(template=template_list,
                            search=search_img,
                            event_template=template_event,
                            event_search=search_event,
                            ce_template_mask=box_mask_z,
                            ce_keep_rate=ce_keep_rate,
                            return_last_attn=False)

        return out_dict

    def compute_losses(self, pred_dict, gt_dict, return_status=True):
        # gt gaussian map
        gt_bbox = gt_dict['search_anno'][-1]  # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4)
        gt_gaussian_maps = generate_heatmap(gt_dict['search_anno'], self.cfg.DATA.SEARCH.SIZE, self.cfg.MODEL.BACKBONE.STRIDE)
        gt_gaussian_maps = gt_gaussian_maps[-1].unsqueeze(1)

        # Get boxes
        pred_boxes = pred_dict['pred_boxes']
        if torch.isnan(pred_boxes).any():
            raise ValueError("Network outputs is NAN! Stop Training")
        num_queries = pred_boxes.size(1)
        pred_boxes_vec = box_cxcywh_to_xyxy(pred_boxes).view(-1, 4)  # (B,N,4) --> (BN,4) (x1,y1,x2,y2)
        gt_boxes_vec = box_xywh_to_xyxy(gt_bbox)[:, None, :].repeat((1, num_queries, 1)).view(-1, 4).clamp(min=0.0,
                                                                                                           max=1.0)  # (B,4) --> (B,1,4) --> (B,N,4)
        # compute giou and iou
        try:
            giou_loss, iou = self.objective['giou'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
        except:
            giou_loss, iou = torch.tensor(0.0).cuda(), torch.tensor(0.0).cuda()

        # compute l1 loss
        l1_loss = self.objective['l1'](pred_boxes_vec, gt_boxes_vec)  # (BN,4) (BN,4)
        # compute location loss
        if 'score_map' in pred_dict:
            location_loss = self.objective['focal'](pred_dict['score_map'], gt_gaussian_maps)
        else:
            location_loss = torch.tensor(0.0, device=l1_loss.device)
        # 按照HRCEUTrack,新增loss约束
        
        # rank_rgb_loss = self.rank_rgb_loss(pred_dict,gt_dict['search_anno'], gt_dict['template_anno'])
        rank_event_loss= self.rank_event_loss(pred_dict,gt_dict['search_anno'], gt_dict['template_anno'])

        #######
        # 计算特征一致性损失
        # feature_consistency_loss = F.mse_loss(pred_dict['A_enhanced_x'], pred_dict['A_original_x']) + F.mse_loss(pred_dict['A_enhanced_event_x'], pred_dict['A_original_event_x'])

        # # 标量化特征一致性损失
        # feature_consistency_loss = feature_consistency_loss.mean()
        
        # weighted sum
        loss = self.loss_weight['giou'] * giou_loss + self.loss_weight['l1'] * l1_loss + self.loss_weight['focal'] * location_loss + rank_event_loss * 1.2       # + 1 * feature_consistency_loss 
        if return_status:
            # status for log
            mean_iou = iou.detach().mean()
            status = {"Loss/total": loss.item(),
                      "Loss/giou": giou_loss.item(),
                      "Loss/l1": l1_loss.item(),
                      "Loss/location": location_loss.item(),
                      #'Loss/feature_consistency': feature_consistency_loss.item(),
                    #   "rank_rgb_loss": rank_rgb_loss.item(),
                      "rank_event_loss": rank_event_loss.item(),
                      "IoU": mean_iou.item()}
            return loss, status
        else:
            return loss

#################################### 按照 HRCEUTrack新增 loss  #########################################

    def _random_permute(self,matrix):
        # matrix = random.choice(matrix)
        b, c, h, w = matrix.shape
        idx = [  torch.randperm(c).to(matrix.device) for i in range(b)]
        idx = torch.stack(idx, dim=0)[:, :, None, None].repeat([1,1,h,w])
        # idx = torch.randperm(c)[None,:,None,None].repeat([b,1,h,w]).to(matrix.device)
        matrix01 = torch.gather(matrix, 1, idx)
        return matrix01
    def crop_flag(self, flag, global_index_s, global_index_t,H1 = 64, H2 = 256):
        B,Ls = global_index_s.shape
        B, Lt = global_index_t.shape
        B,C,L1,L2 = flag.shape
        flag_t = flag[:,:,:H1,:]
        flag_s = flag[:,:,H1:,:]

        flag_t = torch.gather(flag_t,2,global_index_t[:,None,:,None].repeat([1,C,1,L2]).long())
        flag_s = torch.gather(flag_s,2,global_index_s[:,None,:,None].repeat([1,C,1,L2]).long())
        flag = torch.cat([flag_t, flag_s], dim = 2)

        flag_t = flag[:,:,:,:H1]
        flag_s = flag[:,:,:,H1:]
        flag_t = torch.gather(flag_t,3,global_index_t[:,None,None,:].repeat([1,C,int(Ls+Lt),1]).long())
        flag_s = torch.gather(flag_s,3,global_index_s[:,None,None,:].repeat([1,C,int(Ls+Lt),1]).long())
        flag = torch.cat([flag_t, flag_s], dim = 3)
        B, C, L11, L12 = flag.shape
        try:
            assert(L11 == int(Lt + Ls))
            assert(L12 == int(Lt + Ls))
        except:
            print('L11:{}, L12:{}, L1:{}, L2:{}'.format(L11, L12, L1, L2))
        return flag
    def crop_fusion(self, flag, attn, global_index_s, global_index_t,H1 = 64, H2 = 256 ):
        flag = self.crop_flag(flag=flag, global_index_s=global_index_s, global_index_t=global_index_t)
        B,C,L1,L2 = flag.shape
        Ba, Ca, La, La2 = attn.shape
        _,idx1 = flag.mean(dim=3,keepdim=False).sort(dim=2,descending=True)
        # print('shape of flag:{}, idx1:{}'.format(flag.shape, idx1[:,:,:32,None].repeat([1,Ca,1,L2]).shape))
        flag = torch.gather(flag,2,idx1[:,:,:32,None].repeat([1,C,1,L2]).long())
        attn = torch.gather(attn,2,idx1[:,:,:32,None].repeat([1,Ca,1,L2]).long())
        _,idx2 = flag.mean(dim=2,keepdim=False).sort(dim=2,descending=True)
        flag = torch.gather(flag,3,idx2[:,:,None,:32].repeat([1,C,32,1]).long())
        attn = torch.gather(attn,3,idx2[:,:,None,:32].repeat([1,Ca,32,1]).long())
        return attn * flag


    def loss_rank_ori(self, outputs, targetsi, temp_annoi=None):
        """Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss
           targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4]
           The target boxes are expected in format (center_x, center_y, h, w), normalized by the image size.
        """
        attn = outputs['attn']
        # print('attn shape:',attn.shape)
        attn1 = torch.cat([attn[:,:,114:344,57:114], attn[:,:,114:344,344:]],dim=3)
        attn1 = attn1.mean(dim=0, keepdim=True).mean(dim=1, keepdim=True)
        # print('attn1 shape',attn1.shape)
        attn2 = torch.cat([attn[:,:,344:,:57], attn[:,:,344:,114:344]],dim=3)
        attn2 = attn2.mean(dim=0, keepdim=True).mean(dim=1, keepdim=True)
        # print('attn2 shape',attn2.shape)
        attn = torch.cat([attn1, attn2],dim=1)
        B, C, H, W = attn.shape
        _,s1,_ = torch.svd(attn.reshape([B*C, H, W]))
        s01 = torch.abs(s1 - 1)
        return torch.mean(s01)

    # def rank_rgb_loss(self, outputs, targetsi, temp_annoi=None):
    #     attn = outputs['attn']
    #     # print('attn.shape',attn.shape) # attn.shape torch.Size([2, 12, 320, 320])
    #     attn1 = torch.cat([attn[:,:,57:172,28:57], attn[:,:,57:172,172:]], dim=3)
    #     attn1 = attn1.mean(dim=0, keepdim=True).mean(dim=1, keepdim=True)
    #     attn2 = torch.cat([attn[:,:,172:,:28], attn[:,:,172:,57:172]], dim=3)
    #     attn2 = attn2.mean(dim=0, keepdim=True).mean(dim=1, keepdim=True)
    #     attn2 = F.interpolate(attn2, size=(115, 144), mode='bilinear', align_corners=False)
    #     attn = torch.cat([attn1, attn2], dim=1)
    #     B, C, H, W = attn.shape
    #     _, s1, _ = torch.svd(attn.reshape([B*C, H, W]))
    #     s01 = torch.abs(s1 - 1)
    #     return torch.mean(s01)

    def rank_event_loss(self, outputs, targetsi, temp_annoi=None):
        attn = outputs['attn']
        print('**********************************')
        print(attn)
        print('**********************************')
        print('attn.shape',attn.shape) 
        attn_show = attn  # 复制一份attn用于显示

        # 整体归一化处理
        attn_min = attn_show.min(dim=2, keepdim=True)[0].min(dim=3, keepdim=True)[0]
        attn_max = attn_show.max(dim=2, keepdim=True)[0].max(dim=3, keepdim=True)[0]
        attn_show = (attn_show - attn_min) / (attn_max - attn_min)

        attn1 = torch.cat([attn[:,:,114:344,57:114], attn[:,:,114:344,344:]],dim=3)
        attn1 = attn1.mean(dim=0, keepdim=True).mean(dim=1, keepdim=True)
        attn2 = torch.cat([attn[:,:,344:,:57], attn[:,:,344:,114:344]],dim=3)
        attn2 = attn2.mean(dim=0, keepdim=True).mean(dim=1, keepdim=True)

        # attn1 = torch.cat([attn[:,:,57:172,28:57], attn[:,:,57:172,172:]], dim=3)
        # attn1 = attn1.mean(dim=0, keepdim=True).mean(dim=1, keepdim=True)
        # attn2 = torch.cat([attn[:,:,172:,:28], attn[:,:,172:,57:172]], dim=3)
        # attn2 = attn2.mean(dim=0, keepdim=True).mean(dim=1, keepdim=True)
        # attn2 = F.interpolate(attn2, size=(115, 144), mode='bilinear', align_corners=False)
        attn = torch.cat([attn1, attn2], dim=1)
        B, C, H, W = attn.shape
        _, s1, _ = torch.svd(attn.reshape([B*C, H, W]))
        s02 = torch.abs(s1 - 1)

        # 设置阈值找到高响应区域
        threshold = 0.8 * attn_show.max().item()
        high_resp_mask = attn_show > threshold
        # labeled, num_features = label(high_resp_mask[0].cpu().numpy())
        labeled, num_features = label(high_resp_mask[0].cpu().numpy(), return_num=True)
        # 输出高响应区域的边界框
        props = regionprops(labeled)
        plt.figure(figsize=(10, 10))
        ax = plt.gca()
        for prop in props:
            bbox = prop.bbox
            print(f'High response region: {bbox}')  # 输出边界框信息
            rect = patches.Rectangle((bbox[1], bbox[0]), bbox[3] - bbox[1], bbox[2] - bbox[0], 
                                    linewidth=2, edgecolor='r', facecolor='none')
            ax.add_patch(rect)

        # 展示注意力图和标记的高响应区域
        for i in range(attn_show.shape[1]):
            plt.subplot(attn_show.shape[1] // 3 + 1, 3, i + 1)
            plt.imshow(attn_show[0, i].detach().cpu().numpy(), cmap='viridis', interpolation='nearest')
            plt.title(f'Channel {i+1}')
            plt.axis('off')
        plt.suptitle('Attention Maps')
        plt.show()
        return torch.mean(s02)
，其中我们将attn可视化之后，发现注意力特征可视化大致可以分为16个部分，然后作者将横轴定义为key，纵轴定义为qurey，然后key分为：模板（RGB），模板（event）,搜索（RGB），搜索（event）,query轴从上到下分别是模板（RGB），模板（event）,搜索（RGB），搜索（event）