|
import os |
|
import numpy as np |
|
from plyfile import PlyData, PlyElement |
|
|
|
|
|
def save_point_cloud(pcd, rgb, filename, binary=True): |
|
"""Save an RGB point cloud as a PLY file. |
|
:paras |
|
@pcd: Nx3 matrix, the XYZ coordinates |
|
@rgb: NX3 matrix, the rgb colors for each 3D point |
|
""" |
|
assert pcd.shape[0] == rgb.shape[0] |
|
|
|
if rgb is None: |
|
gray_concat = np.tile(np.array([128], dtype=np.uint8), (pcd.shape[0], 3)) |
|
points_3d = np.hstack((pcd, gray_concat)) |
|
else: |
|
points_3d = np.hstack((pcd, rgb)) |
|
python_types = (float, float, float, int, int, int) |
|
npy_types = [('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('red', 'u1'), ('green', 'u1'), |
|
('blue', 'u1')] |
|
if binary is True: |
|
|
|
vertices = [] |
|
for row_idx in range(points_3d.shape[0]): |
|
cur_point = points_3d[row_idx] |
|
vertices.append(tuple(dtype(point) for dtype, point in zip(python_types, cur_point))) |
|
vertices_array = np.array(vertices, dtype=npy_types) |
|
el = PlyElement.describe(vertices_array, 'vertex') |
|
|
|
|
|
PlyData([el]).write(filename) |
|
else: |
|
x = np.squeeze(points_3d[:, 0]) |
|
y = np.squeeze(points_3d[:, 1]) |
|
z = np.squeeze(points_3d[:, 2]) |
|
r = np.squeeze(points_3d[:, 3]) |
|
g = np.squeeze(points_3d[:, 4]) |
|
b = np.squeeze(points_3d[:, 5]) |
|
|
|
ply_head = 'ply\n' \ |
|
'format ascii 1.0\n' \ |
|
'element vertex %d\n' \ |
|
'property float x\n' \ |
|
'property float y\n' \ |
|
'property float z\n' \ |
|
'property uchar red\n' \ |
|
'property uchar green\n' \ |
|
'property uchar blue\n' \ |
|
'end_header' % r.shape[0] |
|
|
|
np.savetxt(filename, np.column_stack((x, y, z, r, g, b)), fmt="%d %d %d %d %d %d", header=ply_head, comments='') |
|
|