572 lines
20 KiB
Python
572 lines
20 KiB
Python
"""
|
||
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
|
||
|