#!/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()