File size: 1,476 Bytes
55e58d1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
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 rgb [0, 1] then change to [0, 255]
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)
# print(points, points.shape)
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() |