bev-project/mmdet3d/models/heads/vector_map/maptr_head.py

"""
MapTR风格的矢量地图预测Head
集成到BEVFusion框架中

基于MapTRv2的实现思路，适配BEVFusion的多任务架构
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from mmcv.cnn import build_norm_layer
from mmdet.models import HEADS


@HEADS.register_module()
class MapTRHead(nn.Module):
    """
    MapTR风格的矢量地图预测Head
    
    输入: BEV特征 (B, C, H, W)
    输出: 矢量化的地图元素
    
    Args:
        in_channels: BEV特征通道数
        num_classes: 地图元素类别数（如3: divider, boundary, crossing）
        num_queries: 预测的矢量数量
        num_points: 每个矢量的采样点数
        embed_dims: Transformer embedding维度
        num_decoder_layers: Transformer decoder层数
        num_heads: Multi-head attention的head数
    """
    
    def __init__(
        self,
        in_channels=256,
        num_classes=3,
        num_queries=50,
        num_points=20,
        embed_dims=256,
        num_decoder_layers=6,
        num_heads=8,
        dropout=0.1,
        loss_cls_weight=2.0,
        loss_reg_weight=5.0,
        loss_chamfer_weight=2.0,
        score_threshold=0.3,
        nms_threshold=0.5,
        **kwargs
    ):
        super().__init__()
        
        self.in_channels = in_channels
        self.num_classes = num_classes
        self.num_queries = num_queries
        self.num_points = num_points
        self.embed_dims = embed_dims
        self.score_threshold = score_threshold
        self.nms_threshold = nms_threshold
        
        # 损失权重
        self.loss_cls_weight = loss_cls_weight
        self.loss_reg_weight = loss_reg_weight
        self.loss_chamfer_weight = loss_chamfer_weight
        
        # 1. Input projection (BEV features -> Transformer dim)
        self.input_proj = nn.Conv2d(in_channels, embed_dims, kernel_size=1)
        
        # 2. Query Embedding (可学习的query，代表潜在的地图元素)
        self.query_embed = nn.Embedding(num_queries, embed_dims)
        
        # 3. Positional Encoding for BEV
        self.bev_pos_embed = PositionEmbeddingSine(
            num_feats=embed_dims // 2,
            normalize=True
        )
        
        # 4. Transformer Decoder
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=embed_dims,
            nhead=num_heads,
            dim_feedforward=embed_dims * 4,
            dropout=dropout,
            activation='relu',
            batch_first=True,
        )
        self.decoder = nn.TransformerDecoder(
            decoder_layer,
            num_layers=num_decoder_layers
        )
        
        # 5. 预测头
        # 5.1 分类头 (预测矢量的类别)
        self.cls_head = nn.Linear(embed_dims, num_classes + 1)  # +1 for background
        
        # 5.2 点坐标回归头 (预测矢量的点序列)
        self.reg_head = nn.Sequential(
            nn.Linear(embed_dims, embed_dims),
            nn.ReLU(inplace=True),
            nn.Linear(embed_dims, embed_dims),
            nn.ReLU(inplace=True),
            nn.Linear(embed_dims, num_points * 2)  # (x, y) * num_points
        )
        
        # 5.3 置信度头
        self.score_head = nn.Linear(embed_dims, 1)
        
        # 类别名称映射
        self.class_names = ['divider', 'boundary', 'ped_crossing']
    
    def forward(self, bev_features, img_metas=None):
        """
        前向传播
        
        Args:
            bev_features: (B, C, H, W) - 来自BEV decoder的特征
            img_metas: 元数据（可选）
        
        Returns:
            dict: 包含分类、回归、置信度预测
        """
        B, C, H, W = bev_features.shape
        
        # 1. Project BEV features
        bev_feat = self.input_proj(bev_features)  # (B, embed_dims, H, W)
        
        # 2. Flatten BEV features for Transformer
        # (B, embed_dims, H, W) -> (B, H*W, embed_dims)
        bev_flatten = bev_feat.flatten(2).permute(0, 2, 1)
        
        # 3. 添加位置编码
        pos_embed = self.bev_pos_embed(bev_feat)  # (B, embed_dims, H, W)
        pos_embed = pos_embed.flatten(2).permute(0, 2, 1)  # (B, H*W, embed_dims)
        
        # 4. 准备Query
        query_embed = self.query_embed.weight.unsqueeze(0).repeat(B, 1, 1)
        # query_embed: (B, num_queries, embed_dims)
        
        # 5. Transformer Decoder
        # tgt: query embeddings
        # memory: BEV features + positional encoding
        hs = self.decoder(
            tgt=query_embed,
            memory=bev_flatten + pos_embed,  # 加上位置编码
        )
        # hs: (B, num_queries, embed_dims)
        
        # 6. 预测
        # 6.1 分类 (哪种类型的地图元素)
        cls_scores = self.cls_head(hs)  # (B, num_queries, num_classes+1)
        
        # 6.2 点坐标 (矢量的形状)
        reg_preds_flat = self.reg_head(hs)  # (B, num_queries, num_points*2)
        reg_preds = reg_preds_flat.view(B, self.num_queries, self.num_points, 2)
        # Sigmoid归一化到[0, 1]（相对于BEV范围）
        reg_preds = reg_preds.sigmoid()
        
        # 6.3 置信度 (这个预测有多可靠)
        obj_scores = self.score_head(hs).sigmoid()  # (B, num_queries, 1)
        
        return {
            'cls_scores': cls_scores,
            'reg_preds': reg_preds,
            'obj_scores': obj_scores,
        }
    
    def loss(self, predictions, gt_vectors_labels, gt_vectors_points):
        """
        计算损失
        
        Args:
            predictions: forward的输出
            gt_vectors_labels: (B, max_num_vectors) - GT类别标签
            gt_vectors_points: (B, max_num_vectors, num_points, 2) - GT点坐标（归一化）
        
        Returns:
            dict: 各项损失
        """
        cls_scores = predictions['cls_scores']  # (B, num_queries, num_classes+1)
        reg_preds = predictions['reg_preds']    # (B, num_queries, num_points, 2)
        obj_scores = predictions['obj_scores']  # (B, num_queries, 1)
        
        B = cls_scores.shape[0]
        device = cls_scores.device
        
        # 1. Hungarian Matching (为预测找到最佳匹配的GT)
        matched_indices = self.hungarian_matcher(
            cls_scores, reg_preds, obj_scores,
            gt_vectors_labels, gt_vectors_points
        )
        
        losses = {}
        
        # 2. 分类损失
        target_classes = torch.full(
            (B, self.num_queries),
            self.num_classes,  # background class
            dtype=torch.long,
            device=device
        )
        
        # 填充匹配的GT类别
        for b, (pred_idx, gt_idx) in enumerate(matched_indices):
            if len(pred_idx) > 0:
                valid_gt = gt_vectors_labels[b, gt_idx] >= 0
                target_classes[b, pred_idx[valid_gt]] = gt_vectors_labels[b, gt_idx[valid_gt]]
        
        # Focal Loss
        losses['loss_cls'] = self.focal_loss(
            cls_scores.flatten(0, 1),  # (B*num_queries, num_classes+1)
            target_classes.flatten()    # (B*num_queries,)
        ) * self.loss_cls_weight
        
        # 3. 点坐标回归损失 (只对匹配上的计算)
        loss_reg = 0
        num_pos = 0
        
        for b, (pred_idx, gt_idx) in enumerate(matched_indices):
            if len(pred_idx) > 0:
                # 过滤有效GT
                valid_mask = gt_vectors_labels[b, gt_idx] >= 0
                if valid_mask.sum() > 0:
                    pred_points = reg_preds[b, pred_idx[valid_mask]]
                    gt_points = gt_vectors_points[b, gt_idx[valid_mask]]
                    
                    # L1 Loss
                    loss_reg += F.l1_loss(pred_points, gt_points, reduction='sum')
                    num_pos += valid_mask.sum().item()
        
        losses['loss_reg'] = (loss_reg / max(num_pos, 1)) * self.loss_reg_weight
        
        # 4. Chamfer Distance Loss (点集距离)
        loss_chamfer = 0
        for b, (pred_idx, gt_idx) in enumerate(matched_indices):
            if len(pred_idx) > 0:
                valid_mask = gt_vectors_labels[b, gt_idx] >= 0
                if valid_mask.sum() > 0:
                    pred_pts = reg_preds[b, pred_idx[valid_mask]]  # (N, num_points, 2)
                    gt_pts = gt_vectors_points[b, gt_idx[valid_mask]]
                    
                    # 计算Chamfer距离
                    chamfer = self.chamfer_distance(pred_pts, gt_pts)
                    loss_chamfer += chamfer
        
        losses['loss_chamfer'] = (loss_chamfer / B) * self.loss_chamfer_weight
        
        return losses
    
    def hungarian_matcher(self, cls_scores, reg_preds, obj_scores, 
                         gt_labels, gt_points):
        """
        Hungarian匹配算法
        
        为每个样本中的预测找到最佳匹配的GT
        
        Returns:
            list of tuples: [(pred_indices, gt_indices), ...]
        """
        from scipy.optimize import linear_sum_assignment
        
        B = cls_scores.shape[0]
        matched_indices = []
        
        for b in range(B):
            # 过滤有效的GT (class >= 0)
            valid_gt_mask = gt_labels[b] >= 0
            valid_gt_idx = valid_gt_mask.nonzero(as_tuple=True)[0]
            num_valid_gt = len(valid_gt_idx)
            
            if num_valid_gt == 0:
                # 没有有效GT
                matched_indices.append((torch.tensor([], dtype=torch.long), 
                                       torch.tensor([], dtype=torch.long)))
                continue
            
            # 计算cost matrix
            # Cost = alpha * cost_cls + beta * cost_reg
            
            # 1. 分类cost
            cls_prob = cls_scores[b].softmax(dim=-1)  # (num_queries, num_classes+1)
            valid_gt_labels = gt_labels[b, valid_gt_idx]
            # 负对数似然作为cost
            cost_cls = -cls_prob[:, valid_gt_labels].transpose(0, 1)
            # cost_cls: (num_valid_gt, num_queries)
            
            # 2. 点坐标cost (L1距离)
            pred_pts = reg_preds[b].unsqueeze(0)  # (1, num_queries, num_points, 2)
            gt_pts = gt_points[b, valid_gt_idx].unsqueeze(1)  # (num_valid_gt, 1, num_points, 2)
            cost_reg = (pred_pts - gt_pts).abs().sum(dim=-1).mean(dim=-1)
            # cost_reg: (num_valid_gt, num_queries)
            
            # 3. 总cost
            cost = cost_cls + 5.0 * cost_reg
            
            # 4. Hungarian算法求最优匹配
            gt_idx_matched, pred_idx_matched = linear_sum_assignment(
                cost.detach().cpu().numpy()
            )
            
            # 转回原始GT索引
            gt_idx_original = valid_gt_idx[gt_idx_matched]
            
            matched_indices.append((
                torch.tensor(pred_idx_matched, dtype=torch.long),
                gt_idx_original
            ))
        
        return matched_indices
    
    def focal_loss(self, inputs, targets, alpha=0.25, gamma=2.0):
        """
        Focal Loss用于分类
        
        处理类别不平衡问题（背景类很多，地图元素类较少）
        """
        ce_loss = F.cross_entropy(inputs, targets, reduction='none')
        p_t = torch.exp(-ce_loss)
        focal_weight = alpha * (1 - p_t) ** gamma
        loss = focal_weight * ce_loss
        return loss.mean()
    
    def chamfer_distance(self, pred_points, gt_points):
        """
        Chamfer距离：衡量两个点集的相似度
        
        Args:
            pred_points: (N, num_points, 2)
            gt_points: (N, num_points, 2)
        
        Returns:
            chamfer_dist: scalar
        """
        # 计算点对点距离矩阵
        # (N, num_points, 2) -> (N, num_points, num_points)
        dist_matrix = torch.cdist(pred_points, gt_points)
        
        # 从pred到gt的最小距离
        min_dist_pred_to_gt = dist_matrix.min(dim=-1)[0]  # (N, num_points)
        
        # 从gt到pred的最小距离
        min_dist_gt_to_pred = dist_matrix.min(dim=-2)[0]  # (N, num_points)
        
        # Chamfer距离 = 双向最小距离之和
        chamfer = min_dist_pred_to_gt.mean() + min_dist_gt_to_pred.mean()
        
        return chamfer
    
    def get_vectors(self, predictions, img_metas):
        """
        后处理：从模型预测提取最终的矢量地图
        
        Args:
            predictions: forward的输出
            img_metas: 元数据（包含BEV范围等信息）
        
        Returns:
            list of dict: 每个样本的矢量地图
                [
                    {
                        'vectors': [
                            {'class': 0, 'class_name': 'divider', 
                             'points': [[x1,y1], ...], 'score': 0.95},
                            ...
                        ]
                    },
                    ...
                ]
        """
        cls_scores = predictions['cls_scores']  # (B, num_queries, num_classes+1)
        reg_preds = predictions['reg_preds']    # (B, num_queries, num_points, 2)
        obj_scores = predictions['obj_scores']  # (B, num_queries, 1)
        
        B = cls_scores.shape[0]
        results = []
        
        for b in range(B):
            # 1. 获取分类结果
            cls_pred = cls_scores[b].argmax(dim=-1)  # (num_queries,)
            cls_conf = cls_scores[b].softmax(dim=-1).max(dim=-1)[0]  # (num_queries,)
            scores = obj_scores[b].squeeze(-1)  # (num_queries,)
            
            # 2. 过滤
            # 过滤背景类
            valid_mask = cls_pred < self.num_classes
            # 过滤低置信度
            valid_mask &= (cls_conf > self.score_threshold)
            valid_mask &= (scores > self.score_threshold)
            
            valid_idx = valid_mask.nonzero(as_tuple=True)[0]
            
            # 3. NMS (去除重复的矢量)
            if len(valid_idx) > 0:
                keep_idx = self.vector_nms(
                    reg_preds[b, valid_idx],
                    scores[valid_idx],
                    iou_threshold=self.nms_threshold
                )
                final_idx = valid_idx[keep_idx]
            else:
                final_idx = torch.tensor([], dtype=torch.long)
            
            # 4. 反归一化到真实坐标
            vectors = []
            for idx in final_idx:
                # 归一化坐标[0,1] -> BEV坐标
                points_norm = reg_preds[b, idx].detach().cpu().numpy()  # (num_points, 2)
                
                # 获取BEV范围（从img_metas或使用默认值）
                if img_metas is not None and b < len(img_metas):
                    x_range = img_metas[b].get('x_range', [-50, 50])
                    y_range = img_metas[b].get('y_range', [-50, 50])
                else:
                    x_range, y_range = [-50, 50], [-50, 50]
                
                points_real = self.denormalize_points(points_norm, x_range, y_range)
                
                class_id = cls_pred[idx].item()
                vectors.append({
                    'class': class_id,
                    'class_name': self.get_class_name(class_id),
                    'points': points_real.tolist(),
                    'score': scores[idx].item(),
                })
            
            results.append({'vectors': vectors})
        
        return results
    
    def denormalize_points(self, points_norm, x_range, y_range):
        """将归一化坐标[0,1]转换为真实BEV坐标"""
        points = points_norm.copy()
        points[:, 0] = points[:, 0] * (x_range[1] - x_range[0]) + x_range[0]
        points[:, 1] = points[:, 1] * (y_range[1] - y_range[0]) + y_range[0]
        return points
    
    def vector_nms(self, points, scores, iou_threshold=0.5):
        """
        矢量NMS：去除重叠的矢量预测
        
        使用Chamfer距离作为相似度度量
        """
        if len(points) == 0:
            return []
        
        N = len(points)
        
        # 计算所有矢量对之间的相似度
        similarities = torch.zeros(N, N).to(points.device)
        
        for i in range(N):
            for j in range(i + 1, N):
                # Chamfer距离
                dist = torch.cdist(points[i:i+1], points[j:j+1])[0]  # (num_points, num_points)
                chamfer = dist.min(dim=0)[0].mean() + dist.min(dim=1)[0].mean()
                
                # 转为相似度（距离越小，相似度越高）
                similarity = torch.exp(-chamfer * 10)  # 放缩因子
                similarities[i, j] = similarity
                similarities[j, i] = similarity
        
        # NMS
        keep = []
        order = scores.argsort(descending=True).cpu().numpy()
        
        while len(order) > 0:
            i = order[0]
            keep.append(i)
            
            if len(order) == 1:
                break
            
            # 找到与当前矢量相似度低的（保留）
            remaining_idx = order[1:]
            remaining_mask = similarities[i, remaining_idx].cpu().numpy() <= iou_threshold
            order = remaining_idx[remaining_mask]
        
        return keep
    
    def get_class_name(self, class_id):
        """获取类别名称"""
        if 0 <= class_id < len(self.class_names):
            return self.class_names[class_id]
        return 'unknown'


class PositionEmbeddingSine(nn.Module):
    """
    正弦位置编码
    
    为BEV特征的每个位置生成唯一的位置编码
    """
    
    def __init__(self, num_feats=128, temperature=10000, normalize=True):
        super().__init__()
        self.num_feats = num_feats
        self.temperature = temperature
        self.normalize = normalize
    
    def forward(self, x):
        """
        Args:
            x: (B, C, H, W) - BEV features
        
        Returns:
            pos: (B, num_feats*2, H, W) - 位置编码
        """
        B, C, H, W = x.shape
        device = x.device
        
        # 创建坐标网格
        y_embed = torch.arange(H, dtype=torch.float32, device=device)
        x_embed = torch.arange(W, dtype=torch.float32, device=device)
        
        if self.normalize:
            # 归一化到[0, 1]
            y_embed = y_embed / H
            x_embed = x_embed / W
        
        # 计算频率
        dim_t = torch.arange(self.num_feats, dtype=torch.float32, device=device)
        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_feats)
        
        # y方向的位置编码
        pos_y = y_embed[:, None] / dim_t  # (H, num_feats)
        pos_y = torch.stack([
            pos_y[:, 0::2].sin(),
            pos_y[:, 1::2].cos()
        ], dim=2).flatten(1)  # (H, num_feats)
        
        # x方向的位置编码
        pos_x = x_embed[:, None] / dim_t  # (W, num_feats)
        pos_x = torch.stack([
            pos_x[:, 0::2].sin(),
            pos_x[:, 1::2].cos()
        ], dim=2).flatten(1)  # (W, num_feats)
        
        # 组合y和x的位置编码
        pos = torch.zeros(B, self.num_feats * 2, H, W, device=device)
        pos[:, :self.num_feats, :, :] = pos_y[:, :, None].repeat(1, 1, W).unsqueeze(0).repeat(B, 1, 1, 1)
        pos[:, self.num_feats:, :, :] = pos_x[:, None, :].repeat(1, H, 1).unsqueeze(0).repeat(B, 1, 1, 1)
        
        return pos


# 评估相关函数

def evaluate_vector_map(pred_vectors, gt_vectors, iou_thresholds=[0.5, 0.75]):
    """
    评估矢量地图预测性能
    
    Args:
        pred_vectors: 预测的矢量列表
        gt_vectors: GT矢量列表
        iou_thresholds: IoU阈值列表
    
    Returns:
        dict: 评估指标
    """
    metrics = {}
    
    for iou_thr in iou_thresholds:
        # 计算AP
        ap = compute_vector_ap(pred_vectors, gt_vectors, iou_threshold=iou_thr)
        metrics[f'AP@{iou_thr}'] = ap
    
    # Chamfer距离
    cd = compute_chamfer_distance_metric(pred_vectors, gt_vectors)
    metrics['chamfer_distance'] = cd
    
    return metrics