在"""
Basic mamda_fetrack 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.utils.box_ops import box_xyxy_to_cxcywh
from lib.models.mamba_fetrack.models_mamba import create_block
from timm.models import create_model
import torch.nn.functional as F
from lib.models.mamba_fetrack.mamba_cross import CrossMamba
from thop import profile
from .noise import LDDiffuse

class Mamba_FEtrack(nn.Module):
    """ This is the base class for mamda_fetrack """

    def __init__(self, visionmamba, cross_mamba, 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 = visionmamba
        self.cross_mamba= cross_mamba
        self.box_head = 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)
        # 初始化 LDDiffuse
        self.lddiffuse = LDDiffuse(total_timesteps=1000, sampling_timesteps=1)
    def forward(self, template: torch.Tensor,
                search: torch.Tensor,
                event_template: torch.Tensor,       
                event_search: torch.Tensor,          
                ce_template_mask=None,
                ce_keep_rate=None,
                return_last_attn=False,
                ):
        rgb_feature = self.backbone.forward_features( z=template, x=search,                                                                     #[B, 320, 384]
                                                inference_params=None, if_random_cls_token_position=False, if_random_token_rank=False)
        event_feature = self.backbone.forward_features(z=event_template, x=event_search,                                                        #[B, 320, 384]
                                                inference_params=None, if_random_cls_token_position=False, if_random_token_rank=False)
        
        timesteps = torch.randint(0, self.lddiffuse.total_timesteps, (rgb_feature.size(0),), dtype=torch.long, device=rgb_feature.device)

        # 正向扩散过程
        diffused_rgb_feature = self.lddiffuse.forward_diffusion(rgb_feature, timesteps)
        diffused_event_feature = self.lddiffuse.forward_diffusion(event_feature, timesteps)
        ############# 在选定时间步上，被噪声处理后的扩散特征。 #############

        # 特征融合
        residual_event_f = 0
        residual_rgb_f = 0
        event_f = self.cross_mamba(diffused_event_feature, residual_event_f, diffused_rgb_feature) + event_feature
        rgb_f = self.cross_mamba(diffused_rgb_feature, residual_rgb_f, diffused_event_feature) + rgb_feature

        # 逆向扩散过程
        reconstructed_rgb = self.lddiffuse.reverse_diffusion(diffused_rgb_feature, timesteps, rgb_feature)
        reconstructed_event = self.lddiffuse.reverse_diffusion(diffused_event_feature, timesteps, event_feature)
        ############## 通过逆向扩散恢复的未被扩散特征，意图将特征恢复为接近原始输入的状态。 ##############
        
        
        
        # 使用逆向扩散后的特征进行预测
        event_search = reconstructed_event[:, -self.feat_len_s:]
        rgb_search = reconstructed_rgb[:, -self.feat_len_s:]
        x = torch.cat((event_search, rgb_search), dim=-1)

        feat_last = x             
        if isinstance(x, list):
            feat_last = x[-1]
        out = self.forward_head(feat_last, None)
       
        out['backbone_feat'] = x
        return out

    def forward_head(self, cat_feature, gt_score_map=None):
        """
        cat_feature: output embeddings of the backbone, it can be (HW1+HW2, B, C) or (HW2, B, C)
        """
        search_feature = cat_feature
        opt = (search_feature.unsqueeze(-1)).permute((0, 3, 2, 1)).contiguous()           # opt.shape = torch.Size([B, 1, 384, 256])
        bs, Nq, C, HW = opt.size()
        opt_feat = opt.view(-1, C, self.feat_sz_s, self.feat_sz_s)                       # opt_feat.shape = torch.Size([B, 384, 16, 16])

        if self.head_type == "CORNER":
            # run the corner head
            pred_box, score_map = self.box_head(opt_feat, True)
            outputs_coord = box_xyxy_to_cxcywh(pred_box)
            outputs_coord_new = outputs_coord.view(bs, Nq, 4)
            out = {'pred_boxes': outputs_coord_new,
                   'score_map': score_map,
                   }
            return out

        elif self.head_type == "CENTER":
            # run the center head
           
            score_map_ctr, bbox, size_map, offset_map = self.box_head(opt_feat, gt_score_map)
            outputs_coord = bbox
            outputs_coord_new = outputs_coord.view(bs, Nq, 4)
            out = {'pred_boxes': outputs_coord_new,
                   'score_map': score_map_ctr,
                   'size_map': size_map,
                   'offset_map': offset_map}
            return out
        else:
            raise NotImplementedError


def build_mamba_fetrack(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')
    # pretrained_path = os.path.join(current_dir, '../../../pretrained_models')
    if cfg.MODEL.PRETRAIN_FILE and ('Mamba_FETrack' not in cfg.MODEL.PRETRAIN_FILE) and training:
        pretrained = os.path.join(pretrained_path, cfg.MODEL.PRETRAIN_FILE)
    else:
        pretrained = ''
    
    backbone = create_model( model_name= cfg.MODEL.BACKBONE.TYPE, pretrained= pretrained, num_classes=1000,
            drop_rate=0.0, drop_path_rate=cfg.TRAIN.DROP_PATH_RATE, drop_block_rate=None, img_size=256
            )
    hidden_dim = 384
    cross_mamba = CrossMamba(hidden_dim)
    box_head = build_box_head(cfg, hidden_dim*2)
    model = Mamba_FEtrack(
        backbone,
        cross_mamba,
        box_head,
        aux_loss=False,
        head_type=cfg.MODEL.HEAD.TYPE,
    )
   
    if 'Mamba_FETrack' 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
中集成了import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import matplotlib.pyplot as plt
import gc


def cosine_beta_schedule(timesteps, s=0.008):
    """
    cosine schedule
    as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
    """
    steps = timesteps + 1
    x = torch.linspace(0, timesteps, steps, dtype=torch.float64)
    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * math.pi * 0.5) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
    return torch.clip(betas, 0, 0.999)

def exists(x):
    return x is not None

def default(val, d):
    if exists(val):
        return val
    return d() if callable(d) else d


def extract(a, t, x_shape):
    """extract the appropriate  t  index for a batch of indices"""
    batch_size = t.shape[0]
    out = a.gather(-1, t)
    return out.reshape(batch_size, *((1,) * (len(x_shape) - 1)))


class LDDiffuse(nn.Module):
    def __init__(self, total_timesteps=1000, sampling_timesteps=5):
        super(LDDiffuse, self).__init__()
        
        # 初始化扩散参数
        self.total_timesteps = total_timesteps
        self.sampling_timesteps = sampling_timesteps

        # 计算betas和alphas
        betas = cosine_beta_schedule(self.total_timesteps)
        alphas = 1.0 - betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)
        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)

        self.register_buffer('betas', betas)
        self.register_buffer('alphas_cumprod', alphas_cumprod)
        self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev)

        # 存储计算中可能使用的中间变量
        self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))
        self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1.0 - alphas_cumprod))
        self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1.0 / alphas_cumprod))

    def forward_diffusion(self, x_start, t, noise=None):
        """
        前向扩散过程：给定初始输入，生成其在时间 t 的扩散状态。
        """
        if noise is None:
            noise = torch.randn_like(x_start)
        sqrt_alphas_cumprod_t = extract(self.sqrt_alphas_cumprod, t, x_start.shape)
        sqrt_one_minus_alphas_cumprod_t = extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
        return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise

    def reverse_diffusion(self, xt, t, pred_x0):
        """
        反向扩散过程：通过预测的初始状态逆推在时间 t 的状态。
        """
        reciprocal_sqrt_alphas = extract(self.sqrt_recip_alphas_cumprod, t, xt.shape)
        reciprocal_sqrt_alpha_minus_one = extract(self.sqrt_recip_alphas_cumprod - 1, t, xt.shape)
        return (reciprocal_sqrt_alphas * xt - pred_x0) / reciprocal_sqrt_alpha_minus_one

    def forward(self, x, t, noise=None, reverse=False, pred_x0=None):
            """
            通用前向传播接口，控制前向或反向扩散过程。

            Args:
                x: 输入特征，中间状态xt或起始状态x_start。
                t: 时间步骤张量。
                noise: 可选噪声张量，仅在前向扩散中使用。
                reverse: 布尔值，指定是否进行反向扩散。
                pred_x0: 预测的x0，仅在反向扩散中使用。
            
            Returns:
                处理后的特征。
            """
            if reverse:
                assert pred_x0 is not None, "反向扩散需要给定初始预测x0"
                return self.reverse_diffusion(x, t, pred_x0)
            else:
                return self.forward_diffusion(x, t, noise)，因此在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


class Mamba_FETrackActor(BaseActor):
    """ Actor for training Mamba_FETrack 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)
        # 引入扩散模型, 需要增加扩散损失
        
        # weighted sum
        loss = self.loss_weight['giou'] * giou_loss + self.loss_weight['l1'] * l1_loss + self.loss_weight['focal'] * location_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(),
                      "IoU": mean_iou.item()}
            return loss, status
        else:
            return loss
中需要增加扩散损失，给出修改代码。