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()
|
||
|