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.
58 lines
1.9 KiB
58 lines
1.9 KiB
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))
|
|
|