275 lines
10 KiB
Python
275 lines
10 KiB
Python
|
#!/usr/bin/env python
|
|||
|
|
"""
|
|||
|
|
可视化单batch推理结果 - BEV分割和3D检测
|
|||
|
|
"""
|
|||
|
|
import os
|
|||
|
|
import pickle
|
|||
|
|
import numpy as np
|
|||
|
|
import matplotlib.pyplot as plt
|
|||
|
|
import torch
|
|||
|
|
from PIL import Image
|
|||
|
|
import argparse
|
|||
|
|
try:
|
|||
|
|
import cv2
|
|||
|
|
except ImportError:
|
|||
|
|
cv2 = None
|
|||
|
|
print("Warning: cv2 not available, some visualization features may be limited")
|
|||
|
|
|
|||
|
|
# NuScenes类别颜色映射
|
|||
|
|
NUSCENES_COLORS = {
|
|||
|
|
'drivable_area': '#00FF00', # 绿色 - 可行驶区域
|
|||
|
|
'ped_crossing': '#FFFF00', # 黄色 - 人行横道
|
|||
|
|
'walkway': '#FFA500', # 橙色 - 人行道
|
|||
|
|
'stop_line': '#FF0000', # 红色 - 停止线
|
|||
|
|
'carpark_area': '#800080', # 紫色 - 停车区
|
|||
|
|
'divider': '#000000', # 黑色 - 分隔线
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# BEV分割类别索引
|
|||
|
|
MAP_CLASSES = [
|
|||
|
|
'drivable_area', # 0
|
|||
|
|
'ped_crossing', # 1
|
|||
|
|
'walkway', # 2
|
|||
|
|
'stop_line', # 3
|
|||
|
|
'carpark_area', # 4
|
|||
|
|
'divider', # 5
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
def visualize_bev_segmentation(pred_masks, gt_masks=None, sample_idx=0, save_path=None):
|
|||
|
|
"""
|
|||
|
|
可视化BEV分割结果
|
|||
|
|
Args:
|
|||
|
|
pred_masks: 预测分割mask, shape [6, H, W]
|
|||
|
|
gt_masks: ground truth分割mask, shape [6, H, W]
|
|||
|
|
sample_idx: 样本索引
|
|||
|
|
save_path: 保存路径
|
|||
|
|
"""
|
|||
|
|
fig, axes = plt.subplots(2, 4, figsize=(20, 10))
|
|||
|
|
fig.suptitle(f'BEV分割结果 - 样本 {sample_idx}', fontsize=16)
|
|||
|
|
|
|||
|
|
# 预测结果
|
|||
|
|
for i in range(6):
|
|||
|
|
ax = axes[0, i]
|
|||
|
|
mask = pred_masks[i].cpu().numpy() if torch.is_tensor(pred_masks[i]) else pred_masks[i]
|
|||
|
|
|
|||
|
|
# 创建彩色mask
|
|||
|
|
colored_mask = np.zeros((*mask.shape, 3), dtype=np.uint8)
|
|||
|
|
color = NUSCENES_COLORS[MAP_CLASSES[i]]
|
|||
|
|
# 将hex颜色转换为RGB
|
|||
|
|
r = int(color[1:3], 16)
|
|||
|
|
g = int(color[3:5], 16)
|
|||
|
|
b = int(color[5:7], 16)
|
|||
|
|
|
|||
|
|
colored_mask[mask > 0.5] = [r, g, b]
|
|||
|
|
|
|||
|
|
ax.imshow(colored_mask)
|
|||
|
|
ax.set_title(f'{MAP_CLASSES[i]} (Pred)')
|
|||
|
|
ax.axis('off')
|
|||
|
|
|
|||
|
|
# Ground Truth (如果提供)
|
|||
|
|
if gt_masks is not None:
|
|||
|
|
for i in range(6):
|
|||
|
|
if i < 4: # 只显示前4个类别
|
|||
|
|
ax = axes[1, i]
|
|||
|
|
mask = gt_masks[i].cpu().numpy() if torch.is_tensor(gt_masks[i]) else gt_masks[i]
|
|||
|
|
|
|||
|
|
colored_mask = np.zeros((*mask.shape, 3), dtype=np.uint8)
|
|||
|
|
color = NUSCENES_COLORS[MAP_CLASSES[i]]
|
|||
|
|
r = int(color[1:3], 16)
|
|||
|
|
g = int(color[3:5], 16)
|
|||
|
|
b = int(color[5:7], 16)
|
|||
|
|
|
|||
|
|
colored_mask[mask > 0.5] = [r, g, b]
|
|||
|
|
|
|||
|
|
ax.imshow(colored_mask)
|
|||
|
|
ax.set_title(f'{MAP_CLASSES[i]} (GT)')
|
|||
|
|
ax.axis('off')
|
|||
|
|
|
|||
|
|
# 第4个位置显示GT叠加结果
|
|||
|
|
ax = axes[1, 3]
|
|||
|
|
combined_gt = np.zeros((gt_masks.shape[1], gt_masks.shape[2], 3), dtype=np.uint8)
|
|||
|
|
for i in range(6):
|
|||
|
|
mask = gt_masks[i].cpu().numpy() if torch.is_tensor(gt_masks[i]) else gt_masks[i]
|
|||
|
|
if mask.max() > 0.5:
|
|||
|
|
color = NUSCENES_COLORS[MAP_CLASSES[i]]
|
|||
|
|
r = int(color[1:3], 16)
|
|||
|
|
g = int(color[3:5], 16)
|
|||
|
|
b = int(color[5:7], 16)
|
|||
|
|
combined_gt[mask > 0.5] = [r, g, b]
|
|||
|
|
|
|||
|
|
ax.imshow(combined_gt)
|
|||
|
|
ax.set_title('Ground Truth 叠加')
|
|||
|
|
ax.axis('off')
|
|||
|
|
else:
|
|||
|
|
# 如果没有GT，在第二行显示预测的统计信息
|
|||
|
|
axes[1, 0].text(0.5, 0.5, 'Ground Truth\n不可用',
|
|||
|
|
transform=axes[1, 0].transAxes, ha='center', va='center',
|
|||
|
|
fontsize=12, bbox=dict(boxstyle="round,pad=0.3", facecolor="lightgray"))
|
|||
|
|
axes[1, 0].set_title('Ground Truth (N/A)')
|
|||
|
|
axes[1, 0].axis('off')
|
|||
|
|
|
|||
|
|
# 隐藏其他子图
|
|||
|
|
for i in range(1, 4):
|
|||
|
|
axes[1, i].axis('off')
|
|||
|
|
|
|||
|
|
plt.tight_layout()
|
|||
|
|
if save_path:
|
|||
|
|
plt.savefig(save_path, dpi=300, bbox_inches='tight')
|
|||
|
|
print(f"BEV分割可视化已保存: {save_path}")
|
|||
|
|
plt.show()
|
|||
|
|
|
|||
|
|
def visualize_3d_detection(boxes_3d, scores_3d, labels_3d, sample_idx=0, save_path=None):
|
|||
|
|
"""
|
|||
|
|
可视化3D检测结果
|
|||
|
|
Args:
|
|||
|
|
boxes_3d: 3D检测框
|
|||
|
|
scores_3d: 检测置信度
|
|||
|
|
labels_3d: 检测类别标签
|
|||
|
|
sample_idx: 样本索引
|
|||
|
|
save_path: 保存路径
|
|||
|
|
"""
|
|||
|
|
fig, ax = plt.subplots(1, 1, figsize=(12, 10))
|
|||
|
|
|
|||
|
|
# 创建空的BEV图像 (598x598)
|
|||
|
|
bev_image = np.ones((598, 598, 3), dtype=np.uint8) * 240 # 浅灰色背景
|
|||
|
|
|
|||
|
|
# BEV坐标范围: [-50, 50] x [-50, 50] 米
|
|||
|
|
# 分辨率: 0.167米/像素 (从之前的配置)
|
|||
|
|
bev_range = 50.0
|
|||
|
|
bev_resolution = 0.167
|
|||
|
|
bev_size = int(2 * bev_range / bev_resolution) # 598
|
|||
|
|
|
|||
|
|
detection_count = 0
|
|||
|
|
if len(boxes_3d) > 0:
|
|||
|
|
print(f"正在可视化 {len(boxes_3d)} 个3D检测框...")
|
|||
|
|
|
|||
|
|
# 将3D框投影到BEV平面
|
|||
|
|
for i, (box, score, label) in enumerate(zip(boxes_3d, scores_3d, labels_3d)):
|
|||
|
|
if torch.is_tensor(box):
|
|||
|
|
box = box.cpu().numpy()
|
|||
|
|
if torch.is_tensor(score):
|
|||
|
|
score = score.item() if score.numel() == 1 else score.cpu().numpy()
|
|||
|
|
if torch.is_tensor(label):
|
|||
|
|
label = label.item() if label.numel() == 1 else label.cpu().numpy()
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# LiDAR坐标系: x(前), y(左), z(上)
|
|||
|
|
center_x, center_y = box[0], box[1]
|
|||
|
|
|
|||
|
|
# 转换为BEV像素坐标
|
|||
|
|
# BEV原点在图像中心，x轴向右，y轴向下
|
|||
|
|
bev_x = int((center_x + bev_range) / bev_resolution)
|
|||
|
|
bev_y = int((bev_range - center_y) / bev_resolution) # y轴翻转
|
|||
|
|
|
|||
|
|
if 0 <= bev_x < bev_size and 0 <= bev_y < bev_size:
|
|||
|
|
# 根据置信度设置颜色深浅
|
|||
|
|
alpha = min(1.0, score) # 置信度越高颜色越深
|
|||
|
|
color_intensity = int(255 * alpha)
|
|||
|
|
|
|||
|
|
# 使用红色系表示检测框
|
|||
|
|
color = (255, 255 - color_intensity, 255 - color_intensity) # 红到粉色渐变
|
|||
|
|
|
|||
|
|
# 在BEV图像上绘制检测点
|
|||
|
|
if bev_x < bev_image.shape[1] and bev_y < bev_image.shape[0]:
|
|||
|
|
cv2.circle(bev_image, (bev_x, bev_y), 4, color, -1) if cv2 else None
|
|||
|
|
|
|||
|
|
# 添加文本标签
|
|||
|
|
conf_text = f'{score:.2f}'
|
|||
|
|
ax.scatter(bev_x, bev_y, c=[color], s=50, alpha=alpha, edgecolors='black', linewidth=1)
|
|||
|
|
ax.text(bev_x + 5, bev_y - 5, conf_text, fontsize=8,
|
|||
|
|
bbox=dict(boxstyle="round,pad=0.1", facecolor="white", alpha=0.8))
|
|||
|
|
|
|||
|
|
detection_count += 1
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"处理检测框 {i} 时出错: {e}")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
ax.imshow(bev_image)
|
|||
|
|
ax.set_title(f'3D检测结果 - 样本 {sample_idx} ({detection_count}/{len(boxes_3d)} 个检测框显示)', fontsize=14)
|
|||
|
|
else:
|
|||
|
|
ax.imshow(bev_image)
|
|||
|
|
ax.set_title(f'3D检测结果 - 样本 {sample_idx} (0 个检测框)', fontsize=14)
|
|||
|
|
|
|||
|
|
# 设置坐标轴
|
|||
|
|
ax.set_xlabel('BEV X (pixels)')
|
|||
|
|
ax.set_ylabel('BEV Y (pixels)')
|
|||
|
|
ax.grid(True, alpha=0.3)
|
|||
|
|
|
|||
|
|
# 添加坐标范围信息
|
|||
|
|
ax.text(10, 20, f'范围: [-{bev_range}m, +{bev_range}m]', fontsize=10,
|
|||
|
|
bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8))
|
|||
|
|
|
|||
|
|
plt.tight_layout()
|
|||
|
|
if save_path:
|
|||
|
|
plt.savefig(save_path, dpi=300, bbox_inches='tight')
|
|||
|
|
print(f"3D检测可视化已保存: {save_path}")
|
|||
|
|
plt.show()
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
parser = argparse.ArgumentParser(description='可视化单batch推理结果')
|
|||
|
|
parser.add_argument('--result_file', type=str,
|
|||
|
|
default='/data/infer_test/20251120_124755/one_batch_results.pkl',
|
|||
|
|
help='推理结果文件路径')
|
|||
|
|
parser.add_argument('--sample_idx', type=int, default=0,
|
|||
|
|
help='要可视化的样本索引')
|
|||
|
|
parser.add_argument('--output_dir', type=str, default='./visualization_output',
|
|||
|
|
help='可视化结果保存目录')
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
# 创建输出目录
|
|||
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|||
|
|
|
|||
|
|
# 加载推理结果
|
|||
|
|
print(f"加载推理结果: {args.result_file}")
|
|||
|
|
with open(args.result_file, 'rb') as f:
|
|||
|
|
results = pickle.load(f)
|
|||
|
|
|
|||
|
|
print(f"结果包含 {len(results)} 个样本")
|
|||
|
|
print(f"样本 {args.sample_idx} 的keys: {list(results[args.sample_idx].keys())}")
|
|||
|
|
|
|||
|
|
sample = results[args.sample_idx]
|
|||
|
|
|
|||
|
|
# 1. 可视化BEV分割
|
|||
|
|
if 'masks_bev' in sample:
|
|||
|
|
print(f"BEV分割形状: {sample['masks_bev'].shape}")
|
|||
|
|
pred_masks = sample['masks_bev']
|
|||
|
|
gt_masks = sample.get('gt_masks_bev', None)
|
|||
|
|
|
|||
|
|
seg_save_path = os.path.join(args.output_dir, f'bev_segmentation_sample_{args.sample_idx}.png')
|
|||
|
|
visualize_bev_segmentation(pred_masks, gt_masks, args.sample_idx, seg_save_path)
|
|||
|
|
|
|||
|
|
# 2. 可视化3D检测
|
|||
|
|
if 'boxes_3d' in sample and 'scores_3d' in sample and 'labels_3d' in sample:
|
|||
|
|
boxes_3d = sample['boxes_3d']
|
|||
|
|
scores_3d = sample['scores_3d']
|
|||
|
|
labels_3d = sample['labels_3d']
|
|||
|
|
|
|||
|
|
print(f"3D检测框数量: {len(boxes_3d)}")
|
|||
|
|
if len(scores_3d) > 0:
|
|||
|
|
print(f"检测置信度范围: {scores_3d.min():.3f} - {scores_3d.max():.3f}")
|
|||
|
|
|
|||
|
|
det_save_path = os.path.join(args.output_dir, f'3d_detection_sample_{args.sample_idx}.png')
|
|||
|
|
visualize_3d_detection(boxes_3d, scores_3d, labels_3d, args.sample_idx, det_save_path)
|
|||
|
|
|
|||
|
|
# 3. 打印统计信息
|
|||
|
|
print("\n" + "="*50)
|
|||
|
|
print(f"样本 {args.sample_idx} 统计信息:")
|
|||
|
|
print("="*50)
|
|||
|
|
|
|||
|
|
if 'masks_bev' in sample:
|
|||
|
|
pred_masks = sample['masks_bev']
|
|||
|
|
for i, class_name in enumerate(MAP_CLASSES):
|
|||
|
|
mask = pred_masks[i]
|
|||
|
|
pixel_count = (mask > 0.5).sum().item()
|
|||
|
|
percentage = pixel_count / (mask.shape[0] * mask.shape[1]) * 100
|
|||
|
|
print(".2f")
|
|||
|
|
|
|||
|
|
if 'boxes_3d' in sample:
|
|||
|
|
print(f"3D检测框数量: {len(sample['boxes_3d'])}")
|
|||
|
|
|
|||
|
|
print(f"\n可视化结果已保存到: {args.output_dir}")
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
main()
|