256 lines
7.6 KiB
Python
256 lines
7.6 KiB
Python
|
"""
|
|||
|
|
将BEVFusion实时输出转换为CARLA可用格式
|
|||
|
|
|
|||
|
|
用于在CARLA仿真中实时显示BEVFusion的检测和地图结果
|
|||
|
|
|
|||
|
|
使用方法:
|
|||
|
|
python tools/convert_bevfusion_output_to_carla.py \
|
|||
|
|
--bevfusion-output results/predictions.pkl \
|
|||
|
|
--output carla_visualization.json
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import pickle
|
|||
|
|
import argparse
|
|||
|
|
import json
|
|||
|
|
import numpy as np
|
|||
|
|
|
|||
|
|
|
|||
|
|
class BEVFusionToCarlaConverter:
|
|||
|
|
"""BEVFusion输出转CARLA可视化格式"""
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
def convert_detection_results(self, bevfusion_output):
|
|||
|
|
"""
|
|||
|
|
转换3D检测结果
|
|||
|
|
|
|||
|
|
BEVFusion格式:
|
|||
|
|
{
|
|||
|
|
'boxes_3d': (N, 9), # x,y,z,dx,dy,dz,yaw,vx,vy
|
|||
|
|
'scores_3d': (N,),
|
|||
|
|
'labels_3d': (N,)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
CARLA格式:
|
|||
|
|
[
|
|||
|
|
{
|
|||
|
|
'type': 'vehicle',
|
|||
|
|
'location': {'x': ..., 'y': ..., 'z': ...},
|
|||
|
|
'extent': {'x': ..., 'y': ..., 'z': ...},
|
|||
|
|
'rotation': {'pitch': 0, 'yaw': ..., 'roll': 0},
|
|||
|
|
'velocity': {'x': ..., 'y': ...},
|
|||
|
|
'confidence': 0.95
|
|||
|
|
},
|
|||
|
|
...
|
|||
|
|
]
|
|||
|
|
"""
|
|||
|
|
carla_objects = []
|
|||
|
|
|
|||
|
|
boxes = bevfusion_output['boxes_3d']
|
|||
|
|
scores = bevfusion_output['scores_3d']
|
|||
|
|
labels = bevfusion_output['labels_3d']
|
|||
|
|
|
|||
|
|
class_names = ['car', 'truck', 'construction_vehicle', 'bus', 'trailer',
|
|||
|
|
'barrier', 'motorcycle', 'bicycle', 'pedestrian', 'traffic_cone']
|
|||
|
|
|
|||
|
|
for i in range(len(boxes)):
|
|||
|
|
x, y, z, dx, dy, dz, yaw, vx, vy = boxes[i]
|
|||
|
|
|
|||
|
|
carla_obj = {
|
|||
|
|
'type': class_names[labels[i]],
|
|||
|
|
'location': {
|
|||
|
|
'x': float(x),
|
|||
|
|
'y': float(y),
|
|||
|
|
'z': float(z)
|
|||
|
|
},
|
|||
|
|
'extent': {
|
|||
|
|
'x': float(dx / 2),
|
|||
|
|
'y': float(dy / 2),
|
|||
|
|
'z': float(dz / 2)
|
|||
|
|
},
|
|||
|
|
'rotation': {
|
|||
|
|
'pitch': 0.0,
|
|||
|
|
'yaw': float(np.degrees(yaw)),
|
|||
|
|
'roll': 0.0
|
|||
|
|
},
|
|||
|
|
'velocity': {
|
|||
|
|
'x': float(vx),
|
|||
|
|
'y': float(vy)
|
|||
|
|
},
|
|||
|
|
'confidence': float(scores[i])
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
carla_objects.append(carla_obj)
|
|||
|
|
|
|||
|
|
return carla_objects
|
|||
|
|
|
|||
|
|
def convert_segmentation_results(self, bevfusion_output):
|
|||
|
|
"""
|
|||
|
|
转换BEV分割结果
|
|||
|
|
|
|||
|
|
BEVFusion格式:
|
|||
|
|
{
|
|||
|
|
'masks_bev': (H, W, C) # 每个类别的mask
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
CARLA格式:
|
|||
|
|
{
|
|||
|
|
'drivable_area': [[x,y], ...], # 多边形顶点
|
|||
|
|
'lanes': [[[x,y], ...], ...], # 车道线列表
|
|||
|
|
...
|
|||
|
|
}
|
|||
|
|
"""
|
|||
|
|
masks = bevfusion_output['masks_bev']
|
|||
|
|
|
|||
|
|
# 从mask提取轮廓
|
|||
|
|
import cv2
|
|||
|
|
|
|||
|
|
carla_segmentation = {}
|
|||
|
|
|
|||
|
|
class_names = ['drivable_area', 'ped_crossing', 'walkway',
|
|||
|
|
'stop_line', 'carpark_area', 'divider']
|
|||
|
|
|
|||
|
|
for i, class_name in enumerate(class_names):
|
|||
|
|
mask = masks[:, :, i]
|
|||
|
|
|
|||
|
|
# 提取轮廓
|
|||
|
|
contours, _ = cv2.findContours(
|
|||
|
|
mask.astype(np.uint8),
|
|||
|
|
cv2.RETR_EXTERNAL,
|
|||
|
|
cv2.CHAIN_APPROX_SIMPLE
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 转为坐标点列表
|
|||
|
|
polygons = []
|
|||
|
|
for contour in contours:
|
|||
|
|
if len(contour) > 10: # 过滤小轮廓
|
|||
|
|
# 从像素坐标转为米坐标
|
|||
|
|
points = contour.squeeze()
|
|||
|
|
points_meter = self.pixel_to_meter(points, masks.shape)
|
|||
|
|
polygons.append(points_meter.tolist())
|
|||
|
|
|
|||
|
|
carla_segmentation[class_name] = polygons
|
|||
|
|
|
|||
|
|
return carla_segmentation
|
|||
|
|
|
|||
|
|
def convert_vector_map_results(self, bevfusion_output):
|
|||
|
|
"""
|
|||
|
|
转换矢量地图结果(MapTR输出)
|
|||
|
|
|
|||
|
|
BEVFusion/MapTR格式:
|
|||
|
|
{
|
|||
|
|
'vectors': [
|
|||
|
|
{
|
|||
|
|
'class': 0, # divider
|
|||
|
|
'points': [[x,y], ...],
|
|||
|
|
'score': 0.95
|
|||
|
|
},
|
|||
|
|
...
|
|||
|
|
]
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
CARLA格式:
|
|||
|
|
{
|
|||
|
|
'lanes': [
|
|||
|
|
{
|
|||
|
|
'points': [[x,y,z], ...],
|
|||
|
|
'type': 'divider',
|
|||
|
|
'confidence': 0.95
|
|||
|
|
},
|
|||
|
|
...
|
|||
|
|
]
|
|||
|
|
}
|
|||
|
|
"""
|
|||
|
|
vectors = bevfusion_output.get('vectors', [])
|
|||
|
|
|
|||
|
|
carla_vectors = []
|
|||
|
|
|
|||
|
|
class_map = {0: 'divider', 1: 'boundary', 2: 'ped_crossing'}
|
|||
|
|
|
|||
|
|
for vec in vectors:
|
|||
|
|
points_2d = np.array(vec['points'])
|
|||
|
|
|
|||
|
|
# 添加z坐标（假设在地面）
|
|||
|
|
points_3d = np.column_stack([
|
|||
|
|
points_2d,
|
|||
|
|
np.zeros(len(points_2d))
|
|||
|
|
])
|
|||
|
|
|
|||
|
|
carla_vec = {
|
|||
|
|
'points': points_3d.tolist(),
|
|||
|
|
'type': class_map.get(vec['class'], 'unknown'),
|
|||
|
|
'confidence': vec.get('score', 1.0)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
carla_vectors.append(carla_vec)
|
|||
|
|
|
|||
|
|
return {'lanes': carla_vectors}
|
|||
|
|
|
|||
|
|
def pixel_to_meter(self, points_pixel, shape):
|
|||
|
|
"""像素坐标转米坐标"""
|
|||
|
|
H, W = shape[:2]
|
|||
|
|
|
|||
|
|
# 假设BEV范围 [-50, 50] x [-50, 50]
|
|||
|
|
points_meter = points_pixel.copy().astype(float)
|
|||
|
|
points_meter[:, 0] = (points_pixel[:, 0] / W) * 100 - 50
|
|||
|
|
points_meter[:, 1] = (points_pixel[:, 1] / H) * 100 - 50
|
|||
|
|
|
|||
|
|
return points_meter
|
|||
|
|
|
|||
|
|
def convert_full_output(self, bevfusion_output, output_path):
|
|||
|
|
"""转换完整的BEVFusion输出"""
|
|||
|
|
carla_format = {}
|
|||
|
|
|
|||
|
|
# 1. 检测结果
|
|||
|
|
if 'boxes_3d' in bevfusion_output:
|
|||
|
|
carla_format['objects'] = self.convert_detection_results(bevfusion_output)
|
|||
|
|
|
|||
|
|
# 2. 分割结果
|
|||
|
|
if 'masks_bev' in bevfusion_output:
|
|||
|
|
carla_format['segmentation'] = self.convert_segmentation_results(bevfusion_output)
|
|||
|
|
|
|||
|
|
# 3. 矢量地图
|
|||
|
|
if 'vectors' in bevfusion_output or 'vector_map' in bevfusion_output:
|
|||
|
|
vec_output = bevfusion_output.get('vector_map', bevfusion_output)
|
|||
|
|
carla_format['vector_map'] = self.convert_vector_map_results(vec_output)
|
|||
|
|
|
|||
|
|
# 保存JSON
|
|||
|
|
with open(output_path, 'w') as f:
|
|||
|
|
json.dump(carla_format, f, indent=2)
|
|||
|
|
|
|||
|
|
print(f"✅ CARLA格式文件已保存: {output_path}")
|
|||
|
|
|
|||
|
|
return carla_format
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
parser = argparse.ArgumentParser(description='BEVFusion输出转CARLA格式')
|
|||
|
|
parser.add_argument('--bevfusion-output', type=str, required=True,
|
|||
|
|
help='BEVFusion预测结果文件(.pkl)')
|
|||
|
|
parser.add_argument('--output', type=str, required=True,
|
|||
|
|
help='输出CARLA格式文件(.json)')
|
|||
|
|
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
# 加载BEVFusion输出
|
|||
|
|
print(f"加载BEVFusion输出: {args.bevfusion_output}")
|
|||
|
|
with open(args.bevfusion_output, 'rb') as f:
|
|||
|
|
bevfusion_output = pickle.load(f)
|
|||
|
|
|
|||
|
|
# 转换
|
|||
|
|
converter = BEVFusionToCarlaConverter()
|
|||
|
|
converter.convert_full_output(bevfusion_output, args.output)
|
|||
|
|
|
|||
|
|
print(f"\n在CARLA中使用:")
|
|||
|
|
print(f" import carla")
|
|||
|
|
print(f" # 加载JSON并可视化")
|
|||
|
|
print(f" with open('{args.output}') as f:")
|
|||
|
|
print(f" data = json.load(f)")
|
|||
|
|
print(f" # 在CARLA中绘制...")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
main()
|
|||
|
|
|