227 lines
7.2 KiB
Python
227 lines
7.2 KiB
Python
#!/usr/bin/env python
|
||
"""
|
||
简单BEV分割可视化 - 不依赖mmdet3d库
|
||
"""
|
||
import os
|
||
import pickle
|
||
import numpy as np
|
||
import matplotlib.pyplot as plt
|
||
|
||
def load_results_safely(filepath):
|
||
"""安全加载结果，避免CUDA库问题"""
|
||
# 临时禁用CUDA相关的导入
|
||
import sys
|
||
original_modules = sys.modules.copy()
|
||
|
||
try:
|
||
# 阻止可能导致CUDA导入的模块
|
||
blocked_modules = [
|
||
'mmdet3d', 'torch', 'torchvision', 'torchaudio',
|
||
'mmcv', 'mmdet', 'torchpack'
|
||
]
|
||
for mod in blocked_modules:
|
||
if mod in sys.modules:
|
||
del sys.modules[mod]
|
||
|
||
with open(filepath, 'rb') as f:
|
||
results = pickle.load(f)
|
||
return results, True
|
||
except Exception as e:
|
||
print(f"标准加载失败: {e}")
|
||
return None, False
|
||
finally:
|
||
# 恢复原始模块状态
|
||
sys.modules.update(original_modules)
|
||
|
||
def visualize_bev_masks(masks_bev, sample_idx=0, save_dir='bev_visualization'):
|
||
"""可视化BEV分割结果"""
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
|
||
# 类别名称
|
||
class_names = ['drivable_area', 'ped_crossing', 'walkway', 'stop_line', 'carpark_area', 'divider']
|
||
colors = ['green', 'yellow', 'orange', 'red', 'purple', 'black']
|
||
|
||
# 如果是torch tensor，转换为numpy
|
||
if hasattr(masks_bev, 'cpu'):
|
||
masks_bev = masks_bev.cpu().numpy()
|
||
|
||
fig, axes = plt.subplots(2, 3, figsize=(15, 10))
|
||
fig.suptitle(f'BEV分割结果 - 样本 {sample_idx}', fontsize=16)
|
||
|
||
for i in range(6):
|
||
row, col = i // 3, i % 3
|
||
ax = axes[row, col]
|
||
|
||
mask = masks_bev[i]
|
||
if hasattr(mask, 'cpu'):
|
||
mask = mask.cpu().numpy()
|
||
|
||
# 显示mask (0-1之间的概率)
|
||
im = ax.imshow(mask, cmap='viridis', vmin=0, vmax=1)
|
||
ax.set_title(f'{class_names[i]}\n激活像素: {(mask > 0.5).sum()}')
|
||
ax.axis('off')
|
||
|
||
# 计算统计信息
|
||
activated_pixels = (mask > 0.5).sum()
|
||
total_pixels = mask.size
|
||
percentage = activated_pixels / total_pixels * 100
|
||
print("15s")
|
||
|
||
# 添加颜色条
|
||
plt.tight_layout()
|
||
cbar_ax = fig.add_axes([0.92, 0.15, 0.02, 0.7])
|
||
fig.colorbar(im, cax=cbar_ax, label='预测概率')
|
||
|
||
# 保存图像
|
||
save_path = os.path.join(save_dir, f'bev_segmentation_sample_{sample_idx}.png')
|
||
plt.savefig(save_path, dpi=300, bbox_inches='tight')
|
||
print(f"\nBEV分割可视化已保存: {save_path}")
|
||
plt.show()
|
||
|
||
# 创建叠加视图
|
||
fig, ax = plt.subplots(1, 1, figsize=(10, 10))
|
||
ax.set_title(f'BEV分割叠加视图 - 样本 {sample_idx}')
|
||
|
||
# 创建彩色叠加图像
|
||
height, width = masks_bev[0].shape
|
||
overlay = np.zeros((height, width, 3), dtype=np.uint8)
|
||
|
||
for i in range(6):
|
||
mask = masks_bev[i]
|
||
if hasattr(mask, 'cpu'):
|
||
mask = mask.cpu().numpy()
|
||
|
||
# 只显示高置信度区域
|
||
binary_mask = mask > 0.5
|
||
|
||
# 为每个类别分配颜色
|
||
if class_names[i] == 'drivable_area':
|
||
color = [0, 255, 0] # 绿色
|
||
elif class_names[i] == 'ped_crossing':
|
||
color = [255, 255, 0] # 黄色
|
||
elif class_names[i] == 'walkway':
|
||
color = [255, 165, 0] # 橙色
|
||
elif class_names[i] == 'stop_line':
|
||
color = [255, 0, 0] # 红色
|
||
elif class_names[i] == 'carpark_area':
|
||
color = [128, 0, 128] # 紫色
|
||
else: # divider
|
||
color = [0, 0, 0] # 黑色
|
||
|
||
# 应用颜色 (只在mask区域)
|
||
for c in range(3):
|
||
overlay[binary_mask, c] = color[c]
|
||
|
||
ax.imshow(overlay)
|
||
ax.set_title(f'BEV分割叠加 (高置信度区域) - 样本 {sample_idx}')
|
||
ax.axis('off')
|
||
|
||
# 添加图例
|
||
legend_elements = []
|
||
for i, name in enumerate(class_names):
|
||
if class_names[i] == 'drivable_area':
|
||
color = [0, 255, 0]
|
||
elif class_names[i] == 'ped_crossing':
|
||
color = [255, 255, 0]
|
||
elif class_names[i] == 'walkway':
|
||
color = [255, 165, 0]
|
||
elif class_names[i] == 'stop_line':
|
||
color = [255, 0, 0]
|
||
elif class_names[i] == 'carpark_area':
|
||
color = [128, 0, 128]
|
||
else:
|
||
color = [0, 0, 0]
|
||
|
||
import matplotlib.patches as mpatches
|
||
patch = mpatches.Patch(color=np.array(color)/255, label=name)
|
||
legend_elements.append(patch)
|
||
|
||
ax.legend(handles=legend_elements, bbox_to_anchor=(1.05, 1), loc='upper left')
|
||
|
||
overlay_path = os.path.join(save_dir, f'bev_overlay_sample_{sample_idx}.png')
|
||
plt.savefig(overlay_path, dpi=300, bbox_inches='tight')
|
||
print(f"BEV叠加可视化已保存: {overlay_path}")
|
||
plt.show()
|
||
|
||
def analyze_3d_detection(sample, sample_idx=0):
|
||
"""分析3D检测结果"""
|
||
print(f"\n=== 3D检测结果分析 - 样本 {sample_idx} ===")
|
||
|
||
if 'boxes_3d' in sample:
|
||
boxes_3d = sample['boxes_3d']
|
||
print(f"3D检测框数量: {len(boxes_3d)}")
|
||
|
||
if len(boxes_3d) > 0:
|
||
if hasattr(boxes_3d, 'cpu'):
|
||
boxes_3d = boxes_3d.cpu().numpy()
|
||
|
||
print(f"检测框形状: {boxes_3d.shape}")
|
||
|
||
# 分析检测框的范围
|
||
if boxes_3d.shape[1] >= 7: # x, y, z, w, l, h, rot
|
||
x_coords = boxes_3d[:, 0]
|
||
y_coords = boxes_3d[:, 1]
|
||
z_coords = boxes_3d[:, 2]
|
||
|
||
print(".2f")
|
||
print(".2f")
|
||
print(".2f")
|
||
|
||
if 'scores_3d' in sample:
|
||
scores_3d = sample['scores_3d']
|
||
if len(scores_3d) > 0:
|
||
if hasattr(scores_3d, 'cpu'):
|
||
scores_3d = scores_3d.cpu().numpy()
|
||
print(".3f")
|
||
|
||
if 'labels_3d' in sample:
|
||
labels_3d = sample['labels_3d']
|
||
if len(labels_3d) > 0:
|
||
if hasattr(labels_3d, 'cpu'):
|
||
labels_3d = labels_3d.cpu().numpy()
|
||
unique_labels, counts = np.unique(labels_3d, return_counts=True)
|
||
print("检测类别分布:")
|
||
for label, count in zip(unique_labels, counts):
|
||
print(f" 类别 {int(label)}: {count} 个")
|
||
|
||
def main():
|
||
# 推理结果文件
|
||
result_file = '/data/infer_test/20251120_124755/one_batch_results.pkl'
|
||
|
||
print("开始分析推理结果...")
|
||
print(f"结果文件: {result_file}")
|
||
|
||
# 加载结果
|
||
results, success = load_results_safely(result_file)
|
||
|
||
if not success or results is None:
|
||
print("无法加载结果文件")
|
||
return
|
||
|
||
print(f"成功加载! 包含 {len(results)} 个样本")
|
||
|
||
# 分析第一个样本
|
||
sample_idx = 0
|
||
if sample_idx < len(results):
|
||
sample = results[sample_idx]
|
||
print(f"\n样本 {sample_idx} 包含字段: {list(sample.keys())}")
|
||
|
||
# 可视化BEV分割
|
||
if 'masks_bev' in sample:
|
||
print("\n开始BEV分割可视化...")
|
||
visualize_bev_masks(sample['masks_bev'], sample_idx)
|
||
|
||
# 分析3D检测
|
||
analyze_3d_detection(sample, sample_idx)
|
||
|
||
# 分析Ground Truth (如果有)
|
||
if 'gt_masks_bev' in sample:
|
||
gt_masks = sample['gt_masks_bev']
|
||
print(f"\nGround Truth BEV分割形状: {gt_masks.shape}")
|
||
|
||
else:
|
||
print(f"样本索引 {sample_idx} 超出范围")
|
||
|
||
if __name__ == '__main__':
|
||
main()
|