create_ground_planes

python create_ground_planes.py

import argparse
import numpy as np
import os
from tqdm import tqdm
from det3d.datasets.utils.ground_plane_detection import fit_plane_LSE_RANSAC

def load_points(bin_file):
    points = np.fromfile(bin_file, dtype=np.float32).reshape(-1, 4)
    points = np.concatenate((points[:, :3], np.ones((points.shape[0], 1))), axis=1)
    return points

def get_ground_plane(points):
    p_set_1 = points
    p1, inlier_list1, outlier_list1 = fit_plane_LSE_RANSAC(
        p_set_1, return_outlier_list=True
    )
    p_set_2 = p_set_1[outlier_list1, :]
    p2, inlier_list2, outlier_list2 = fit_plane_LSE_RANSAC(
        p_set_2, return_outlier_list=True
    )
    p_set_3 = p_set_2[outlier_list2, :]
    p3, inlier_list3, outlier_list3 = fit_plane_LSE_RANSAC(
        p_set_3, return_outlier_list=True
    )
    p_set_4 = p_set_3[outlier_list3, :]
    p4, inlier_list4, outlier_list4 = fit_plane_LSE_RANSAC(
        p_set_4, return_outlier_list=True
    )
    return [p1, p2, p3, p4]

def create_ground_plane(root_path):
    bin_folder = os.path.join(root_path, "training", "velodyne_reduced")
    gp_folder = os.path.join(root_path, "training", "planes")
    if not os.path.exists(gp_folder):
        os.makedirs(gp_folder)
    bin_file_list = os.listdir(bin_folder)

    for i in tqdm(range(len(bin_file_list))):
        gp_file = os.path.join(gp_folder, bin_file_list[i].replace("bin", "txt"))
        bin_file = os.path.join(bin_folder, bin_file_list[i])
        points = load_points(bin_file)
        planes = get_ground_plane(points)
        np.savetxt(gp_file, planes)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Create Ground Planes")
    parser.add_argument("root_path", help="dataset root folder")
    args = parser.parse_args()
    create_ground_plane(args.root_path)

# python create_ground_planes.py /data/Datasets/KITTI/object