207 lines
7.1 KiB
Python
207 lines
7.1 KiB
Python
|
#!/usr/bin/env python
|
|||
|
|
"""
|
|||
|
|
基于推理结果进行可视化 - 参考visualize_single.py
|
|||
|
|
"""
|
|||
|
|
import argparse
|
|||
|
|
import copy
|
|||
|
|
import os
|
|||
|
|
import pickle
|
|||
|
|
|
|||
|
|
import mmcv
|
|||
|
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
from mmcv import Config
|
|||
|
|
from mmcv.runner import load_checkpoint
|
|||
|
|
from torchpack.utils.config import configs
|
|||
|
|
|
|||
|
|
from mmdet3d.core import LiDARInstance3DBoxes
|
|||
|
|
from mmdet3d.core.utils import visualize_camera, visualize_lidar, visualize_map
|
|||
|
|
from mmdet3d.datasets import build_dataloader, build_dataset
|
|||
|
|
|
|||
|
|
|
|||
|
|
def recursive_eval(obj, globals=None):
|
|||
|
|
if globals is None:
|
|||
|
|
globals = copy.deepcopy(obj)
|
|||
|
|
|
|||
|
|
if isinstance(obj, dict):
|
|||
|
|
for key in obj:
|
|||
|
|
obj[key] = recursive_eval(obj[key], globals)
|
|||
|
|
elif isinstance(obj, list):
|
|||
|
|
for k, val in enumerate(obj):
|
|||
|
|
obj[k] = recursive_eval(val, globals)
|
|||
|
|
elif isinstance(obj, str) and obj.startswith("${") and obj.endswith("}"):
|
|||
|
|
obj = eval(obj[2:-1], globals)
|
|||
|
|
obj = recursive_eval(obj, globals)
|
|||
|
|
|
|||
|
|
return obj
|
|||
|
|
|
|||
|
|
|
|||
|
|
def load_inference_results(result_file):
|
|||
|
|
"""加载推理结果"""
|
|||
|
|
print(f"加载推理结果: {result_file}")
|
|||
|
|
with open(result_file, 'rb') as f:
|
|||
|
|
results = pickle.load(f)
|
|||
|
|
print(f"结果包含 {len(results)} 个样本")
|
|||
|
|
return results
|
|||
|
|
|
|||
|
|
|
|||
|
|
def visualize_sample(sample_idx, sample, cfg, metas=None, out_dir="viz_inference"):
|
|||
|
|
"""可视化单个样本"""
|
|||
|
|
os.makedirs(out_dir, exist_ok=True)
|
|||
|
|
|
|||
|
|
print(f"\n--- 可视化样本 {sample_idx} ---")
|
|||
|
|
|
|||
|
|
# 准备检测结果
|
|||
|
|
bboxes = None
|
|||
|
|
labels = None
|
|||
|
|
scores = None
|
|||
|
|
|
|||
|
|
if 'boxes_3d' in sample and len(sample['boxes_3d']) > 0:
|
|||
|
|
# 3D检测框
|
|||
|
|
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(".3f")
|
|||
|
|
|
|||
|
|
# 转换为LiDARInstance3DBoxes格式
|
|||
|
|
if torch.is_tensor(boxes_3d):
|
|||
|
|
boxes_3d = boxes_3d.cpu().numpy()
|
|||
|
|
if torch.is_tensor(scores_3d):
|
|||
|
|
scores_3d = scores_3d.cpu().numpy()
|
|||
|
|
if torch.is_tensor(labels_3d):
|
|||
|
|
labels_3d = labels_3d.cpu().numpy()
|
|||
|
|
|
|||
|
|
# 过滤低置信度检测
|
|||
|
|
if len(scores_3d) > 0:
|
|||
|
|
keep = scores_3d >= 0.1 # 置信度阈值
|
|||
|
|
boxes_3d = boxes_3d[keep]
|
|||
|
|
scores_3d = scores_3d[keep]
|
|||
|
|
labels_3d = labels_3d[keep]
|
|||
|
|
|
|||
|
|
if len(boxes_3d) > 0:
|
|||
|
|
# 调整z坐标 (底面到中心)
|
|||
|
|
boxes_3d = boxes_3d.copy()
|
|||
|
|
boxes_3d[:, 2] -= boxes_3d[:, 5] / 2
|
|||
|
|
bboxes = LiDARInstance3DBoxes(boxes_3d, box_dim=9)
|
|||
|
|
scores = scores_3d
|
|||
|
|
labels = labels_3d.astype(np.int64)
|
|||
|
|
print(f"过滤后检测框数量: {len(boxes_3d)}")
|
|||
|
|
|
|||
|
|
# 准备分割结果
|
|||
|
|
masks = None
|
|||
|
|
if 'masks_bev' in sample:
|
|||
|
|
masks_bev = sample['masks_bev']
|
|||
|
|
print(f"BEV分割形状: {masks_bev.shape}")
|
|||
|
|
|
|||
|
|
if torch.is_tensor(masks_bev):
|
|||
|
|
masks_bev = masks_bev.cpu().numpy()
|
|||
|
|
|
|||
|
|
# 二值化分割结果 (阈值0.5)
|
|||
|
|
masks = masks_bev >= 0.5
|
|||
|
|
masks = masks.astype(np.bool)
|
|||
|
|
|
|||
|
|
# 创建虚拟的metas信息（用于可视化）
|
|||
|
|
if metas is None:
|
|||
|
|
metas = {
|
|||
|
|
"timestamp": f"sample_{sample_idx}",
|
|||
|
|
"token": f"token_{sample_idx}",
|
|||
|
|
"lidar2image": [np.eye(4)], # 单位矩阵
|
|||
|
|
"filename": ["dummy_image.jpg"], # 虚拟图像路径
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 可视化分割结果 (BEV地图)
|
|||
|
|
if masks is not None:
|
|||
|
|
map_out_dir = os.path.join(out_dir, "map")
|
|||
|
|
os.makedirs(map_out_dir, exist_ok=True)
|
|||
|
|
|
|||
|
|
map_path = os.path.join(map_out_dir, f"sample_{sample_idx:03d}.png")
|
|||
|
|
visualize_map(
|
|||
|
|
map_path,
|
|||
|
|
masks,
|
|||
|
|
classes=cfg.map_classes if hasattr(cfg, 'map_classes') else
|
|||
|
|
['drivable_area', 'ped_crossing', 'walkway', 'stop_line', 'carpark_area', 'divider']
|
|||
|
|
)
|
|||
|
|
print(f"BEV分割可视化已保存: {map_path}")
|
|||
|
|
|
|||
|
|
# 打印分割统计
|
|||
|
|
class_names = ['drivable_area', 'ped_crossing', 'walkway', 'stop_line', 'carpark_area', 'divider']
|
|||
|
|
print("BEV分割统计:")
|
|||
|
|
for i, name in enumerate(class_names[:len(masks)]):
|
|||
|
|
pixel_count = masks[i].sum()
|
|||
|
|
percentage = pixel_count / masks[i].size * 100
|
|||
|
|
print("15s")
|
|||
|
|
|
|||
|
|
# 创建虚拟点云数据用于LiDAR可视化
|
|||
|
|
if bboxes is not None:
|
|||
|
|
# 创建虚拟的LiDAR点云 (只用于可视化)
|
|||
|
|
# BEV范围: [-50, 50] x [-50, 50] x [-5, 3]
|
|||
|
|
num_points = 10000
|
|||
|
|
points = np.random.rand(num_points, 5)
|
|||
|
|
points[:, 0] = (points[:, 0] - 0.5) * 100 # x: -50 to 50
|
|||
|
|
points[:, 1] = (points[:, 1] - 0.5) * 100 # y: -50 to 50
|
|||
|
|
points[:, 2] = points[:, 2] * 8 - 5 # z: -5 to 3
|
|||
|
|
points[:, 3] = np.ones(num_points) # intensity
|
|||
|
|
points[:, 4] = np.zeros(num_points) # ring (unused)
|
|||
|
|
|
|||
|
|
lidar_out_dir = os.path.join(out_dir, "lidar")
|
|||
|
|
os.makedirs(lidar_out_dir, exist_ok=True)
|
|||
|
|
|
|||
|
|
lidar_path = os.path.join(lidar_out_dir, f"sample_{sample_idx:03d}.png")
|
|||
|
|
visualize_lidar(
|
|||
|
|
lidar_path,
|
|||
|
|
points,
|
|||
|
|
bboxes=bboxes,
|
|||
|
|
labels=labels,
|
|||
|
|
xlim=[-50, 50], # BEV范围
|
|||
|
|
ylim=[-50, 50],
|
|||
|
|
classes=cfg.object_classes if hasattr(cfg, 'object_classes') else
|
|||
|
|
['car', 'truck', 'construction_vehicle', 'bus', 'trailer',
|
|||
|
|
'motorcycle', 'bicycle', 'pedestrian', 'traffic_cone']
|
|||
|
|
)
|
|||
|
|
print(f"LiDAR检测可视化已保存: {lidar_path}")
|
|||
|
|
|
|||
|
|
print(f"样本 {sample_idx} 可视化完成")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
parser = argparse.ArgumentParser(description='可视化推理结果')
|
|||
|
|
parser.add_argument('--result-file', type=str,
|
|||
|
|
default='/data/infer_test/20251120_124755/one_batch_results.pkl',
|
|||
|
|
help='推理结果文件路径')
|
|||
|
|
parser.add_argument('--config', type=str,
|
|||
|
|
default='configs/nuscenes/det/transfusion/secfpn/camera+lidar/swint_v0p075/multitask_BEV2X_phase4b_rmtppad_segmentation.yaml',
|
|||
|
|
help='配置文件路径')
|
|||
|
|
parser.add_argument('--sample-idx', type=int, default=0,
|
|||
|
|
help='要可视化的样本索引')
|
|||
|
|
parser.add_argument('--out-dir', type=str, default='viz_inference_results',
|
|||
|
|
help='可视化结果保存目录')
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
# 加载配置
|
|||
|
|
configs.load(args.config, recursive=True)
|
|||
|
|
cfg = Config(recursive_eval(configs), filename=args.config)
|
|||
|
|
|
|||
|
|
# 加载推理结果
|
|||
|
|
results = load_inference_results(args.result_file)
|
|||
|
|
|
|||
|
|
if args.sample_idx >= len(results):
|
|||
|
|
print(f"样本索引 {args.sample_idx} 超出范围 (总共 {len(results)} 个样本)")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 可视化指定样本
|
|||
|
|
sample = results[args.sample_idx]
|
|||
|
|
visualize_sample(args.sample_idx, sample, cfg, out_dir=args.out_dir)
|
|||
|
|
|
|||
|
|
print("\n✅ 可视化完成!")
|
|||
|
|
print(f"结果保存在: {args.out_dir}")
|
|||
|
|
print("\n包含的文件:")
|
|||
|
|
print("- map/sample_XXX.png: BEV分割结果")
|
|||
|
|
print("- lidar/sample_XXX.png: LiDAR检测结果")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
main()
|