Waymoのデータセットをいじり倒したい！part①

はじめに

みなさん、こんにちは！今回は初学者の私がWaymoのデータセットを使って、前処理を行った方法についてご紹介します。Waymoのデータセットは非常に大規模で、前処理には多くのステップが必要です。この記事では、私が直面した苦労とこれからのプランについてもお話しします。

今までの苦労

1. データの読み込みと確認

まず、Waymoのデータセットを読み込むために、
pyarrowを使用して.parquetファイルを読み込みました。しかし、データの構造が複雑で、正しく読み込むのに苦労しました。特に、ポイントクラウドデータが正しく読み込まれない問題が発生しました。

2. データの前処理

データの前処理では、ポイントクラウドデータとラベルデータを正しく処理する必要がありました。特に、ポイントクラウドデータのパディングやラベルデータのパディングに苦労しました。また、複数のポイントクラウドを保存する際に、データの形式やサイズに問題が発生しました。

3. データの保存

前処理されたデータを保存する際にも、データの形式やサイズに注意が必要でした。特に、複数のポイントクラウドを保存する際に、データが正しく保存されない問題が発生しました。

4. データの確認

前処理されたデータが正しく保存されているか確認するために、データを読み込み、内容を確認しました。ポイントクラウドが正しく読み込まれなかったり、複数のポイントクラウドを保存できなかったりする問題が発生しました。

ソースコード

以下に前処理を実行したコードを記載します。ディレクトリは変えてください。

import pandas as pd
import os
from tqdm import tqdm
import pyarrow.parquet as pq
import numpy as np
import json
import pickle

# データセットのディレクトリ
data_dir = "INPUT_YOUR_DIRECTORY"
output_dir = "OUTPUT_YOUR_DIRECTORY"

# データの読み込み関数
def load_parquet_file(file_path):
    try:
        # pyarrowを使用してデータを読み込む
        table = pq.read_table(file_path)
        df = table.to_pandas()
        return df
    except Exception as e:
        print(f"Error reading {file_path}: {e}")
        return None

# ポイントクラウドデータの前処理関数
def preprocess_point_cloud(df_lidar, df_labels):
    point_clouds = []
    labels = []
    # LiDARデータの前処理
    range_image_strs = df_lidar['[LiDARComponent].range_image_return1.values'].values
    for range_image_str in range_image_strs:
        if isinstance(range_image_str, str):
            range_image = np.array(json.loads(range_image_str))
        else:
            range_image = np.array(range_image_str)
        if range_image.ndim == 1:
            range_image = range_image.reshape(-1, 4)
        valid_points = range_image[(range_image[:, 0] != -1) & (range_image[:, 3] != -1)]
        if valid_points.ndim == 2 and valid_points.shape[1] >= 4:
            x = valid_points[:, 0]
            y = valid_points[:, 1]
            z = valid_points[:, 2]
            intensity = valid_points[:, 3]
            points = np.vstack((x, y, z, intensity)).T
            point_clouds.append(points)
        else:
            print("Invalid range_image shape")
    
    # ラベルデータの前処理
    print(df_labels.columns)  # カラム名を表示して確認
    for _, row in df_labels.iterrows():
        # 適切なカラム名を使用
        label = row['[LiDARBoxComponent].type']  # 'type'カラムを使用
        box = row[['[LiDARBoxComponent].box.center.x', '[LiDARBoxComponent].box.center.y', '[LiDARBoxComponent].box.center.z', '[LiDARBoxComponent].box.size.x', '[LiDARBoxComponent].box.size.y', '[LiDARBoxComponent].box.size.z']].values
        labels.append((label, box))
    
    return point_clouds, labels

# 各ファイルの前処理を行う関数
def process_all_files(lidar_folder, label_folder, output_folder):
    all_point_clouds = []
    all_labels = []
    for root, _, files in os.walk(lidar_folder):
        for file in tqdm(files, desc=f"Processing files in {root}"):
            if file.endswith('.parquet'):
                lidar_file_path = os.path.join(root, file)
                label_file_path = os.path.join(label_folder, file)
                df_lidar = load_parquet_file(lidar_file_path)
                df_labels = load_parquet_file(label_file_path)
                if df_lidar is not None and df_labels is not None:
                    print(f"Loaded files: {lidar_file_path}, {label_file_path}")
                    point_clouds, labels = preprocess_point_cloud(df_lidar, df_labels)
                    print(f"Processed {len(point_clouds)} point clouds")
                    if point_clouds:
                        all_point_clouds.extend(point_clouds)
                        all_labels.extend(labels)
                else:
                    print(f"Failed to load files: {lidar_file_path}, {label_file_path}")
    # 前処理されたデータを保存
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    with open(os.path.join(output_folder, 'processed_point_clouds.pkl'), 'wb') as f:
        pickle.dump((all_point_clouds, all_labels), f)
    
    # ポイントクラウドの数とラベルの数を表示
    print(f"Total number of point clouds: {len(all_point_clouds)}")
    print(f"Total number of labels: {len(all_labels)}")

# 各ディレクトリのパス
directory = "training"
lidar_folder = os.path.join(data_dir, directory, "lidar")
label_folder = os.path.join(data_dir, directory, "lidar_box")
output_folder = os.path.join(output_dir, directory)

# 前処理を実行
if os.path.exists(lidar_folder) and os.path.exists(label_folder):
    process_all_files(lidar_folder, label_folder, output_folder)
else:
    print(f"Folder does not exist: {lidar_folder} or {label_folder}")

途中経過

これからのプラン

1. モデルのトレーニング

前処理が完了したデータを使用して、物体検出モデルをトレーニングします。PointPillarsやPointRCNNなどのモデルを使用して、高精度な物体検出を目指します。

2. リアルタイム物体検出

トレーニングされたモデルを使用して、リアルタイムで物体検出を行います。CARLAシミュレーターを使用して、実際の運転シナリオでモデルをテストします。

3. モデルの評価と改善

モデルの精度を評価し、必要に応じて改善を行います。データの前処理やモデルのハイパーパラメータを調整して、より高精度な物体検出を目指します。

まとめ

自分が思っていた以上に前処理のコードを完成させるのに時間がかかりました。これからも、多くのステップと苦労が伴うと思います。次回の記事では、モデルのトレーニングについて詳しく紹介しますので、お楽しみに！もしみなさんも初学者時代に苦労したことなどあったら教えてください！

それではまた！