|
|
import argparse |
|
|
from pathlib import Path |
|
|
import open3d as o3d |
|
|
from tqdm.contrib.concurrent import thread_map, process_map |
|
|
import numpy as np |
|
|
|
|
|
def process_scene(data): |
|
|
scene_dir, output_dir = data |
|
|
point_cloud = o3d.io.read_point_cloud(scene_dir) |
|
|
xyz = np.asarray(point_cloud.points) |
|
|
|
|
|
rgb = np.array(point_cloud.colors) |
|
|
|
|
|
rgb = np.clip(rgb, 0, 255)[:, :3] |
|
|
|
|
|
|
|
|
if not len(rgb): |
|
|
return None |
|
|
|
|
|
if rgb.max() <= 1: |
|
|
rgb = (rgb * 255) |
|
|
|
|
|
points = np.concatenate([xyz, rgb], axis=1).astype(np.float32) |
|
|
output_path = output_dir / f"{scene_dir.parent.name}_point.bin" |
|
|
print(f"saving {scene_dir} to {output_path}") |
|
|
points.tofile(output_path) |
|
|
|
|
|
|
|
|
|
|
|
return output_path |
|
|
|
|
|
def load_scan(pcd_path): |
|
|
pcd_data = np.fromfile(pcd_path, dtype=np.float32).reshape(-1, 6) |
|
|
return pcd_data |
|
|
|
|
|
def main(): |
|
|
parser = argparse.ArgumentParser() |
|
|
parser.add_argument("-i", "--input", type=str, required=True) |
|
|
parser.add_argument("-o", "--output", type=str, required=True) |
|
|
args = parser.parse_args() |
|
|
|
|
|
input_dir = Path(args.input) / "processed" |
|
|
output_dir = Path(args.output) |
|
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True) |
|
|
|
|
|
input_files = list(input_dir.glob("*/*.ply")) |
|
|
data = [*zip(input_files, [output_dir] * len(input_files))] |
|
|
|
|
|
process_map(process_scene, data, max_workers=16) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
main() |
