a fork of shap-e for gc
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

59 lines
1.9 KiB

2 years ago
import struct
from typing import BinaryIO, Optional
import numpy as np
from shap_e.util.io import buffered_writer
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))