162 lines
6.1 KiB
Python
162 lines
6.1 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, Dict, List, Tuple
|
|
|
|
import numpy as np
|
|
|
|
from .io_gltf import GltfDocument
|
|
from .accessors import access_data, append_accessor_and_bufferview
|
|
from .transforms import compute_mesh_global_transforms, transform_local_vertices_to_world, set_node_trs_identity
|
|
|
|
|
|
def get_all_meshes_triangles(doc: GltfDocument, transform_to_global: bool = True):
|
|
if not doc.meshes():
|
|
raise ValueError("No meshes found in the GLB/GLTF.")
|
|
|
|
all_vertices_groups: List[Dict[str, Any]] = []
|
|
vertices_count = 0
|
|
all_pos_data: List[np.ndarray] = []
|
|
all_indices_data: List[np.ndarray] = []
|
|
|
|
mesh_transforms: Dict[int, np.ndarray] = {}
|
|
if transform_to_global:
|
|
mesh_transforms = compute_mesh_global_transforms(doc)
|
|
|
|
for i_m, mesh_json in enumerate(doc.meshes()):
|
|
for i_p, primitive in enumerate(mesh_json.get("primitives", [])):
|
|
if "attributes" in primitive and "indices" in primitive:
|
|
attributes = primitive["attributes"]
|
|
if "POSITION" in attributes:
|
|
pos_accessor_index = int(attributes["POSITION"])
|
|
pos_data = access_data(doc, pos_accessor_index)
|
|
if pos_data.dtype != np.float32:
|
|
pos_data = pos_data.astype(np.float32)
|
|
|
|
if transform_to_global and i_m in mesh_transforms:
|
|
transform_matrix = mesh_transforms[i_m]
|
|
identity = np.eye(4, dtype=np.float32)
|
|
if not np.allclose(transform_matrix, identity, atol=1e-6):
|
|
fake_node = {"matrix": transform_matrix.flatten(order="F").tolist()}
|
|
pos_data = transform_local_vertices_to_world(pos_data, fake_node)
|
|
|
|
all_vertices_groups.append({
|
|
"mesh_index": i_m,
|
|
"primitive_index": i_p,
|
|
"vertices_range": (vertices_count, vertices_count + pos_data.shape[0]),
|
|
"global_transform_applied": transform_to_global and i_m in mesh_transforms,
|
|
})
|
|
all_pos_data.append(pos_data)
|
|
|
|
indices_accessor_index = int(primitive["indices"])
|
|
indices_data = access_data(doc, indices_accessor_index)
|
|
if indices_data.dtype != np.int64:
|
|
indices_data = indices_data.astype(np.int64)
|
|
all_indices_data.append(indices_data + vertices_count)
|
|
|
|
vertices_count += pos_data.shape[0]
|
|
|
|
return (
|
|
np.concatenate(all_pos_data, axis=0),
|
|
np.concatenate(all_indices_data, axis=0).reshape(-1, 3),
|
|
all_vertices_groups,
|
|
)
|
|
|
|
|
|
def _choose_index_component_type(num_vertices: int) -> int:
|
|
# 5123 = UNSIGNED_SHORT, 5125 = UNSIGNED_INT
|
|
return 5123 if num_vertices <= 65535 else 5125
|
|
|
|
|
|
def merge_all_meshes_to_single(doc: GltfDocument) -> bool:
|
|
if len(doc.meshes()) <= 1:
|
|
return True
|
|
|
|
mesh_transforms = compute_mesh_global_transforms(doc)
|
|
|
|
merged_vertices: List[np.ndarray] = []
|
|
merged_indices: List[np.ndarray] = []
|
|
vertex_offset = 0
|
|
|
|
for mesh_index, mesh in enumerate(doc.meshes()):
|
|
for primitive in mesh.get("primitives", []):
|
|
if "attributes" in primitive and "POSITION" in primitive["attributes"]:
|
|
pos_accessor_index = int(primitive["attributes"]["POSITION"])
|
|
vertices = access_data(doc, pos_accessor_index).astype(np.float32)
|
|
if mesh_index in mesh_transforms:
|
|
transform_matrix = mesh_transforms[mesh_index]
|
|
identity = np.eye(4, dtype=np.float32)
|
|
if not np.allclose(transform_matrix, identity, atol=1e-6):
|
|
pos_h = np.ones((vertices.shape[0], 4), dtype=np.float32)
|
|
pos_h[:, :3] = vertices
|
|
vertices = (pos_h @ transform_matrix.T)[:, :3].astype(np.float32)
|
|
|
|
if "indices" in primitive:
|
|
indices_accessor_index = int(primitive["indices"])
|
|
indices = access_data(doc, indices_accessor_index).astype(np.int64)
|
|
indices = indices + vertex_offset
|
|
else:
|
|
indices = np.arange(len(vertices), dtype=np.int64) + vertex_offset
|
|
indices = indices.reshape(-1, 3)
|
|
|
|
merged_vertices.append(vertices)
|
|
merged_indices.append(indices)
|
|
vertex_offset += len(vertices)
|
|
|
|
if not merged_vertices:
|
|
return False
|
|
|
|
final_vertices = np.concatenate(merged_vertices, axis=0)
|
|
final_indices = np.concatenate(merged_indices, axis=0)
|
|
|
|
acc_pos, _ = append_accessor_and_bufferview(
|
|
doc,
|
|
final_vertices.astype(np.float32),
|
|
component_type=5126,
|
|
element_type="VEC3",
|
|
target=34962,
|
|
)
|
|
|
|
index_component_type = _choose_index_component_type(final_vertices.shape[0])
|
|
index_dtype = np.uint16 if index_component_type == 5123 else np.uint32
|
|
acc_idx, _ = append_accessor_and_bufferview(
|
|
doc,
|
|
final_indices.astype(index_dtype).reshape(-1),
|
|
component_type=index_component_type,
|
|
element_type="SCALAR",
|
|
target=34963,
|
|
)
|
|
|
|
merged_mesh = {
|
|
"name": "Merged_Mesh",
|
|
"primitives": [
|
|
{
|
|
"attributes": {"POSITION": acc_pos},
|
|
"indices": acc_idx,
|
|
"mode": 4,
|
|
}
|
|
],
|
|
}
|
|
|
|
doc.json()["meshes"] = [merged_mesh]
|
|
|
|
update_nodes_for_merged_mesh(doc)
|
|
return True
|
|
|
|
|
|
def update_nodes_for_merged_mesh(doc: GltfDocument) -> None:
|
|
nodes = doc.nodes()
|
|
meshes = doc.meshes()
|
|
merged_node = {"name": "Merged_Mesh_Node", "mesh": 0}
|
|
nodes.append(merged_node)
|
|
merged_node_index = len(nodes) - 1
|
|
|
|
scene_index = doc.json().get("scene", 0)
|
|
scenes = doc.scenes()
|
|
if not scenes:
|
|
doc.json()["scenes"] = [{"nodes": [merged_node_index]}]
|
|
doc.json()["scene"] = 0
|
|
else:
|
|
scenes[scene_index].setdefault("nodes", [])
|
|
scenes[scene_index]["nodes"].append(merged_node_index)
|
|
|