Spaces:
Running on Zero
Running on Zero
File size: 2,521 Bytes
e340a84 | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | 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)
|