145 lines
4.4 KiB
Python
145 lines
4.4 KiB
Python
#!/usr/bin/env python
|
||
"""
|
||
只评估BEV分割的脚本，避免3D检测为空的问题
|
||
"""
|
||
import os
|
||
import sys
|
||
import torch
|
||
|
||
# 设置环境
|
||
os.environ['PATH'] = '/opt/conda/bin:' + os.environ.get('PATH', '')
|
||
os.environ['LD_LIBRARY_PATH'] = '/opt/conda/lib/python3.8/site-packages/torch/lib:/opt/conda/lib:/usr/local/cuda/lib64:' + os.environ.get('LD_LIBRARY_PATH', '')
|
||
os.environ['PYTHONPATH'] = '/workspace/bevfusion:' + os.environ.get('PYTHONPATH', '')
|
||
|
||
def main():
|
||
import mmcv
|
||
from mmcv import Config
|
||
from mmdet3d.datasets import build_dataloader, build_dataset
|
||
from mmdet3d.models import build_model
|
||
from mmcv.runner import load_checkpoint
|
||
from mmdet3d.apis import single_gpu_test
|
||
|
||
# 加载配置
|
||
config_file = 'configs/nuscenes/det/transfusion/secfpn/camera+lidar/swint_v0p075/multitask_BEV2X_phase4b_rmtppad_segmentation.yaml'
|
||
cfg = Config.fromfile(config_file)
|
||
|
||
# 设置测试模式
|
||
cfg.data.val.test_mode = True
|
||
|
||
# 构建数据集
|
||
print("构建数据集...")
|
||
from mmdet3d.datasets import NuScenesDataset
|
||
|
||
val_config = {
|
||
'type': 'NuScenesDataset',
|
||
'dataset_root': cfg.dataset_root,
|
||
'ann_file': cfg.dataset_root + 'nuscenes_infos_val.pkl',
|
||
'pipeline': cfg.evaluation.pipeline,
|
||
'object_classes': cfg.object_classes,
|
||
'map_classes': cfg.map_classes,
|
||
'modality': {
|
||
'use_lidar': True,
|
||
'use_camera': True,
|
||
'use_radar': False,
|
||
'use_map': False,
|
||
'use_external': False
|
||
},
|
||
'test_mode': True,
|
||
'use_valid_flag': False,
|
||
'box_type_3d': 'LiDAR',
|
||
'load_interval': cfg.data.val.load_interval if hasattr(cfg.data.val, 'load_interval') else 1
|
||
}
|
||
|
||
dataset = build_dataset(val_config)
|
||
print(f"数据集大小: {len(dataset)}")
|
||
|
||
# 构建数据加载器
|
||
data_loader = build_dataloader(
|
||
dataset,
|
||
samples_per_gpu=1,
|
||
workers_per_gpu=0,
|
||
dist=False,
|
||
shuffle=False,
|
||
)
|
||
|
||
# 构建模型
|
||
print("构建模型...")
|
||
cfg.model.train_cfg = None
|
||
model = build_model(cfg.model, test_cfg=cfg.get("test_cfg"))
|
||
print(f"模型heads: {list(model.heads.keys())}")
|
||
|
||
# 加载checkpoint
|
||
checkpoint_path = 'runs/run-4c8ec7e5-fabdc997/epoch_1.pth'
|
||
print(f"加载checkpoint: {checkpoint_path}")
|
||
checkpoint = load_checkpoint(model, checkpoint_path, map_location="cpu")
|
||
|
||
if "CLASSES" in checkpoint.get("meta", {}):
|
||
model.CLASSES = checkpoint["meta"]["CLASSES"]
|
||
else:
|
||
model.CLASSES = dataset.CLASSES
|
||
|
||
# 单GPU测试
|
||
print("开始BEV分割评估...")
|
||
model = model.cuda()
|
||
outputs = single_gpu_test(model, data_loader)
|
||
|
||
print(f"推理完成，结果数量: {len(outputs)}")
|
||
|
||
# 保存结果
|
||
output_file = '/data/eval_fast/bev_only_results.pkl'
|
||
mmcv.dump(outputs, output_file)
|
||
print(f"结果已保存至: {output_file}")
|
||
|
||
# 手动评估BEV分割
|
||
print("开始BEV分割评估...")
|
||
|
||
# 过滤只包含分割结果的输出
|
||
seg_results = []
|
||
for sample in outputs:
|
||
if 'masks_bev' in sample and 'gt_masks_bev' in sample:
|
||
seg_results.append({
|
||
'masks_bev': sample['masks_bev'],
|
||
'gt_masks_bev': sample['gt_masks_bev']
|
||
})
|
||
|
||
print(f"找到 {len(seg_results)} 个有效的分割样本")
|
||
|
||
if seg_results:
|
||
# 计算分割指标
|
||
total_iou = 0
|
||
valid_samples = 0
|
||
|
||
for sample in seg_results:
|
||
pred_mask = sample['masks_bev']
|
||
gt_mask = sample['gt_masks_bev']
|
||
|
||
if hasattr(pred_mask, 'cpu'):
|
||
pred_mask = pred_mask.cpu().numpy()
|
||
if hasattr(gt_mask, 'cpu'):
|
||
gt_mask = gt_mask.cpu().numpy()
|
||
|
||
# 计算mIoU (简化版，只计算第一个类别)
|
||
pred_flat = pred_mask[0].flatten() # 假设第一个类别是主要类别
|
||
gt_flat = gt_mask[0].flatten()
|
||
|
||
# 简单IoU计算
|
||
intersection = ((pred_flat > 0.5) & (gt_flat > 0.5)).sum()
|
||
union = ((pred_flat > 0.5) | (gt_flat > 0.5)).sum()
|
||
|
||
if union > 0:
|
||
iou = intersection / union
|
||
total_iou += iou
|
||
valid_samples += 1
|
||
|
||
if valid_samples > 0:
|
||
avg_iou = total_iou / valid_samples
|
||
print(".4f")
|
||
else:
|
||
print("没有有效的分割样本进行评估")
|
||
|
||
print("BEV分割评估完成！")
|
||
|
||
if __name__ == '__main__':
|
||
main()
|
||
|