wehub-resource-sync
269 lines
9.5 KiB
Python
269 lines
9.5 KiB
Python
from typing import Tuple, Dict
|
|
import numpy as np
|
|
from trimesh import grouping, util, remesh
|
|
import struct
|
|
import re
|
|
from plyfile import PlyData, PlyElement
|
|
|
|
|
|
def read_ply(filename):
|
|
"""
|
|
Read a PLY file and return vertices, triangle faces, and quad faces.
|
|
|
|
Args:
|
|
filename (str): The file path to read from.
|
|
|
|
Returns:
|
|
vertices (np.ndarray): Array of shape [N, 3] containing vertex positions.
|
|
tris (np.ndarray): Array of shape [M, 3] containing triangle face indices (empty if none).
|
|
quads (np.ndarray): Array of shape [K, 4] containing quad face indices (empty if none).
|
|
"""
|
|
with open(filename, 'rb') as f:
|
|
# Read the header until 'end_header' is encountered
|
|
header_bytes = b""
|
|
while True:
|
|
line = f.readline()
|
|
if not line:
|
|
raise ValueError("PLY header not found")
|
|
header_bytes += line
|
|
if b"end_header" in line:
|
|
break
|
|
header = header_bytes.decode('utf-8')
|
|
|
|
# Determine if the file is in ASCII or binary format
|
|
is_ascii = "ascii" in header
|
|
|
|
# Extract the number of vertices and faces from the header using regex
|
|
vertex_match = re.search(r'element vertex (\d+)', header)
|
|
if vertex_match:
|
|
num_vertices = int(vertex_match.group(1))
|
|
else:
|
|
raise ValueError("Vertex count not found in header")
|
|
|
|
face_match = re.search(r'element face (\d+)', header)
|
|
if face_match:
|
|
num_faces = int(face_match.group(1))
|
|
else:
|
|
raise ValueError("Face count not found in header")
|
|
|
|
vertices = []
|
|
tris = []
|
|
quads = []
|
|
|
|
if is_ascii:
|
|
# For ASCII format, read each line of vertex data (each line contains 3 floats)
|
|
for _ in range(num_vertices):
|
|
line = f.readline().decode('utf-8').strip()
|
|
if not line:
|
|
continue
|
|
parts = line.split()
|
|
vertices.append([float(parts[0]), float(parts[1]), float(parts[2])])
|
|
|
|
# Read face data, where the first number indicates the number of vertices for the face
|
|
for _ in range(num_faces):
|
|
line = f.readline().decode('utf-8').strip()
|
|
if not line:
|
|
continue
|
|
parts = line.split()
|
|
count = int(parts[0])
|
|
indices = list(map(int, parts[1:]))
|
|
if count == 3:
|
|
tris.append(indices)
|
|
elif count == 4:
|
|
quads.append(indices)
|
|
else:
|
|
# Skip faces with other numbers of vertices (can be extended as needed)
|
|
pass
|
|
else:
|
|
# For binary format: read directly from the binary stream
|
|
# Each vertex consists of 3 floats (12 bytes per vertex)
|
|
for _ in range(num_vertices):
|
|
data = f.read(12)
|
|
if len(data) < 12:
|
|
raise ValueError("Insufficient vertex data")
|
|
v = struct.unpack('<fff', data)
|
|
vertices.append(v)
|
|
|
|
# Read face data from the binary stream
|
|
for _ in range(num_faces):
|
|
# First, read 1 byte indicating the number of vertices in the face
|
|
count_data = f.read(1)
|
|
if len(count_data) < 1:
|
|
raise ValueError("Failed to read face vertex count")
|
|
count = struct.unpack('<B', count_data)[0]
|
|
if count == 3:
|
|
data = f.read(12) # 3 * 4 bytes
|
|
if len(data) < 12:
|
|
raise ValueError("Insufficient data for triangle face")
|
|
indices = struct.unpack('<3i', data)
|
|
tris.append(indices)
|
|
elif count == 4:
|
|
data = f.read(16) # 4 * 4 bytes
|
|
if len(data) < 16:
|
|
raise ValueError("Insufficient data for quad face")
|
|
indices = struct.unpack('<4i', data)
|
|
quads.append(indices)
|
|
else:
|
|
# For faces with a different number of vertices, read count*4 bytes
|
|
data = f.read(count * 4)
|
|
# Skip or extend processing as needed
|
|
raise ValueError(f"Unsupported face with {count} vertices")
|
|
|
|
# Convert lists to torch.Tensor
|
|
vertices = np.array(vertices, dtype=np.float32)
|
|
tris = np.array(tris, dtype=np.int32) if len(tris) > 0 else np.empty((0, 3), dtype=np.int32)
|
|
quads = np.array(quads, dtype=np.int32) if len(quads) > 0 else np.empty((0, 4), dtype=np.int32)
|
|
|
|
return vertices, tris, quads
|
|
|
|
|
|
def write_ply(
|
|
filename: str,
|
|
vertices: np.ndarray,
|
|
tris: np.ndarray,
|
|
quads: np.ndarray,
|
|
vertex_colors: np.ndarray = None,
|
|
ascii: bool = False
|
|
):
|
|
"""
|
|
Write a mesh to a PLY file, with the option to save in ASCII or binary format,
|
|
and optional per-vertex colors.
|
|
|
|
Args:
|
|
filename (str): The filename to write to.
|
|
vertices (np.ndarray): [N, 3] The vertex positions.
|
|
tris (np.ndarray): [M, 3] The triangle indices.
|
|
quads (np.ndarray): [K, 4] The quad indices.
|
|
vertex_colors (np.ndarray, optional): [N, 3] or [N, 4] UInt8 colors for each vertex (RGB or RGBA).
|
|
ascii (bool): If True, write in ASCII format; otherwise binary little-endian.
|
|
"""
|
|
import struct
|
|
|
|
num_vertices = len(vertices)
|
|
num_faces = len(tris) + len(quads)
|
|
|
|
# Build header
|
|
header_lines = [
|
|
"ply",
|
|
f"format {'ascii 1.0' if ascii else 'binary_little_endian 1.0'}",
|
|
f"element vertex {num_vertices}",
|
|
"property float x",
|
|
"property float y",
|
|
"property float z",
|
|
]
|
|
|
|
# Add vertex color properties if provided
|
|
has_color = vertex_colors is not None
|
|
if has_color:
|
|
# Expect uint8 values 0-255
|
|
header_lines += [
|
|
"property uchar red",
|
|
"property uchar green",
|
|
"property uchar blue",
|
|
]
|
|
# Include alpha if RGBA
|
|
if vertex_colors.shape[1] == 4:
|
|
header_lines.append("property uchar alpha")
|
|
|
|
header_lines += [
|
|
f"element face {num_faces}",
|
|
"property list uchar int vertex_index",
|
|
"end_header",
|
|
""
|
|
]
|
|
header = "\n".join(header_lines)
|
|
|
|
mode = 'w' if ascii else 'wb'
|
|
with open(filename, mode) as f:
|
|
# Write header
|
|
if ascii:
|
|
f.write(header)
|
|
else:
|
|
f.write(header.encode('utf-8'))
|
|
|
|
# Write vertex data
|
|
for i, v in enumerate(vertices):
|
|
if ascii:
|
|
line = f"{v[0]} {v[1]} {v[2]}"
|
|
if has_color:
|
|
col = vertex_colors[i]
|
|
line += ' ' + ' '.join(str(int(c)) for c in col)
|
|
f.write(line + '\n')
|
|
else:
|
|
# pack xyz as floats
|
|
f.write(struct.pack('<fff', *v))
|
|
if has_color:
|
|
col = vertex_colors[i]
|
|
# pack as uchar
|
|
if col.shape[0] == 3:
|
|
f.write(struct.pack('<BBB', *col))
|
|
else:
|
|
f.write(struct.pack('<BBBB', *col))
|
|
|
|
# Write face data
|
|
if ascii:
|
|
for tri in tris:
|
|
f.write(f"3 {tri[0]} {tri[1]} {tri[2]}\n")
|
|
for quad in quads:
|
|
f.write(f"4 {quad[0]} {quad[1]} {quad[2]} {quad[3]}\n")
|
|
else:
|
|
for tri in tris:
|
|
f.write(struct.pack('<B3i', 3, *tri))
|
|
for quad in quads:
|
|
f.write(struct.pack('<B4i', 4, *quad))
|
|
|
|
|
|
def write_pbr_ply(
|
|
filename: str,
|
|
vertices: np.ndarray,
|
|
faces: np.ndarray,
|
|
base_color: np.ndarray,
|
|
metallic: np.ndarray,
|
|
roughness: np.ndarray,
|
|
alpha: np.ndarray,
|
|
ascii: bool = False
|
|
):
|
|
"""
|
|
Write a mesh to a PLY file, with the option to save in ASCII or binary format,
|
|
and optional per-vertex colors.
|
|
|
|
Args:
|
|
filename (str): The filename to write to.
|
|
vertices (np.ndarray): [N, 3] The vertex positions.
|
|
faces (np.ndarray): [M, 3] The triangle indices.
|
|
base_color (np.ndarray): [N, 3] UInt8 colors for each vertex (RGB).
|
|
metallic (np.ndarray): [N] UInt8 values for metallicness.
|
|
roughness (np.ndarray): [N] UInt8 values for roughness.
|
|
alpha (np.ndarray): [N] UInt8 values for alpha.
|
|
ascii (bool): If True, write in ASCII format; otherwise binary little-endian.
|
|
"""
|
|
vertex_dtype = [
|
|
('x', 'f4'), ('y', 'f4'), ('z', 'f4'),
|
|
('red', 'u1'), ('green', 'u1'), ('blue', 'u1'),
|
|
('metallic', 'u1'), ('roughness', 'u1'), ('alpha', 'u1')
|
|
]
|
|
|
|
vertex_data = np.empty(len(vertices), dtype=vertex_dtype)
|
|
vertex_data['x'] = vertices[:, 0]
|
|
vertex_data['y'] = vertices[:, 1]
|
|
vertex_data['z'] = vertices[:, 2]
|
|
vertex_data['red'] = base_color[:, 0]
|
|
vertex_data['green'] = base_color[:, 1]
|
|
vertex_data['blue'] = base_color[:, 2]
|
|
vertex_data['metallic'] = metallic
|
|
vertex_data['roughness'] = roughness
|
|
vertex_data['alpha'] = alpha
|
|
|
|
face_dtype = [
|
|
('vertex_indices', 'i4', (3,))
|
|
]
|
|
|
|
face_data = np.empty(len(faces), dtype=face_dtype)
|
|
face_data['vertex_indices'] = faces
|
|
|
|
ply_data = PlyData([
|
|
PlyElement.describe(vertex_data,'vertex'),
|
|
PlyElement.describe(face_data, 'face'),
|
|
], text=ascii)
|
|
ply_data.write(filename)
|