File size: 2,129 Bytes
7accb91 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
import cv2
import torch
import numpy as np
from scipy.spatial.transform import Rotation as SCR
import math
from navsim.common.dataclasses import AgentInput, EgoStatus, Cameras, Camera, Lidar
OPENCV2IMU = np.array([[0, 0, 1], [-1, 0, 0], [0, -1, 0]])
OPENCV2LIDAR = np.array([[1, 0, 0], [0, 0, 1], [0, -1, 0]])
nusc_cameras = ['CAM_FRONT', 'CAM_FRONT_RIGHT', 'CAM_FRONT_LEFT']
nuplan_cameras = ["cam_f0", "cam_r0", "cam_l0"]
def fov2focal(fov, pixels):
return pixels / (2 * math.tan(fov / 2))
def get_intrinsic(intrin_dict):
fovx, fovy = intrin_dict['fovx'], intrin_dict['fovy']
h, w = intrin_dict['H'], intrin_dict['W']
K = np.eye(4)
K[0, 0], K[1, 1] = fov2focal(fovx, w), fov2focal(fovy, h)
K[0, 2], K[1, 2] = intrin_dict['cx'], intrin_dict['cy']
return K
def parse_raw(raw_data):
obs, info = raw_data
imgs = {}
raw_imgs = {}
for cam in nusc_cameras:
im = obs['rgb'][cam]
resize_im = cv2.resize(im, (1920, 1080))
raw_imgs[cam] = im
imgs[cam] = resize_im
velo, acc = np.zeros(2), np.zeros(2)
command = np.zeros(4)
# print('cmd', info['command'])
# command[info['command']] = 1
command[1] = 1
ego_pose = OPENCV2IMU @ info['ego_pos']
# yaw = -info['ego_rot'][1]
yaw = -info['ego_steer']
forward_velo = info['ego_velo']
forward_acc = info['accelerate']
velo[0] = forward_velo * np.cos(yaw)
velo[1] = forward_velo * np.sin(yaw)
acc[0] = forward_acc * np.cos(yaw)
acc[1] = forward_acc * np.sin(yaw)
ego_status = EgoStatus(ego_pose, velo, acc, command)
cameras_dict = {}
for nusc_cam, nuplan_cam in zip(nusc_cameras, nuplan_cameras):
cameras_dict[nuplan_cam] = Camera(image=imgs[nusc_cam])
cameras = Cameras(*([None] * 8))
for cam_name, cam in cameras_dict.items():
setattr(cameras, cam_name, cam)
lidar = Lidar(np.zeros((6, 100), dtype=np.float32))
agent_input = AgentInput([ego_status], [cameras], [lidar])
data = {
'input': agent_input,
'raw_imgs': raw_imgs,
}
return data
|