t2_voxel_down_sample

t2_voxel_down_sample.py

import os
import sys
import argparse

import numpy as np
import open3d as o3d

sys.path.append('./tools/python_toolbox/evaluation')

# Tau value for each scene in Table 1 of the paper.
scene_tau_dict = {
    'Family': 0.003,
    'Francis': 0.005,
    'Horse': 0.003,
    'Lighthouse': 0.010,
    'M60': 0.005,
    'Panther': 0.005,
    'Playground': 0.010,
    'Train': 0.005,
    'Auditorium': 0.010,
    'Ballroom': 0.010,
    'Courtroom': 0.010,
    'Museum': 0.010,
    'Palace': 0.030,
    'Temple': 0.015,
}


def main():
    # parse arguments
    parser = argparse.ArgumentParser(description='Batch script for point cloud down-sampling')
    parser.add_argument('-i', '--input', type=str, required=True,
                        help='Input folder which stores point clouds')
    parser.add_argument('-o', '--output', type=str, default=None,
                        help='Output folder to store down-sampled point clouds')
    args = parser.parse_args()

    # output directory
    out_dir = args.output if args.output is not None else os.path.join(args.input, 'downsample')
    os.makedirs(out_dir, exist_ok=True)

    # list of point clouds
    pcl_files = [f for f in os.listdir(args.input) if f.endswith('.ply')]
    pcl_files = sorted(pcl_files)
    scene_names = [os.path.splitext(os.path.basename(f))[0] for f in pcl_files]
    pcl_files = [os.path.join(args.input, f) for f in pcl_files]
    pcl_files = {name: f for name, f in zip(scene_names, pcl_files)}

    for name, f in pcl_files.items():
        print(name + ': ' + f)

        print('Now loading ... ', end='', flush=True)
        pcl = o3d.io.read_point_cloud(f)
        d_tau = scene_tau_dict[name]
        print('--> OK!')

        print('dTau = %f' % (d_tau))
        print('Down-sampling ... ', end='', flush=True)
        pcl_down = pcl.voxel_down_sample(d_tau / 2.0)
        print('--> OK!')

        print('Now writing ... ', end='', flush=True)
        out_file = os.path.join(out_dir, os.path.basename(f))
        o3d.io.write_point_cloud(out_file, pcl_down)
        print('--> OK!')

        print('Saved to:', out_file)
        print('')


if __name__ == '__main__':
    main()