176 lines
5.3 KiB
Python
Executable File
176 lines
5.3 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""
|
||
可视化nuScenes推理结果的简化脚本
|
||
直接使用pickle文件生成可视化
|
||
"""
|
||
|
||
import pickle
|
||
import numpy as np
|
||
import matplotlib.pyplot as plt
|
||
from pathlib import Path
|
||
import argparse
|
||
|
||
|
||
def visualize_detection_bev(results, idx=0, output_path='detection_bev.png'):
|
||
"""可视化3D检测结果（BEV视图）"""
|
||
fig, ax = plt.subplots(figsize=(12, 12))
|
||
|
||
result = results[idx]
|
||
|
||
if 'boxes_3d' in result:
|
||
boxes = result['boxes_3d']
|
||
scores = result['scores_3d']
|
||
labels = result['labels_3d']
|
||
|
||
# 类别名称
|
||
class_names = [
|
||
'car', 'truck', 'construction_vehicle', 'bus', 'trailer',
|
||
'barrier', 'motorcycle', 'bicycle', 'pedestrian', 'traffic_cone'
|
||
]
|
||
|
||
colors = plt.cm.tab10(np.linspace(0, 1, 10))
|
||
|
||
for box, score, label in zip(boxes, scores, labels):
|
||
if score < 0.3:
|
||
continue
|
||
|
||
# 提取中心和尺寸
|
||
center = box[:2]
|
||
size = box[3:5]
|
||
yaw = box[6]
|
||
|
||
# 绘制框
|
||
l, w = size[0], size[1]
|
||
corners = np.array([
|
||
[-l/2, -w/2], [l/2, -w/2],
|
||
[l/2, w/2], [-l/2, w/2], [-l/2, -w/2]
|
||
])
|
||
|
||
rot = np.array([
|
||
[np.cos(yaw), -np.sin(yaw)],
|
||
[np.sin(yaw), np.cos(yaw)]
|
||
])
|
||
corners = corners @ rot.T + center
|
||
|
||
color = colors[label % 10]
|
||
ax.plot(corners[:, 0], corners[:, 1], color=color, linewidth=2)
|
||
ax.scatter(center[0], center[1], color=color, s=50)
|
||
|
||
# 标签
|
||
class_name = class_names[label] if label < len(class_names) else f'cls_{label}'
|
||
ax.text(center[0], center[1], f'{class_name}\n{score:.2f}',
|
||
fontsize=8, color='white',
|
||
bbox=dict(boxstyle='round', facecolor=color, alpha=0.7))
|
||
|
||
ax.set_xlim(-50, 50)
|
||
ax.set_ylim(-50, 50)
|
||
ax.set_aspect('equal')
|
||
ax.grid(True, alpha=0.3)
|
||
ax.set_xlabel('X (meters)', fontsize=14)
|
||
ax.set_ylabel('Y (meters)', fontsize=14)
|
||
ax.set_title('3D Detection Results (BEV View)', fontsize=16)
|
||
|
||
plt.tight_layout()
|
||
plt.savefig(output_path, dpi=150)
|
||
plt.close()
|
||
|
||
print(f"✅ 已保存: {output_path}")
|
||
|
||
|
||
def visualize_segmentation(results, idx=0, output_path='segmentation.png'):
|
||
"""可视化BEV分割结果"""
|
||
fig, ax = plt.subplots(figsize=(10, 10))
|
||
|
||
result = results[idx]
|
||
|
||
if 'seg_pred' in result:
|
||
seg_pred = result['seg_pred'] # (C, H, W)
|
||
|
||
# 创建彩色图
|
||
h, w = seg_pred.shape[1:]
|
||
color_map = np.zeros((h, w, 3), dtype=np.uint8)
|
||
|
||
colors = {
|
||
0: [128, 64, 128], # drivable_area
|
||
1: [244, 35, 232], # ped_crossing
|
||
2: [70, 70, 70], # walkway
|
||
3: [220, 20, 60], # stop_line
|
||
4: [157, 234, 50], # carpark_area
|
||
5: [255, 255, 0], # divider
|
||
}
|
||
|
||
class_names = [
|
||
'drivable_area', 'ped_crossing', 'walkway',
|
||
'stop_line', 'carpark_area', 'divider'
|
||
]
|
||
|
||
for idx, name in enumerate(class_names):
|
||
if idx < seg_pred.shape[0]:
|
||
mask = seg_pred[idx] > 0.5
|
||
color_map[mask] = colors[idx]
|
||
|
||
ax.imshow(color_map)
|
||
ax.set_title('BEV Segmentation', fontsize=16)
|
||
ax.axis('off')
|
||
|
||
# 图例
|
||
from matplotlib.patches import Patch
|
||
legend_elements = [
|
||
Patch(facecolor=np.array(colors[i])/255, label=class_names[i])
|
||
for i in range(len(class_names))
|
||
]
|
||
ax.legend(handles=legend_elements, loc='upper right')
|
||
|
||
plt.tight_layout()
|
||
plt.savefig(output_path, dpi=150)
|
||
plt.close()
|
||
|
||
print(f"✅ 已保存: {output_path}")
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument('--results', default='results.pkl', help='结果pickle文件')
|
||
parser.add_argument('--output-dir', default='visualizations', help='输出目录')
|
||
parser.add_argument('--samples', type=int, default=5, help='可视化样本数')
|
||
args = parser.parse_args()
|
||
|
||
print(f"\n{'='*80}")
|
||
print("可视化nuScenes推理结果")
|
||
print(f"{'='*80}\n")
|
||
|
||
# 创建输出目录
|
||
Path(args.output_dir).mkdir(exist_ok=True)
|
||
|
||
# 加载结果
|
||
print(f"加载结果: {args.results}")
|
||
with open(args.results, 'rb') as f:
|
||
results = pickle.load(f)
|
||
|
||
print(f"结果数量: {len(results)}")
|
||
print(f"可视化前 {min(args.samples, len(results))} 个样本\n")
|
||
|
||
# 可视化
|
||
for i in range(min(args.samples, len(results))):
|
||
print(f"处理样本 {i+1}/{min(args.samples, len(results))}...")
|
||
|
||
det_path = Path(args.output_dir) / f"sample_{i:04d}_detection.png"
|
||
seg_path = Path(args.output_dir) / f"sample_{i:04d}_segmentation.png"
|
||
|
||
visualize_detection_bev(results, i, str(det_path))
|
||
visualize_segmentation(results, i, str(seg_path))
|
||
|
||
print(f"\n{'='*80}")
|
||
print(f"✅ 完成！结果保存在: {args.output_dir}/")
|
||
print(f"{'='*80}\n")
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|
||
|
||
|
||
|
||
|
||
|
||
|