241 lines
8.1 KiB
Python
241 lines
8.1 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
"""
|
|||
|
|
BEVFusion 正确的可视化脚本
|
|||
|
|
使用正确的数据键名: masks_bev, boxes_3d
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import pickle
|
|||
|
|
import numpy as np
|
|||
|
|
import cv2
|
|||
|
|
from pathlib import Path
|
|||
|
|
import argparse
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
import matplotlib.pyplot as plt
|
|||
|
|
import matplotlib.patches as mpatches
|
|||
|
|
import torch
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 分割类别名称（nuScenes定义）
|
|||
|
|
SEGMENTATION_CLASSES = [
|
|||
|
|
'drivable_area', # 0
|
|||
|
|
'ped_crossing', # 1
|
|||
|
|
'walkway', # 2
|
|||
|
|
'stop_line', # 3
|
|||
|
|
'carpark_area', # 4
|
|||
|
|
'divider', # 5
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
# 分割类别颜色
|
|||
|
|
SEGMENTATION_COLORS = [
|
|||
|
|
[255, 195, 128], # drivable_area - 橙色
|
|||
|
|
[255, 99, 71], # ped_crossing - 番茄红
|
|||
|
|
[255, 140, 0], # walkway - 深橙色
|
|||
|
|
[180, 62, 39], # stop_line - 棕色
|
|||
|
|
[233, 150, 70], # carpark_area - 棕橙色
|
|||
|
|
[160, 82, 45], # divider - 褐色
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
|
|||
|
|
class BEVVisualizer:
|
|||
|
|
def __init__(self, result_file, frame_interval=1):
|
|||
|
|
"""初始化可视化器"""
|
|||
|
|
print(f"加载推理结果: {result_file}")
|
|||
|
|
with open(result_file, 'rb') as f:
|
|||
|
|
self.results = pickle.load(f)
|
|||
|
|
print(f"成功加载 {len(self.results)} 个样本")
|
|||
|
|
self.frame_interval = frame_interval
|
|||
|
|
|
|||
|
|
def visualize_single_bev(self, sample_idx, save_path=None, show_info=True):
|
|||
|
|
"""可视化单个样本的BEV结果"""
|
|||
|
|
result = self.results[sample_idx]
|
|||
|
|
|
|||
|
|
# 创建图像
|
|||
|
|
fig, axes = plt.subplots(1, 2, figsize=(20, 10))
|
|||
|
|
fig.suptitle(f'Sample {sample_idx}', fontsize=16, fontweight='bold')
|
|||
|
|
|
|||
|
|
# 1. 可视化分割（BEV）
|
|||
|
|
if 'masks_bev' in result:
|
|||
|
|
masks_bev = result['masks_bev']
|
|||
|
|
|
|||
|
|
# 转换为numpy
|
|||
|
|
if hasattr(masks_bev, 'cpu'):
|
|||
|
|
masks_bev = masks_bev.cpu().numpy()
|
|||
|
|
|
|||
|
|
# (C, H, W) -> (H, W)，取最大概率的类别
|
|||
|
|
seg_map = np.argmax(masks_bev, axis=0)
|
|||
|
|
|
|||
|
|
# 创建彩色可视化
|
|||
|
|
h, w = seg_map.shape
|
|||
|
|
vis_img = np.zeros((h, w, 3), dtype=np.uint8)
|
|||
|
|
|
|||
|
|
for class_id in range(len(SEGMENTATION_CLASSES)):
|
|||
|
|
mask = seg_map == class_id
|
|||
|
|
if np.any(mask):
|
|||
|
|
vis_img[mask] = SEGMENTATION_COLORS[class_id]
|
|||
|
|
|
|||
|
|
axes[0].imshow(vis_img)
|
|||
|
|
axes[0].set_title('BEV Segmentation Map', fontsize=14)
|
|||
|
|
axes[0].axis('off')
|
|||
|
|
|
|||
|
|
# 添加图例
|
|||
|
|
patches = []
|
|||
|
|
for i, cls_name in enumerate(SEGMENTATION_CLASSES):
|
|||
|
|
color = np.array(SEGMENTATION_COLORS[i]) / 255.0
|
|||
|
|
mask = seg_map == i
|
|||
|
|
count = np.sum(mask)
|
|||
|
|
if count > 0:
|
|||
|
|
patches.append(mpatches.Patch(color=color,
|
|||
|
|
label=f'{cls_name}: {count} px'))
|
|||
|
|
|
|||
|
|
if patches:
|
|||
|
|
axes[0].legend(handles=patches, loc='upper right',
|
|||
|
|
fontsize=10, framealpha=0.8)
|
|||
|
|
else:
|
|||
|
|
axes[0].text(0.5, 0.5, 'No Segmentation',
|
|||
|
|
ha='center', va='center', fontsize=20)
|
|||
|
|
axes[0].axis('off')
|
|||
|
|
|
|||
|
|
# 2. 可视化GT分割
|
|||
|
|
if 'gt_masks_bev' in result:
|
|||
|
|
gt_masks = result['gt_masks_bev']
|
|||
|
|
|
|||
|
|
# 转换为numpy
|
|||
|
|
if hasattr(gt_masks, 'cpu'):
|
|||
|
|
gt_masks = gt_masks.cpu().numpy()
|
|||
|
|
|
|||
|
|
# 创建彩色可视化
|
|||
|
|
h, w = gt_masks.shape[1], gt_masks.shape[2] if len(gt_masks.shape) == 3 else gt_masks.shape
|
|||
|
|
if len(gt_masks.shape) == 3:
|
|||
|
|
# (C, H, W)
|
|||
|
|
gt_map = np.argmax(gt_masks, axis=0)
|
|||
|
|
else:
|
|||
|
|
gt_map = gt_masks
|
|||
|
|
|
|||
|
|
vis_img_gt = np.zeros((h, w, 3), dtype=np.uint8)
|
|||
|
|
for class_id in range(min(len(SEGMENTATION_CLASSES), gt_map.max() + 1)):
|
|||
|
|
mask = gt_map == class_id
|
|||
|
|
if np.any(mask):
|
|||
|
|
vis_img_gt[mask] = SEGMENTATION_COLORS[class_id]
|
|||
|
|
|
|||
|
|
axes[1].imshow(vis_img_gt)
|
|||
|
|
axes[1].set_title('Ground Truth Segmentation', fontsize=14)
|
|||
|
|
axes[1].axis('off')
|
|||
|
|
else:
|
|||
|
|
axes[1].text(0.5, 0.5, 'No GT Data',
|
|||
|
|
ha='center', va='center', fontsize=20)
|
|||
|
|
axes[1].axis('off')
|
|||
|
|
|
|||
|
|
plt.tight_layout()
|
|||
|
|
|
|||
|
|
if save_path:
|
|||
|
|
plt.savefig(save_path, dpi=100, bbox_inches='tight')
|
|||
|
|
plt.close()
|
|||
|
|
if show_info:
|
|||
|
|
print(f"保存: {save_path}")
|
|||
|
|
else:
|
|||
|
|
plt.show()
|
|||
|
|
|
|||
|
|
def generate_video(self, output_path, fps=10):
|
|||
|
|
"""生成视频"""
|
|||
|
|
print(f"开始生成视频: {output_path}")
|
|||
|
|
|
|||
|
|
# 计算采样索引
|
|||
|
|
sample_indices = list(range(0, len(self.results), self.frame_interval))
|
|||
|
|
print(f"总样本数: {len(self.results)}")
|
|||
|
|
print(f"帧间隔: {self.frame_interval}")
|
|||
|
|
print(f"将生成: {len(sample_indices)} 帧")
|
|||
|
|
print()
|
|||
|
|
|
|||
|
|
# 创建临时目录
|
|||
|
|
temp_dir = Path(output_path).parent / 'temp_frames'
|
|||
|
|
temp_dir.mkdir(exist_ok=True)
|
|||
|
|
|
|||
|
|
# 生成所有帧
|
|||
|
|
frames = []
|
|||
|
|
for idx in tqdm(sample_indices, desc="生成帧"):
|
|||
|
|
frame_path = temp_dir / f'frame_{idx:05d}.png'
|
|||
|
|
self.visualize_single_bev(idx, save_path=frame_path, show_info=False)
|
|||
|
|
frames.append(str(frame_path))
|
|||
|
|
|
|||
|
|
print()
|
|||
|
|
print(f"生成了 {len(frames)} 帧，开始合成视频...")
|
|||
|
|
|
|||
|
|
# 读取第一帧确定尺寸
|
|||
|
|
first_frame = cv2.imread(frames[0])
|
|||
|
|
h, w = first_frame.shape[:2]
|
|||
|
|
|
|||
|
|
# 创建视频写入器
|
|||
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
|||
|
|
video_writer = cv2.VideoWriter(output_path, fourcc, fps, (w, h))
|
|||
|
|
|
|||
|
|
# 写入所有帧
|
|||
|
|
for frame_path in tqdm(frames, desc="写入视频"):
|
|||
|
|
frame = cv2.imread(frame_path)
|
|||
|
|
video_writer.write(frame)
|
|||
|
|
|
|||
|
|
video_writer.release()
|
|||
|
|
print(f"✅ 视频已保存: {output_path}")
|
|||
|
|
|
|||
|
|
# 清理临时文件
|
|||
|
|
print("清理临时文件...")
|
|||
|
|
for frame_path in frames:
|
|||
|
|
os.remove(frame_path)
|
|||
|
|
temp_dir.rmdir()
|
|||
|
|
print("✅ 完成！")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
parser = argparse.ArgumentParser(description='BEVFusion结果可视化（修复版）')
|
|||
|
|
parser.add_argument('--result-file', type=str,
|
|||
|
|
default='results_epoch19.pkl',
|
|||
|
|
help='推理结果文件')
|
|||
|
|
parser.add_argument('--output-dir', type=str,
|
|||
|
|
default='visualizations',
|
|||
|
|
help='输出目录')
|
|||
|
|
parser.add_argument('--mode', type=str,
|
|||
|
|
choices=['quick', 'video'],
|
|||
|
|
default='quick',
|
|||
|
|
help='可视化模式')
|
|||
|
|
parser.add_argument('--num-samples', type=int, default=10,
|
|||
|
|
help='快速可视化的样本数')
|
|||
|
|
parser.add_argument('--video-fps', type=int, default=10,
|
|||
|
|
help='视频帧率')
|
|||
|
|
parser.add_argument('--frame-interval', type=int, default=6,
|
|||
|
|
help='帧间隔（下采样）')
|
|||
|
|
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
# 初始化可视化器
|
|||
|
|
visualizer = BEVVisualizer(
|
|||
|
|
result_file=args.result_file,
|
|||
|
|
frame_interval=args.frame_interval
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|||
|
|
|
|||
|
|
# 执行可视化
|
|||
|
|
if args.mode == 'quick':
|
|||
|
|
print(f"快速可视化前 {args.num_samples} 个样本...")
|
|||
|
|
for idx in tqdm(range(args.num_samples)):
|
|||
|
|
save_path = os.path.join(args.output_dir, f'sample_{idx:04d}_bev.png')
|
|||
|
|
visualizer.visualize_single_bev(idx, save_path=save_path)
|
|||
|
|
print(f"✅ 完成！图像保存在: {args.output_dir}/")
|
|||
|
|
|
|||
|
|
elif args.mode == 'video':
|
|||
|
|
video_path = os.path.join(args.output_dir, 'bevfusion_results.mp4')
|
|||
|
|
visualizer.generate_video(
|
|||
|
|
output_path=video_path,
|
|||
|
|
fps=args.video_fps
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
main()
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|