|
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() |
|
|
