nuScene radar points 그리기

mmdetection3d에 있는 코드를 이해하려면 데이터 구조를 이해해야한다.

 

우선 nuscene을 예로 들면 token구조를 이해해야한다. 

 

그리고 nuscene 에서 points를 그리려면 global 좌표계로 좌표변환을 해야한다. 

 

대부분 로직들은 global 좌표계로 변환 후 실행하고 있다.

 

해달 코드는 radar의 points를 읽어서 global 좌표로 변환 후 그리는 코드이다.

import os
from pyquaternion import Quaternion
import numpy as np
from nuscenes.utils.data_classes import RadarPointCloud
from nuscenes.nuscenes import NuScenes
import matplotlib.pyplot as plt
import cv2

# Initialize NuScenes
nusc = NuScenes(version='v1.0-mini', dataroot='/data/sets/nuscenes', verbose=True)

x_min_value = 1e+4
x_max_value = -1e+4

y_min_value = 1e+4
y_max_value = -1e+4


def load_and_transform_radar_pointcloud(nusc, radar_token):
    pointsensor = nusc.get('sample_data', radar_token)
    pc = RadarPointCloud.from_file(os.path.join(nusc.dataroot, pointsensor['filename']))
    
    cs_record = nusc.get('calibrated_sensor', pointsensor['calibrated_sensor_token'])
    pc.rotate(Quaternion(cs_record['rotation']).rotation_matrix)
    pc.translate(np.array(cs_record['translation']))

    poserecord = nusc.get('ego_pose', pointsensor['ego_pose_token'])
    pc.rotate(Quaternion(poserecord['rotation']).rotation_matrix)
    pc.translate(np.array(poserecord['translation']))
    
    return pc

def plot_radar_data(nusc, sample_token, ax):
    sample_record = nusc.get('sample', sample_token)

    # Radar tokens
    radar_tokens = {
        'front': sample_record['data']['RADAR_FRONT'],
        'front_left': sample_record['data']['RADAR_FRONT_LEFT'],
        'front_right': sample_record['data']['RADAR_FRONT_RIGHT'],
        'back_left': sample_record['data']['RADAR_BACK_LEFT'],
        'back_right': sample_record['data']['RADAR_BACK_RIGHT']
    }

    # Load and transform point clouds
    pcs = {key: load_and_transform_radar_pointcloud(nusc, token) for key, token in radar_tokens.items()}

    # Clear previous plot
    # ax.clear()

    # Plot the point clouds
    for pc in pcs.values():
        ax.plot(pc.points[0], pc.points[1], '.')
        
    global x_min_value, x_max_value,y_min_value,y_max_value
    # Set axis limits
    if np.min(pc.points[0]) < x_min_value:
        x_min_value = np.min(pc.points[0])
        
    if np.max(pc.points[0]) > x_max_value:
        x_max_value = np.max(pc.points[0])
    
    if np.min(pc.points[1]) < y_min_value:
        y_min_value = np.min(pc.points[1])
        
    if np.max(pc.points[1]) > y_max_value:
        y_max_value = np.max(pc.points[1])
    
    ax.set_xlim(x_min_value,x_max_value )
    ax.set_ylim(y_min_value,y_max_value)
            
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_title('Radar Point Clouds')
    ax.grid(True)

def extract_number(filename, name ):
    try:
        number_str = filename.split(f'{name}')[1].split('.')[0]
        return int(number_str)
    except IndexError:
        return None  # "__CAM_BACK__" 뒤에 숫자가 없는 경우 처리

def make_mf4(scene_type):
    # 동영상 작성
    video_filename = os.path.join(output_dir, f'combined_radar_pts_{scene_type}.mp4')
    files = os.listdir(output_dir)
    file_list = []
    for file in files:
        if file.endswith('.jpg'):
            file_list.append(file)
    # 파일 이름에서 숫자를 추출하여 정렬하기
    frame_files = sorted(file_list, key=lambda x: extract_number(x,'combined_radar_pts_') or float('inf'))      

    # 첫 번째 프레임에서 프레임 크기 가져오기
    frame = cv2.imread(os.path.join(output_dir,frame_files[403]))
    height, width, layers = frame.shape

    # 동영상 작성자 초기화
    video = cv2.VideoWriter(video_filename, cv2.VideoWriter_fourcc(*'mp4v'), 10, (width, height))

    # 각 프레임을 동영상 작성자에 추가
    idx = 0
    for frame_file in frame_files:
        frame_path = os.path.join(output_dir, frame_file)
        frame = cv2.imread(frame_path)
        if frame is None:
            print(f"프레임을 읽을 수 없습니다: {frame_path}")
            continue
        idx += 1
        print(f"프레임 {idx}: {frame_path}")
        resized_frame = cv2.resize(frame, (width, height))
        video.write(resized_frame)

    # 동영상 작성 종료
    video.release()

# Set up the plot
fig, ax = plt.subplots(figsize=(10, 10))

scene_type = 'total'

num_scene = len(nusc.scene)

i = 0

enable_write_avi = True

output_dir = './combined_radar_pts'
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

for j in range(num_scene):
    # Iterate through all samples in the scene
    scene = nusc.scene[j]  # Adjust the scene index as needed
    current_sample_token = scene['first_sample_token']

    # 결과 저장 디렉토리
    while current_sample_token:
        plot_radar_data(nusc, current_sample_token, ax)
        # 이미지 저장 
        i=i+1
        output_path = os.path.join(output_dir, f"combined_radar_pts_{i}.jpg")
        plt.savefig(output_path)  # PNG 형식으로 저장    
        
        current_sample = nusc.get('sample', current_sample_token)
        current_sample_token = current_sample['next']  # Move to the next sample
        print(j,i)
        plt.pause(1) 

plt.show()


if enable_write_avi:
    make_mf4(scene_type)

 

 

global  좌표로 변경하는 코드는 간단하다

 

센서 calibration값, 자차 이동 보정 값 이렇게 2번하고 

 

각 센서당 이렇게 총 4번의 matrix연산을 아래와 같이 한다.

# 센서값 반영
cs_record = nusc.get('calibrated_sensor', pointsensor['calibrated_sensor_token'])
pc.rotate(Quaternion(cs_record['rotation']).rotation_matrix)
pc.translate(np.array(cs_record['translation']))

# 자차 값 반영
poserecord = nusc.get('ego_pose', pointsensor['ego_pose_token'])
pc.rotate(Quaternion(poserecord['rotation']).rotation_matrix)
pc.translate(np.array(poserecord['translation']))