| import io |
| import struct |
| from contextlib import contextmanager |
| from typing import BinaryIO, Iterator, Optional |
|
|
| import numpy as np |
|
|
|
|
| def write_ply( |
| raw_f: BinaryIO, |
| coords: np.ndarray, |
| rgb: Optional[np.ndarray] = None, |
| faces: Optional[np.ndarray] = None, |
| ): |
| """ |
| Write a PLY file for a mesh or a point cloud. |
| |
| :param coords: an [N x 3] array of floating point coordinates. |
| :param rgb: an [N x 3] array of vertex colors, in the range [0.0, 1.0]. |
| :param faces: an [N x 3] array of triangles encoded as integer indices. |
| """ |
| with buffered_writer(raw_f) as f: |
| f.write(b"ply\n") |
| f.write(b"format binary_little_endian 1.0\n") |
| f.write(bytes(f"element vertex {len(coords)}\n", "ascii")) |
| f.write(b"property float x\n") |
| f.write(b"property float y\n") |
| f.write(b"property float z\n") |
| if rgb is not None: |
| f.write(b"property uchar red\n") |
| f.write(b"property uchar green\n") |
| f.write(b"property uchar blue\n") |
| if faces is not None: |
| f.write(bytes(f"element face {len(faces)}\n", "ascii")) |
| f.write(b"property list uchar int vertex_index\n") |
| f.write(b"end_header\n") |
|
|
| if rgb is not None: |
| rgb = (rgb * 255.499).round().astype(int) |
| vertices = [ |
| (*coord, *rgb) |
| for coord, rgb in zip( |
| coords.tolist(), |
| rgb.tolist(), |
| ) |
| ] |
| format = struct.Struct("<3f3B") |
| for item in vertices: |
| f.write(format.pack(*item)) |
| else: |
| format = struct.Struct("<3f") |
| for vertex in coords.tolist(): |
| f.write(format.pack(*vertex)) |
|
|
| if faces is not None: |
| format = struct.Struct("<B3I") |
| for tri in faces.tolist(): |
| f.write(format.pack(len(tri), *tri)) |
|
|
|
|
| @contextmanager |
| def buffered_writer(raw_f: BinaryIO) -> Iterator[io.BufferedIOBase]: |
| if isinstance(raw_f, io.BufferedIOBase): |
| yield raw_f |
| else: |
| f = io.BufferedWriter(raw_f) |
| yield f |
| f.flush() |
|
|
