Dual-Radar visualize(C-R-L)

https://github.com/adept-thu/dual-radar

GitHub - adept-thu/Dual-Radar

4D Radar dataset으로 메일로 요청하여 데이터를 받을 수 있다.

 

import numpy as np
import open3d as o3d
import os.path as osp

FRAME_NUM = 57
str_frame_num = f'{FRAME_NUM}'.zfill(6)
header = './training'
calib_file = osp.join(header, 'training_calib', 'calib', str_frame_num + '.txt')
lidar_file = osp.join(header, 'training_robosense', 'robosense', str_frame_num + '.bin')
arbe_file = osp.join(header, 'training_arbe', 'arbe', str_frame_num + '.bin')
label_file = osp.join(header, 'training_label', 'label', str_frame_num + '.txt')

def get_calib_from_file(calib_file):
    with open(calib_file) as f:
        lines = f.readlines()
    obj = lines[0].strip().split(' ')[1:]
    P2 = np.array(obj, dtype=np.float32)
    obj = lines[1].strip().split(' ')[1:]
    R0 = np.array(obj, dtype=np.float32)
    obj = lines[2].strip().split(' ')[1:]
    Tr_velo_to_cam = np.array(obj, dtype=np.float32)
    obj = lines[3].strip().split(' ')[1:]
    Tr_arbe_to_velo = np.array(obj, dtype=np.float32)
    obj = lines[4].strip().split(' ')[1:]
    Tr_ars_to_velo = np.array(obj, dtype=np.float32)

    return {'P2': P2.reshape(3, 4),
            'Tr_arbe2velo': Tr_arbe_to_velo.reshape(3, 4),
            'Tr_ars2velo': Tr_ars_to_velo.reshape(3, 4),
            'R0': R0.reshape(3, 3),
            'Tr_velo2cam': Tr_velo_to_cam.reshape(3, 4)}

def get_objects_from_label(label_file):
    with open(label_file, 'r') as f:
        lines = f.readlines()
    objects = [Object3d(line) for line in lines]
    return objects


def cls_type_to_id(cls_type):
    type_to_id = {'Car': 1, 'Pedestrian': 2, 'Cyclist': 3, 'Van': 4}
    if cls_type not in type_to_id.keys():
        return -1
    return type_to_id[cls_type]

dict_label = { # (consider, logit_idx, rgb, bgr)
    'Sedan':            [True,  1, [0, 1, 0],       [0,255,0]],
    'Bus or Truck':     [True,  2, [1, 0.2, 0],     [0,50,255]],
    'Motorcycle':       [True,  3, [1, 0, 0],       [0,0,255]],
    'Bicycle':          [True,  4, [1, 1, 0],       [0,255,255]],
    'Bicycle Group':    [True,  5, [0, 0.5, 1],     [0,128,255]],
    'Pedestrian':       [True,  6, [0, 0, 1],       [255,0,0]],
    'Pedestrian Group': [True,  7, [0.4, 0, 1],     [255,0,100]],
    'Label':            [True,  0, [0.5, 0.5, 0.5], [128,128,128]],
}

def create_cylinder_mesh(radius, p0, p1, color=[1, 0, 0]):
    cylinder = o3d.geometry.TriangleMesh.create_cylinder(radius=radius, height=np.linalg.norm(np.array(p1)-np.array(p0)))
    cylinder.paint_uniform_color(color)
    frame = np.array(p1) - np.array(p0)
    frame /= np.linalg.norm(frame)
    R = o3d.geometry.get_rotation_matrix_from_xyz((np.arccos(frame[2]), np.arctan2(-frame[0], frame[1]), 0))
    cylinder.rotate(R, center=[0, 0, 0])
    cylinder.translate((np.array(p0) + np.array(p1)) / 2)
    return cylinder

def draw_3d_box_in_cylinder(vis, center, theta, l, w, h, color=[1, 0, 0], cali=[], radius=0.1,  in_cylinder=True):
    R = np.array([[np.cos(theta), -np.sin(theta), 0],
                [np.sin(theta),  np.cos(theta), 0],
                [0,              0,             1]])
    

    corners = np.array([[l/2, w/2, h/2], [l/2, w/2, -h/2], [l/2, -w/2, h/2], [l/2, -w/2, -h/2],
                        [-l/2, w/2, h/2], [-l/2, w/2, -h/2], [-l/2, -w/2, h/2], [-l/2, -w/2, -h/2]])
    
    
    corners_rotated = np.dot(corners, R.T) + center # 이 형태 따라서 치기
    
    #1. (3x4) matrix
    # R_cam_to_velo = cali[:,:3].T
    # t_cam_to_velo = np.dot(cali[:,:3].T, cali[:,3]) # -R.T * t


    # # compute rotation
    # corners_rotated = np.dot(corners_rotated, R_cam_to_velo.T) + t_cam_to_velo

    #2. inverse.Tr_velo_to_cam * corner_rotated (3 x N) -> (4 X N) (3 x N)    
    # cam2velo = np.linalg.inv(np.vstack([cali, [0,0,0,1]]))[:3,:] # 3x4

    # corners_rotated = np.vstack([corners_rotated.T, [1,1,1,1,1,1,1,1]]) #4x8

    # corners_rotated = cam2velo @ corners_rotated #3x8
    # corners_rotated = corners_rotated.T

    lines = [[0, 1], [0, 2], [1, 3], [2, 3], [4, 5], [4, 6], [5, 7], [6, 7],
            [0, 4], [1, 5], [2, 6], [3, 7]]
    line_set = o3d.geometry.LineSet()
    line_set.points = o3d.utility.Vector3dVector(corners_rotated)
    line_set.lines = o3d.utility.Vector2iVector(lines)
    line_set.colors = o3d.utility.Vector3dVector([color for i in range(len(lines))])
    if in_cylinder:
        for line in lines:
            cylinder = create_cylinder_mesh(radius, corners_rotated[line[0]], corners_rotated[line[1]], color)
            vis.add_geometry(cylinder)
    else:
        vis.add_geometry(line_set)

class Object3d(object):
    def __init__(self, line):
        label = line.strip().split(' ')
        self.src = line
        self.cls_type = label[0]
        self.cls_id = cls_type_to_id(self.cls_type)
        self.truncation = float(label[1])
        self.occlusion = float(label[2])  # 0:fully visible 1:partly occluded 2:largely occluded 3:unknown
        self.alpha = float(label[3])
        self.box2d = np.array((float(label[4]), float(label[5]), float(label[6]), float(label[7])), dtype=np.float32)
        self.h = float(label[8])
        self.w = float(label[9])
        self.l = float(label[10])
        self.loc = np.array((float(label[11]), float(label[12]), float(label[13])), dtype=np.float32)
        self.dis_to_cam = np.linalg.norm(self.loc)
        self.ry = float(label[14])
        self.score = float(label[15]) if label.__len__() == 16 else -1.0
        self.level_str = None
        self.level = self.get_kitti_obj_level()
        self.track_id = float(label[16])

    def get_xyzthlwh(self):
        # (self.loc[0], self.loc[2], -self.loc[1]) 
        cali = get_calib_from_file(calib_file=calib_file)["Tr_velo2cam"]
        inv_cali = np.linalg.inv(np.vstack([cali, [0,0,0,1]]))[:3,:]

        calibrated = np.dot(inv_cali, np.array((self.loc[0], self.loc[1], self.loc[2], 1)))
        
        return ((calibrated[0], calibrated[1], calibrated[2]), self.ry, self.l, self.w, self.h, self.cls_id)

    def get_kitti_obj_level(self):
        height = float(self.box2d[3]) - float(self.box2d[1]) + 1
        #fangchange
        if height >= 40 and self.truncation <= 0.15 and self.occlusion <= 0:
            self.level_str = 'Easy'
            return 0  # Easy
        elif height >= 25 and self.truncation <= 0.3 and self.occlusion <= 1:
            self.level_str = 'Moderate'
            return 1  # Moderate
        elif height >= 25 and self.truncation <= 0.5 and self.occlusion <= 2:
            self.level_str = 'Hard'
            return 2  # Hard
        else:
            self.level_str = 'UnKnown'
            return -1

    def generate_corners3d(self):
        """
        generate corners3d representation for this object
        :return corners_3d: (8, 3) corners of box3d in camera coord
        """
        l, h, w = self.l, self.h, self.w
        x_corners = [l / 2, l / 2, -l / 2, -l / 2, l / 2, l / 2, -l / 2, -l / 2]
        y_corners = [0, 0, 0, 0, -h, -h, -h, -h]
        z_corners = [w / 2, -w / 2, -w / 2, w / 2, w / 2, -w / 2, -w / 2, w / 2]

        R = np.array([[np.cos(self.ry), 0, np.sin(self.ry)],
                      [0, 1, 0],
                      [-np.sin(self.ry), 0, np.cos(self.ry)]])
        corners3d = np.vstack([x_corners, y_corners, z_corners])  # (3, 8)
        corners3d = np.dot(R, corners3d).T
        corners3d = corners3d + self.loc
        return corners3d

    def to_str(self):
        print_str = '%s %.3f %.3f %.3f box2d: %s hwl: [%.3f %.3f %.3f] pos: %s ry: %.3f' \
                     % (self.cls_type, self.truncation, self.occlusion, self.alpha, self.box2d, self.h, self.w, self.l,
                        self.loc, self.ry)
        return print_str

    def to_kitti_format(self):
        kitti_str = '%s %.2f %d %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f' \
                    % (self.cls_type, self.truncation, int(self.occlusion), self.alpha, self.box2d[0], self.box2d[1],
                       self.box2d[2], self.box2d[3], self.h, self.w, self.l, self.loc[0], self.loc[1], self.loc[2],
                       self.ry)
        return kitti_str

if __name__ == '__main__':
    points_ldr = np.fromfile(str(lidar_file), dtype=np.float32).reshape(-1, 6)
    points_arbe = np.fromfile(str(arbe_file), dtype=np.float32).reshape(-1, 5)

    # calib (arbe to velo)
    dict_calib = get_calib_from_file(calib_file)
    points_arbe_xyz = np.concatenate((points_arbe[:,:3], np.ones((points_arbe.shape[0],1))), axis=1)
    points_arbe_xyz = points_arbe_xyz @ dict_calib['Tr_arbe2velo'].T
    points_arbe[:,:3] = points_arbe_xyz

    vis = o3d.visualization.Visualizer()
    vis.create_window()

    # ldr
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points_ldr[:,:3])
    vis.add_geometry(pcd)

    # arbe
    pcd_rdr = o3d.geometry.PointCloud()
    pcd_rdr.points = o3d.utility.Vector3dVector(points_arbe[:,:3])
    pcd_rdr.paint_uniform_color([0.,0.,0.])
    vis.add_geometry(pcd_rdr)

    cam2velo = dict_calib['Tr_velo2cam']

    # bbox
    list_label_values = list(dict_label.values())
    list_labels = get_objects_from_label(label_file)
    for temp_label in list_labels:
        loc, th, l, w, h, cls_idx = temp_label.get_xyzthlwh()
        consider, cls_id, rgb, bgr = list_label_values[cls_idx]
        draw_3d_box_in_cylinder(vis, loc, th, l, w, h, rgb, cam2velo, radius=0.2)

    vis.run()
    vis.destroy_window()