Spaces:
Running on Zero
Running on Zero
| import os | |
| import numpy as np | |
| def _maybe_downsample(points, colors=None, max_points=None, seed=0): | |
| pts = np.asarray(points).reshape(-1, 3) | |
| cols = None if colors is None else np.asarray(colors).reshape(-1, 3) | |
| if max_points is None or pts.shape[0] <= int(max_points): | |
| return pts, cols | |
| rng = np.random.default_rng(seed) | |
| keep = rng.choice(pts.shape[0], size=int(max_points), replace=False) | |
| pts = pts[keep] | |
| if cols is not None: | |
| cols = cols[keep] | |
| return pts, cols | |
| def save_pointcloud(path, points, colors=None, max_points=None, seed=0): | |
| os.makedirs(os.path.dirname(path), exist_ok=True) | |
| pts, cols = _maybe_downsample( | |
| points, colors=colors, max_points=max_points, seed=seed | |
| ) | |
| pts = pts.astype(np.float32, copy=False) | |
| if colors is not None: | |
| if cols.max() <= 1.0: | |
| cols = (cols * 255.0).astype(np.uint8) | |
| else: | |
| cols = cols.astype(np.uint8) | |
| has_color = True | |
| else: | |
| cols = None | |
| has_color = False | |
| with open(path, "wb") as f: | |
| f.write(b"ply\n") | |
| f.write(b"format binary_little_endian 1.0\n") | |
| f.write(f"element vertex {pts.shape[0]}\n".encode("ascii")) | |
| f.write(b"property float x\n") | |
| f.write(b"property float y\n") | |
| f.write(b"property float z\n") | |
| if has_color: | |
| f.write(b"property uchar red\n") | |
| f.write(b"property uchar green\n") | |
| f.write(b"property uchar blue\n") | |
| f.write(b"end_header\n") | |
| if has_color: | |
| vertex_dtype = np.dtype( | |
| [ | |
| ("x", "<f4"), | |
| ("y", "<f4"), | |
| ("z", "<f4"), | |
| ("red", "u1"), | |
| ("green", "u1"), | |
| ("blue", "u1"), | |
| ] | |
| ) | |
| vertex_data = np.empty(pts.shape[0], dtype=vertex_dtype) | |
| vertex_data["x"] = pts[:, 0] | |
| vertex_data["y"] = pts[:, 1] | |
| vertex_data["z"] = pts[:, 2] | |
| vertex_data["red"] = cols[:, 0] | |
| vertex_data["green"] = cols[:, 1] | |
| vertex_data["blue"] = cols[:, 2] | |
| vertex_data.tofile(f) | |
| else: | |
| vertex_dtype = np.dtype([("x", "<f4"), ("y", "<f4"), ("z", "<f4")]) | |
| vertex_data = np.empty(pts.shape[0], dtype=vertex_dtype) | |
| vertex_data["x"] = pts[:, 0] | |
| vertex_data["y"] = pts[:, 1] | |
| vertex_data["z"] = pts[:, 2] | |
| vertex_data.tofile(f) | |