データ補正と処理

以下は、点群データを読み込み、補正を行い、新しいPCDファイルを生成するためのコードです。このコードは、座標変換や時間補正を含む複数のステップで構成されています。

<?python
import numpy as np
import open3d as o3d
from scipy.spatial.transform import Rotation as st
import json
import os

def quaternion_to_transformation(qx, qy, qz, qw):
    rot = st.from_quat([qx, qy, qz, qw])
    transformation = np.eye(4)
    transformation[:3, :3] = rot.as_matrix()
    return transformation

def read_pcd_file(file_path):
    pcd = o3d.io.read_point_cloud(file_path)
    points = np.asarray(pcd.points)
    return points

def generate_new_pcd(header, new_points, output_file):
    if not os.path.exists(os.path.dirname(output_file)):
        os.makedirs(os.path.dirname(output_file))
    with open(output_file, 'w') as f:
        for line in header:
            if 'WIDTH' in line:
                line = f"WIDTH {new_points.shape[0]}\n"
            elif 'POINTS' in line:
                line = f"POINTS {new_points.shape[0]}\n"
            f.write(line + '\n')
        np.savetxt(f, new_points, fmt='%.6f')

def motion_compensation(target_time, points, timestamps, poses):
    transformed_points = []
    keep_mask = [True] * len(points)
    for idx, point in enumerate(points):
        time_diff = target_time - timestamps[idx]
        if time_diff != 0 and timestamps[idx] in poses:
            pose = poses[timestamps[idx]]
            trans_matrix = quaternion_to_transformation(pose['qx'], pose['qy'], pose['qz'], pose['qw'])
            trans_matrix[:3, 3] = np.array([pose['tx'], pose['ty'], pose['tz']])
            point_homogeneous = np.append(point, 1.0)
            transformed_point = (trans_matrix @ point_homogeneous)[:3]
            transformed_points.append(transformed_point)
        else:
            keep_mask[idx] = False
            transformed_points.append(point)
    return np.array(transformed_points), keep_mask

# 使用例
source_dir = '/data/pcd_files'
label_dir = '/data/labels'
output_dir = '/data/output'

for label_file in os.listdir(label_dir):
    if label_file.endswith('.json'):
        json_path = os.path.join(label_dir, label_file)
        with open(json_path, 'r') as f:
            data = json.load(f)
        pcd_name = f"refine_lidars/{label_file.replace('.json', '.pcd')}"
        pcd_path = os.path.join(source_dir, pcd_name)
        points = read_pcd_file(pcd_path)
        timestamps = data['frame_info']['lidar_object_info']['timestamps']
        poses = data['frame_info']['lidar_object_info']['poses']

        compensated_points, mask = motion_compensation(1.0, points, timestamps, poses)

        header = [line.strip() for line in open(pcd_path, 'r').readlines() if not line.startswith('DATA')]
        output_path = os.path.join(output_dir, pcd_name)
        generate_new_pcd(header, compensated_points, output_path)

可視化用コード

以下のコードはMongoDBからデータを取得し、画像上に物体の位置を描画します。これにより、点群データとラベルの対応関係を確認できます。

<?python
from pymongo import MongoClient
import cv2
import numpy as np

client = MongoClient('mongodb://root:password@localhost:27017/')
db = client['perceptionDB']
bag_collection = db['OD_data_bag']
frame_collection = db['OD_data_frame']

def load_boxes(boxes):
    gt_box, gt_name, od_id = [], [], []
    for box in boxes:
        if len(box['psr']) == 0:
            continue
        box3d = [
            box['psr']['position']['x'],
            box['psr']['position']['y'],
            box['psr']['position']['z'],
            box['psr']['scale']['x'],
            box['psr']['scale']['y'],
            box['psr']['scale']['z'],
            box['psr']['rotation']['z']
        ]
        gt_box.append(box3d)
        gt_name.append(box['obj_category'])
        od_id.append(box['track_id'])
    return np.array(gt_box), np.array(gt_name), od_id

def draw_box(img, corners, color):
    for i, j in [
        [0, 1], [1, 2], [2, 3], [3, 0],
        [4, 5], [5, 6], [6, 7], [7, 4],
        [0, 4], [1, 5], [2, 6], [3, 7]
    ]:
        cv2.line(img, tuple(corners[i]), tuple(corners[j]), color, 2)

def get_corners(boxes):
    corners_norm = np.stack(np.unravel_index(np.arange(8), [2]*3), axis=1)
    corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]] - np.array([0.5, 0.5, 0.5])
    corners = boxes[:, 3:6].reshape([-1, 1, 3]) * corners_norm.reshape([1, 8, 3])
    corners = rotation_3d_in_aixs(corners, boxes[:, -1])
    corners += boxes[:, :3].reshape(-1, 1, 3)
    return corners

def match_label_and_img(bag_doc, timestamp, save_dir):
    frame_query = {"bag_name": bag_doc['bag_name'], "timestamp": timestamp}
    frame = frame_collection.find_one(frame_query)
    if frame is None:
        print("Frame not found")
        return

    img_paths = frame['global_properties']['images_path']
    boxes = frame['object']
    gt_box, gt_name, od_id = load_boxes(boxes)
    corners = get_corners(gt_box)

    for cam, img_path in img_paths.items():
        img = cv2.imread(img_path)
        extrinsic = np.array(frame['global_properties'][cam]['extrinsic'])
        intrinsic = np.array(frame['global_properties'][cam]['intrinsic'])

        for corner, name, track_id in zip(corners, gt_name, od_id):
            projected_corners = project_3d_to_2d(corner, intrinsic, extrinsic)
            draw_box(img, projected_corners, (0, 255, 0))

        save_path = os.path.join(save_dir, f"{cam}_{timestamp}.jpg")
        cv2.imwrite(save_path, img)

# 使用例
query = {"batch_number": "A"}
all_documents = bag_collection.find(query)
save_dir = './visualizations'
os.makedirs(save_dir, exist_ok=True)

for doc in all_documents:
    timestamp = sample_single_frame(doc)
    match_label_and_img(doc, timestamp, save_dir)