diff --git a/mesh_processor/__init__.py b/mesh_processor/__init__.py new file mode 100644 index 0000000..7250bc4 --- /dev/null +++ b/mesh_processor/__init__.py @@ -0,0 +1,89 @@ +from .io_gltf import ( + GltfDocument, + load_gltf_or_glb, + save_glb, +) +from .accessors import ( + access_data, + update_accessor_binary_data, + append_accessor_and_bufferview, + recompute_accessor_min_max, +) +from .transforms import ( + compose_matrix, + quaternion_matrix, + direction_to_quaternion, + parse_node_transform, + compute_global_transform, + compute_mesh_global_transforms, + transform_positions_inplace, + set_node_trs_identity, + transform_local_vertices_to_world, +) +from .mesh_ops import ( + get_all_meshes_triangles, + merge_all_meshes_to_single, + update_nodes_for_merged_mesh, +) +from .rigging_ops import ( + check_input_file_compatibility, + clean_rigging_data_from_gltf, + clean_skin_data_from_gltf, + insert_generated_data_to_gltf, +) +from .validate import validate_before_save +from .mesh import ( + Mesh, + FastMesh, + TrimeshLike, + TrimeshLikeLoader, + PointCloud, +) + +__all__ = [ + "GltfDocument", + "load_gltf_or_glb", + "save_glb", + "access_data", + "update_accessor_binary_data", + "append_accessor_and_bufferview", + "recompute_accessor_min_max", + "compose_matrix", + "quaternion_matrix", + "direction_to_quaternion", + "parse_node_transform", + "compute_global_transform", + "compute_mesh_global_transforms", + "transform_positions_inplace", + "set_node_trs_identity", + "transform_local_vertices_to_world", + "get_all_meshes_triangles", + "merge_all_meshes_to_single", + "update_nodes_for_merged_mesh", + "check_input_file_compatibility", + "clean_rigging_data_from_gltf", + "clean_skin_data_from_gltf", + "insert_generated_data_to_gltf", + "validate_before_save", + "Mesh", + "FastMesh", + "TrimeshLike", + "TrimeshLikeLoader", + "PointCloud", +] + +try: + from .export_utils import ( + export_to_fastmesh, + export_to_mesh, + export_to_mesh_ultra_fast, + ) + __all__.extend([ + "export_to_fastmesh", + "export_to_mesh", + "export_to_mesh_ultra_fast", + ]) +except ImportError: + pass + + diff --git a/mesh_processor/accessors.py b/mesh_processor/accessors.py new file mode 100644 index 0000000..8889350 --- /dev/null +++ b/mesh_processor/accessors.py @@ -0,0 +1,206 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np + +from .io_gltf import GltfDocument, get_binary_data +from .utils import ensure_little_endian, as_uint8_buffer + + +_ELEMENT_SHAPES: Dict[str, Tuple[int, ...]] = { + "SCALAR": (1,), + "VEC2": (2,), + "VEC3": (3,), + "VEC4": (4,), + "MAT2": (2, 2), + "MAT3": (3, 3), + "MAT4": (4, 4), +} + +_ITEM_TYPES: Dict[int, Any] = { + 5120: np.int8, + 5121: np.uint8, + 5122: np.int16, + 5123: np.uint16, + 5125: np.uint32, + 5126: np.float32, +} + + +def _dtype_itemsize(component_type: int) -> int: + return np.dtype(_ITEM_TYPES[component_type]).itemsize + + +def _validate_stride(stride: int) -> None: + if stride != 0 and stride < 4: + raise ValueError("Stride is too small.") + if stride > 252: + raise ValueError("Stride is too big.") + + +def access_data(doc: GltfDocument, accessor_index: int) -> np.ndarray: + accessor = doc.accessors()[accessor_index] + buffer_view_index = accessor.get("bufferView") + if buffer_view_index is None: + raise NotImplementedError("Undefined buffer view") + + accessor_byte_offset = int(accessor.get("byteOffset", 0)) + component_type = int(accessor["componentType"]) # glTF enum + element_count = int(accessor["count"]) # number of elements + element_type = accessor["type"] + + if accessor.get("sparse") is not None: + raise NotImplementedError("Sparse accessors are not supported") + + buffer_view = doc.bufferViews()[buffer_view_index] + buffer_index = int(buffer_view["buffer"]) # which buffer + buffer_byte_length = int(buffer_view["byteLength"]) # length of this view + element_byte_offset = int(buffer_view.get("byteOffset", 0)) + element_byte_stride = int(buffer_view.get("byteStride", 0)) + _validate_stride(element_byte_stride) + + element_shape = _ELEMENT_SHAPES[element_type] + item_dtype = np.dtype(_ITEM_TYPES[component_type]) + item_count = int(np.prod(element_shape)) + item_size = item_dtype.itemsize + + size = element_count * item_count * item_size + if size > buffer_byte_length: + raise ValueError("Buffer did not have enough data for the accessor") + + buffers = doc.buffers() + if buffer_index >= len(buffers): + raise ValueError("Invalid buffer index") + + binary_data = get_binary_data(doc, buffer_index) + if len(binary_data) < buffers[buffer_index].get("byteLength", 0): + raise ValueError("Not enough binary data for the buffer") + + if element_byte_stride == 0: + element_byte_stride = item_size * item_count + if element_byte_stride < item_size * item_count: + raise ValueError("Items should not overlap") + + # force little-endian + item_dtype_str = np.dtype(item_dtype).newbyteorder("<").str + + dtype = np.dtype( + { + "names": ["element"], + "formats": [str(element_shape) + item_dtype_str], + "offsets": [0], + "itemsize": element_byte_stride, + } + ) + + byte_offset = accessor_byte_offset + element_byte_offset + if byte_offset % item_size != 0: + raise ValueError("Misaligned data") + byte_length = element_count * element_byte_stride + + view = binary_data[byte_offset : byte_offset + byte_length].view(dtype)["element"] + if element_type in ("MAT2", "MAT3", "MAT4"): + view = np.transpose(view, (0, 2, 1)) + return view + + +def update_accessor_binary_data(doc: GltfDocument, accessor_idx: int, new_data: np.ndarray) -> None: + accessor = doc.accessors()[accessor_idx] + buffer_view_idx = accessor.get("bufferView") + if buffer_view_idx is None: + raise NotImplementedError("Undefined buffer view") + + buffer_view = doc.bufferViews()[buffer_view_idx] + buffer_idx = int(buffer_view["buffer"]) # which buffer + byte_offset = int(buffer_view.get("byteOffset", 0)) + accessor_byte_offset = int(accessor.get("byteOffset", 0)) + + binary_data = doc.binary(buffer_idx) + # ensure little-endian float32 or appropriate component + new_bytes = ensure_little_endian(new_data).tobytes() + + start_pos = byte_offset + accessor_byte_offset + end_pos = start_pos + len(new_bytes) + + new_binary_data = binary_data.copy() + new_binary_data[start_pos:end_pos] = as_uint8_buffer(np.frombuffer(new_bytes, dtype=np.uint8)) + doc.set_binary(buffer_idx, new_binary_data) + doc.update_buffer_length(buffer_idx) + + +def recompute_accessor_min_max(doc: GltfDocument, accessor_idx: int) -> None: + accessor = doc.accessors()[accessor_idx] + if accessor.get("type") != "VEC3": + return + data = access_data(doc, accessor_idx) + accessor["min"] = data.min(axis=0).tolist() + accessor["max"] = data.max(axis=0).tolist() + + +def append_accessor_and_bufferview( + doc: GltfDocument, + array: np.ndarray, + component_type: int, + element_type: str, + target: Optional[int] = None, +) -> Tuple[int, int]: + """Append array bytes into buffer 0 and create bufferView + accessor. + + Returns: (accessor_index, buffer_view_index) + """ + if element_type not in _ELEMENT_SHAPES: + raise ValueError("Invalid element type") + if component_type not in _ITEM_TYPES: + raise ValueError("Invalid component type") + + array_le = ensure_little_endian(array, np.dtype(_ITEM_TYPES[component_type])) + element_shape = _ELEMENT_SHAPES[element_type] + + # validate last dimension matches element + num_components = int(np.prod(element_shape)) + if array_le.size % num_components != 0: + raise ValueError("Array size is not divisible by element components") + + if array_le.ndim == 1: + count = array_le.size // num_components + else: + count = array_le.shape[0] + + # append to binary buffer 0 + buf0 = doc.binary(0) + byte_offset = int(buf0.nbytes) + bytes_data = array_le.tobytes() + buf0_new = np.concatenate([buf0, np.frombuffer(bytes_data, dtype=np.uint8)]) + doc.set_binary(0, buf0_new) + + # create bufferView + bv_index = len(doc.bufferViews()) + byte_length = len(bytes_data) + bufferview: Dict[str, Any] = { + "buffer": 0, + "byteOffset": byte_offset, + "byteLength": byte_length, + } + if target is not None: + bufferview["target"] = target + doc.bufferViews().append(bufferview) + + # create accessor + acc_index = len(doc.accessors()) + accessor: Dict[str, Any] = { + "bufferView": bv_index, + "componentType": int(component_type), + "count": int(count), + "type": element_type, + } + # if POSITION like, store min/max + if element_type == "VEC3": + arr2d = array_le.reshape((-1, 3)) + accessor["min"] = arr2d.min(axis=0).tolist() + accessor["max"] = arr2d.max(axis=0).tolist() + doc.accessors().append(accessor) + + doc.update_buffer_length(0) + return acc_index, bv_index + diff --git a/mesh_processor/gltf_ops.py b/mesh_processor/gltf_ops.py new file mode 100644 index 0000000..549e55f --- /dev/null +++ b/mesh_processor/gltf_ops.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Tuple + +import numpy as np + +from .io_gltf import GltfDocument, load_gltf_or_glb, save_glb +from .mesh_ops import get_all_meshes_triangles, merge_all_meshes_to_single +from .transforms import ( + compose_matrix, + quaternion_matrix, + direction_to_quaternion, + compute_mesh_global_transforms, + transform_local_vertices_to_world, + preprocess_transforms_inplace, +) +from .rigging_ops import ( + check_input_file_compatibility, + clean_rigging_data_from_gltf, + insert_generated_data_to_gltf, +) +from .validate import validate_before_save + + +__all__ = [ + "GltfDocument", + "load_gltf_or_glb", + "save_glb", + "get_all_meshes_triangles", + "merge_all_meshes_to_single", + "compose_matrix", + "quaternion_matrix", + "direction_to_quaternion", + "compute_mesh_global_transforms", + "transform_local_vertices_to_world", + "preprocess_transforms_inplace", + "check_input_file_compatibility", + "clean_rigging_data_from_gltf", + "insert_generated_data_to_gltf", + "validate_before_save", +] + diff --git a/mesh_processor/io_gltf.py b/mesh_processor/io_gltf.py new file mode 100644 index 0000000..baf8d05 --- /dev/null +++ b/mesh_processor/io_gltf.py @@ -0,0 +1,186 @@ +from __future__ import annotations + +import json +import struct +from base64 import b64decode +from typing import Any, BinaryIO, Dict, List, Optional, Tuple + +import numpy as np + +from .utils import NpEncoder, pad_to_4_bytes, warn + + +GLTF_MAGIC = 0x46546C67 +JSON_CHUNK_TYPE = 0x4E4F534A +BINARY_CHUNK_TYPE = 0x004E4942 +DATA_URI_PREFIX = "data:application/octet-stream;base64," + + +def read_header(stream: BinaryIO) -> Optional[Tuple[int, int]]: + header = stream.read(12) + if len(header) != 12: + return None + magic, version, length = struct.unpack(" Optional[Tuple[Dict[str, Any], np.ndarray]]: + json_data: Optional[Dict[str, Any]] = None + binary_data: Optional[np.ndarray] = None + + while stream.tell() < length: + chunk_header = stream.read(8) + if len(chunk_header) != 8: + warn(f"Invalid chunk header size: {len(chunk_header)} bytes (expected 8)") + break + chunk_length, chunk_type = struct.unpack(" Dict[str, Any]: + return self._json + + def buffers(self) -> List[Dict[str, Any]]: + return self._json.setdefault("buffers", []) + + def bufferViews(self) -> List[Dict[str, Any]]: + return self._json.setdefault("bufferViews", []) + + def accessors(self) -> List[Dict[str, Any]]: + return self._json.setdefault("accessors", []) + + def nodes(self) -> List[Dict[str, Any]]: + return self._json.setdefault("nodes", []) + + def meshes(self) -> List[Dict[str, Any]]: + return self._json.setdefault("meshes", []) + + def scenes(self) -> List[Dict[str, Any]]: + return self._json.setdefault("scenes", []) + + def materials(self) -> List[Dict[str, Any]]: + return self._json.setdefault("materials", []) + + def images(self) -> List[Dict[str, Any]]: + return self._json.setdefault("images", []) + + def textures(self) -> List[Dict[str, Any]]: + return self._json.setdefault("textures", []) + + def samplers(self) -> List[Dict[str, Any]]: + return self._json.setdefault("samplers", []) + + def binary(self, index: int = 0) -> np.ndarray: + return self._bin.setdefault(index, np.array([], dtype=np.uint8)) + + def set_binary(self, index: int, data: np.ndarray) -> None: + self._bin[index] = data + + def update_buffer_length(self, buffer_index: int = 0) -> None: + if not self.buffers(): + self._json["buffers"] = [{"byteLength": int(self.binary(buffer_index).nbytes)}] + else: + self.buffers()[buffer_index]["byteLength"] = int(self.binary(buffer_index).nbytes) + + +def get_binary_data(doc: GltfDocument, buffer_index: int) -> np.ndarray: + """Return binary data for a buffer index. + + Supports embedded base64 buffers in GLTF as well as GLB's main buffer. + """ + buffers = doc.buffers() + if buffer_index >= len(buffers): + raise IndexError("Buffer index out of range") + binary = doc._bin.get(buffer_index) + if binary is not None and binary.size > 0: + return binary + uri = buffers[buffer_index].get("uri") + if uri and uri.startswith(DATA_URI_PREFIX): + binary_data = b64decode(uri[len(DATA_URI_PREFIX) :]) + arr = np.frombuffer(binary_data, dtype=np.uint8) + doc.set_binary(buffer_index, arr) + return arr + # if no uri, return existing buffer (may be empty for GLB until provided) + return doc.binary(buffer_index) + + +def load_gltf_or_glb(path: str) -> GltfDocument: + with open(path, "rb") as f: + stream = f + version_and_length = read_header(stream) + if version_and_length is None: # GLTF JSON + stream.seek(0) + json_data = json.load(stream) + # embedded or external buffers will be resolved lazily + return GltfDocument(json_data=json_data) + version, length = version_and_length + if version != 2: + raise ValueError(f"Only GLB version 2 is supported, found version {version}") + json_data, binary_data = read_chunks(stream, length) + return GltfDocument(json_data=json_data, binary_buffers={0: binary_data}) + + +def save_glb(doc: GltfDocument, stream: BinaryIO) -> None: + """Write GLB v2 with proper 4-byte alignment for JSON and BIN chunks.""" + # ensure buffers length is correct for buffer 0 + doc.update_buffer_length(0) + + json_bytes = json.dumps(doc.json(), cls=NpEncoder, separators=(",", ":")).encode("utf-8") + json_padded = pad_to_4_bytes(json_bytes, b"\x20") + + bin_bytes = bytes(doc.binary(0)) + bin_padded = pad_to_4_bytes(bin_bytes, b"\x00") + + version = 2 + file_header_size = 12 + chunk_header_size = 8 + file_length = ( + file_header_size + + chunk_header_size + + len(json_padded) + + chunk_header_size + + len(bin_padded) + ) + + # write GLB header + stream.write(struct.pack(" 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) + diff --git a/mesh_processor/rigging_ops.py b/mesh_processor/rigging_ops.py new file mode 100644 index 0000000..1c9abe7 --- /dev/null +++ b/mesh_processor/rigging_ops.py @@ -0,0 +1,340 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np + +from .io_gltf import GltfDocument +from .transforms import ( + compute_mesh_global_transforms, + compose_matrix, + direction_to_quaternion, +) +from .accessors import append_accessor_and_bufferview + +_EPS = np.finfo(float).eps * 4.0 + + +def check_input_file_compatibility(doc: GltfDocument) -> Dict[str, Any]: + info = { + "has_skins": False, + "has_animations": False, + "skin_count": 0, + "animation_count": 0, + "issues": [], + "recommendations": [], + "mesh_transforms_non_identity": [], + } + + if "skins" in doc.json(): + info["has_skins"] = True + info["skin_count"] = len(doc.json()["skins"]) + info["issues"].append(f"File already contains {info['skin_count']} skin(s)") + info["recommendations"].append("Remove existing skin data before auto-rigging") + + if "animations" in doc.json(): + info["has_animations"] = True + info["animation_count"] = len(doc.json()["animations"]) + info["issues"].append( + f"File already contains {info['animation_count']} animation(s)" + ) + info["recommendations"].append( + "Existing animations may conflict with generated rig" + ) + + mesh_transforms = compute_mesh_global_transforms(doc) + for mesh_index, transform in mesh_transforms.items(): + identity = np.eye(4, dtype=np.float32) + if not np.allclose(transform, identity, atol=1e-6): + mesh_name = doc.meshes()[mesh_index].get("name", f"Mesh_{mesh_index}") + info["mesh_transforms_non_identity"].append( + {"mesh_index": mesh_index, "mesh_name": mesh_name, "transform": transform.tolist()} + ) + + if info["mesh_transforms_non_identity"]: + info["issues"].append( + f"{len(info['mesh_transforms_non_identity'])} mesh(es) have non-identity transformations" + ) + info["recommendations"].append( + "Mesh positions will be transformed to global space for auto-rigging" + ) + return info + + +def clean_rigging_data_from_gltf(doc: GltfDocument) -> Dict[str, Any]: + cleanup_info = { + "removed_skins": 0, + "removed_animations": 0, + "removed_skin_refs": 0, + "removed_joint_attrs": 0, + "removed_weight_attrs": 0, + "removed_accessors": 0, + "removed_buffer_views": 0, + "binary_buffer_rebuilt": False, + "operations": [], + } + + # collect accessors to remove + joint_weight_accessors = set() + for mesh in doc.meshes(): + for primitive in mesh.get("primitives", []): + attrs = primitive.get("attributes", {}) + if "JOINTS_0" in attrs: + joint_weight_accessors.add(int(attrs["JOINTS_0"])) + cleanup_info["removed_joint_attrs"] += 1 + if "WEIGHTS_0" in attrs: + joint_weight_accessors.add(int(attrs["WEIGHTS_0"])) + cleanup_info["removed_weight_attrs"] += 1 + + inverse_bind_accessors = set() + if "skins" in doc.json(): + for skin in doc.json()["skins"]: + if "inverseBindMatrices" in skin: + inverse_bind_accessors.add(int(skin["inverseBindMatrices"])) + + all_skin_accessors = joint_weight_accessors | inverse_bind_accessors + + skin_buffer_views = set() + for accessor_idx in all_skin_accessors: + if accessor_idx < len(doc.accessors()): + accessor = doc.accessors()[accessor_idx] + if "bufferView" in accessor: + skin_buffer_views.add(int(accessor["bufferView"])) + + # remove animations + if "animations" in doc.json(): + cleanup_info["removed_animations"] = len(doc.json()["animations"]) + del doc.json()["animations"] + cleanup_info["operations"].append( + f"Removed 'animations' section with {cleanup_info['removed_animations']} elements" + ) + + # remove skins + if "skins" in doc.json(): + cleanup_info["removed_skins"] = len(doc.json()["skins"]) + del doc.json()["skins"] + cleanup_info["operations"].append( + f"Removed 'skins' section with {cleanup_info['removed_skins']} elements" + ) + + # remove node skin refs + for idx, node in enumerate(doc.nodes()): + if "skin" in node: + node.pop("skin") + cleanup_info["removed_skin_refs"] += 1 + cleanup_info["operations"].append(f"Removed skin reference from node index {idx}") + + # remove JOINTS_0/WEIGHTS_0 attributes + for m_idx, mesh in enumerate(doc.meshes()): + for p_idx, primitive in enumerate(mesh.get("primitives", [])): + attrs = primitive.get("attributes", {}) + removed = [] + if "JOINTS_0" in attrs: + attrs.pop("JOINTS_0") + removed.append("JOINTS_0") + if "WEIGHTS_0" in attrs: + attrs.pop("WEIGHTS_0") + removed.append("WEIGHTS_0") + if removed: + cleanup_info["operations"].append( + f"Removed attributes {removed} from mesh {m_idx} primitive {p_idx}" + ) + + # rebuild accessors without skin ones + if all_skin_accessors: + old_accessors = doc.accessors() + new_accessors = [] + old_to_new = {} + for i, acc in enumerate(old_accessors): + if i in all_skin_accessors: + cleanup_info["removed_accessors"] += 1 + else: + old_to_new[i] = len(new_accessors) + new_accessors.append(acc) + doc.json()["accessors"] = new_accessors + + # remap references + for mesh in doc.meshes(): + for primitive in mesh.get("primitives", []): + attrs = primitive.get("attributes", {}) + for key, acc_idx in list(attrs.items()): + if acc_idx in old_to_new: + attrs[key] = old_to_new[acc_idx] + else: + # invalid -> drop attribute + attrs.pop(key) + if "indices" in primitive: + idx = primitive["indices"] + if idx in old_to_new: + primitive["indices"] = old_to_new[idx] + else: + primitive.pop("indices", None) + + # bufferViews cleanup requires rebuilding binary; keep it simple for now + # caller may opt to rebuild or leave holes; we just report counts + cleanup_info["removed_buffer_views"] = len(skin_buffer_views) + return cleanup_info + + +def clean_skin_data_from_gltf(doc: GltfDocument) -> Dict[str, Any]: + return clean_rigging_data_from_gltf(doc) + + +@dataclass +class JointData: + joint_node_index: int + joint_name: str + parent_index: Optional[int] = None + children_indices: Tuple[int, ...] = () + forward_bind_matrix: np.ndarray = field(default_factory=lambda: np.eye(4, dtype=np.float32)) + inverse_bind_matrix: np.ndarray = field(default_factory=lambda: np.eye(4, dtype=np.float32)) + local_transform_matrix: np.ndarray = field(default_factory=lambda: np.eye(4, dtype=np.float32)) + original_index: Optional[int] = None + + +def insert_generated_data_to_gltf( + doc: GltfDocument, + joints: np.ndarray, + conns: np.ndarray, + skins: np.ndarray, + all_vertices_groups: List[Dict[str, Any]], +) -> None: + # A) collect used joints and build remap + used = set() + for v_group in all_vertices_groups: + v0, v1 = v_group["vertices_range"] + skins_data = skins[v0:v1] + top_idx = np.argsort(-skins_data, axis=1)[:, :4] + rows = np.arange(skins_data.shape[0])[:, None] + top_w = skins_data[rows, top_idx] + mask = top_w > _EPS + used.update(top_idx[mask].tolist()) + used_sorted = sorted(used) + old_to_new = {old: new for new, old in enumerate(used_sorted)} + + # B) build joint nodes and skin + current_nodes_len = len(doc.nodes()) + joints_nodes_indices: List[int] = [] + all_joints_data: List[JointData] = [] + joints_global_pos: List[np.ndarray] = [] + for new_joint_idx, old_joint_idx in enumerate(used_sorted): + parent_index = conns[old_joint_idx] if old_joint_idx < len(conns) else old_joint_idx + joint_node_index = current_nodes_len + new_joint_idx + joint_name = f"Joint_{new_joint_idx}" + jd = JointData(joint_node_index, joint_name) + jd.original_index = int(old_joint_idx) + joints_nodes_indices.append(joint_node_index) + all_joints_data.append(jd) + parent_pick = int(parent_index) if parent_index < len(joints) else int(old_joint_idx) + joints_global_pos.append(joints[parent_pick]) + + # build parent/children links + for jd in all_joints_data: + old_idx = int(jd.original_index) + parent_old = int(conns[old_idx]) if old_idx < len(conns) else old_idx + if parent_old == old_idx: + continue + if parent_old in old_to_new: + parent_new = old_to_new[parent_old] + parent_node_index = all_joints_data[parent_new].joint_node_index + jd.parent_index = parent_node_index + parent_children = list(all_joints_data[parent_new].children_indices) + if jd.joint_node_index not in parent_children: + parent_children.append(jd.joint_node_index) + all_joints_data[parent_new].children_indices = tuple(parent_children) + + # compute matrices + inverse_bind_mats: List[np.ndarray] = [] + for i_name, jd in enumerate(all_joints_data): + global_scale = np.array([1.0, 1.0, 1.0], dtype=np.float32) + global_translation = np.asarray(joints_global_pos[i_name], dtype=np.float32) + if jd.children_indices: + offset = np.zeros(3, dtype=np.float32) + for child_index in jd.children_indices: + offset += ( + joints_global_pos[child_index - current_nodes_len] - global_translation + ) + global_rotation = direction_to_quaternion(offset) + elif jd.parent_index is not None: + offset_from_parent = ( + global_translation + - joints_global_pos[jd.parent_index - current_nodes_len] + ) + global_rotation = direction_to_quaternion(offset_from_parent) + else: + global_rotation = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float32) + forward = compose_matrix(global_scale, global_rotation, global_translation) + jd.forward_bind_matrix = forward + jd.inverse_bind_matrix = np.linalg.inv(forward) + inverse_bind_mats.append(jd.inverse_bind_matrix) + + inverse_bind_mats = np.stack(inverse_bind_mats, axis=0).astype(np.float32) + # glTF expects column-major arrays in JSON; we transpose like original code + inverse_bind_mats = np.transpose(inverse_bind_mats, (0, 2, 1)).copy(order="C") + + # C) add vertex joint/weight attributes per groups + for v_group in all_vertices_groups: + v0, v1 = v_group["vertices_range"] + skins_data = skins[v0:v1] + top_idx = np.argsort(-skins_data, axis=1)[:, :4] + rows = np.arange(skins_data.shape[0])[:, None] + weights_data = skins_data[rows, top_idx].astype(np.float32) + + joints_data = np.zeros_like(top_idx, dtype=np.uint16) + for i_row in range(top_idx.shape[0]): + for i_col in range(4): + old = int(top_idx[i_row, i_col]) + if old in old_to_new: + joints_data[i_row, i_col] = np.uint16(old_to_new[old]) + else: + joints_data[i_row, i_col] = np.uint16(0) + weights_data[i_row, i_col] = 0.0 + + # D) normalize weights + sums = np.sum(weights_data, axis=1, keepdims=True) + mask_zero = sums < _EPS + weights_data = np.where(mask_zero, 0.25, weights_data / sums) + + mesh_index = v_group["mesh_index"] + prim_index = v_group["primitive_index"] + primitive = doc.meshes()[mesh_index]["primitives"][prim_index] + + acc_j, _ = append_accessor_and_bufferview( + doc, joints_data.astype(np.uint16), component_type=5123, element_type="VEC4", target=34962 + ) + acc_w, _ = append_accessor_and_bufferview( + doc, weights_data.astype(np.float32), component_type=5126, element_type="VEC4", target=34962 + ) + primitive.setdefault("attributes", {}) + primitive["attributes"]["JOINTS_0"] = acc_j + primitive["attributes"]["WEIGHTS_0"] = acc_w + + # skin JSON + acc_ibm, _ = append_accessor_and_bufferview( + doc, inverse_bind_mats.astype(np.float32), component_type=5126, element_type="MAT4", target=None + ) + doc.json()["skins"] = [ + {"joints": joints_nodes_indices, "inverseBindMatrices": acc_ibm, "name": "Generated_Skin"} + ] + + # add joint nodes + root_joint_nodes: List[int] = [] + for jd in all_joints_data: + node_json = { + "name": jd.joint_name, + "matrix": np.transpose(jd.local_transform_matrix, (1, 0)).copy(order="C").flatten().tolist(), + } + if jd.children_indices: + node_json["children"] = list(jd.children_indices) + doc.nodes().append(node_json) + if jd.parent_index is None: + root_joint_nodes.append(jd.joint_node_index) + + # joints root + doc.nodes().append({"name": "Joints_Root", "children": root_joint_nodes}) + joints_root_index = len(doc.nodes()) - 1 + scene_index = doc.json().get("scene", 0) + doc.scenes()[scene_index].setdefault("nodes", []) + doc.scenes()[scene_index]["nodes"].append(joints_root_index) + diff --git a/mesh_processor/transforms.py b/mesh_processor/transforms.py new file mode 100644 index 0000000..f6d7921 --- /dev/null +++ b/mesh_processor/transforms.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +import numpy as np + +from .io_gltf import GltfDocument +from .accessors import access_data, update_accessor_binary_data, recompute_accessor_min_max + + +def quaternion_matrix(quaternion: np.ndarray) -> np.ndarray: + x, y, z, w = quaternion + x2 = x * x + y2 = y * y + z2 = z * z + xy = x * y + xz = x * z + yz = y * z + wx = w * x + wy = w * y + wz = w * z + return np.array( + [ + [1.0 - 2.0 * (y2 + z2), 2.0 * (xy - wz), 2.0 * (xz + wy), 0.0], + [2.0 * (xy + wz), 1.0 - 2.0 * (x2 + z2), 2.0 * (yz - wx), 0.0], + [2.0 * (xz - wy), 2.0 * (yz + wx), 1.0 - 2.0 * (x2 + y2), 0.0], + [0.0, 0.0, 0.0, 1.0], + ] + ) + + +def compose_matrix(scale: np.ndarray, rotation: np.ndarray, translation: np.ndarray) -> np.ndarray: + S = np.diag([scale[0], scale[1], scale[2], 1.0]) + R = quaternion_matrix(rotation) + T = np.eye(4) + T[:3, 3] = translation + return T @ (R @ S) + + +def direction_to_quaternion(direction: np.ndarray) -> np.ndarray: + if not np.issubdtype(direction.dtype, np.floating): + direction = direction.astype(np.float32) + norm = np.linalg.norm(direction) + if norm < 1e-10: + return np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float32) + direction = direction / norm + forward = np.array([0.0, 1.0, 0.0], dtype=np.float32) + up = np.array([0.0, 0.0, 1.0], dtype=np.float32) + dot = float(np.clip(np.dot(forward, direction), -1.0, 1.0)) + if dot > 0.9999: + return np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float32) + if dot < -0.9999: + return np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32) + axis = np.cross(forward, direction) + axis_norm = np.linalg.norm(axis) + if axis_norm < 1e-10: + if abs(direction[2]) < 0.9: + axis = np.cross(up, direction) + else: + axis = np.cross(np.array([1.0, 0.0, 0.0], dtype=np.float32), direction) + axis = axis / np.linalg.norm(axis) + else: + axis = axis / axis_norm + angle = np.arccos(dot) + half = angle * 0.5 + s = np.sin(half) + c = np.cos(half) + quat = np.array([axis[0] * s, axis[1] * s, axis[2] * s, c], dtype=np.float32) + return quat / np.linalg.norm(quat) + + +def parse_node_transform(node_json: Dict[str, Any]) -> np.ndarray: + if "matrix" in node_json: + m = np.array(node_json["matrix"], dtype=np.float32).reshape(4, 4) + return m.T + t = np.array(node_json.get("translation", [0.0, 0.0, 0.0]), dtype=np.float32) + r = np.array(node_json.get("rotation", [0.0, 0.0, 0.0, 1.0]), dtype=np.float32) + s = np.array(node_json.get("scale", [1.0, 1.0, 1.0]), dtype=np.float32) + return compose_matrix(s, r, t) + + +def compute_global_transform(doc: GltfDocument, node_index: int, parent_transform: Optional[np.ndarray] = None) -> np.ndarray: + node_json = doc.nodes()[node_index] + local = parse_node_transform(node_json) + return parent_transform @ local if parent_transform is not None else local + + +def compute_mesh_global_transforms(doc: GltfDocument) -> Dict[int, np.ndarray]: + mesh_to_transform: Dict[int, np.ndarray] = {} + scene_index = doc.json().get("scene") + if scene_index is None: + return mesh_to_transform + root_nodes = doc.scenes()[scene_index].get("nodes", []) + + def traverse(node_idx: int, parent: Optional[np.ndarray] = None) -> None: + global_m = compute_global_transform(doc, node_idx, parent) + node = doc.nodes()[node_idx] + if "mesh" in node: + mesh_to_transform[int(node["mesh"])] = global_m + for child in node.get("children", []) or []: + traverse(int(child), global_m) + + for root in root_nodes: + traverse(int(root)) + return mesh_to_transform + + +def transform_positions_inplace(doc: GltfDocument, mesh_idx: int, matrix: np.ndarray) -> None: + mesh = doc.meshes()[mesh_idx] + for primitive in mesh.get("primitives", []): + attrs = primitive.get("attributes", {}) + if "POSITION" not in attrs: + continue + acc_idx = int(attrs["POSITION"]) + data = access_data(doc, acc_idx) + pos_h = np.ones((data.shape[0], 4), dtype=np.float32) + pos_h[:, :3] = data.astype(np.float32) + transformed = (pos_h @ matrix.T)[:, :3].astype(np.float32) + update_accessor_binary_data(doc, acc_idx, transformed) + recompute_accessor_min_max(doc, acc_idx) + + +def transform_local_vertices_to_world(pos_data: np.ndarray, node: Dict[str, Any]) -> np.ndarray: + if "matrix" in node: + M = np.array(node["matrix"], dtype=np.float32).reshape(4, 4).T + else: + t = np.array(node.get("translation", [0.0, 0.0, 0.0]), dtype=np.float32) + r = np.array(node.get("rotation", [0.0, 0.0, 0.0, 1.0]), dtype=np.float32) + s = np.array(node.get("scale", [1.0, 1.0, 1.0]), dtype=np.float32) + M = compose_matrix(s, r, t) + pos_h = np.concatenate([pos_data, np.ones((pos_data.shape[0], 1), dtype=np.float32)], axis=1) + out = (M @ pos_h.T).T[:, :3] + return out + + +def set_node_trs_identity(doc: GltfDocument, node_idx: int) -> None: + node = doc.nodes()[node_idx] + for key in ["matrix", "translation", "rotation", "scale"]: + if key in node: + node.pop(key) + + +def preprocess_transforms_inplace(doc: GltfDocument) -> None: + """Apply node transforms to vertex positions and set TRS for all nodes. + + Idea: for each mesh, we calculate the global matrix, transform the POSITION of all + primitives of this mesh and remove the transform for the nodes referring to this mesh. + """ + mesh_transforms = compute_mesh_global_transforms(doc) + # quick exit if all identity + any_non_identity = any( + not np.allclose(M, np.eye(4), atol=1e-6) for M in mesh_transforms.values() + ) + if not any_non_identity: + return + + for mesh_idx, M in mesh_transforms.items(): + if np.allclose(M, np.eye(4), atol=1e-6): + continue + # transform all positions in the mesh + mesh = doc.meshes()[mesh_idx] + for primitive in mesh.get("primitives", []): + attrs = primitive.get("attributes", {}) + if "POSITION" not in attrs: + continue + acc_idx = int(attrs["POSITION"]) + data = access_data(doc, acc_idx) + pos_h = np.ones((data.shape[0], 4), dtype=np.float32) + pos_h[:, :3] = data.astype(np.float32) + transformed = (pos_h @ M.T)[:, :3].astype(np.float32) + update_accessor_binary_data(doc, acc_idx, transformed) + recompute_accessor_min_max(doc, acc_idx) + + # reset TRS for all nodes referring to this mesh + for node_idx, node in enumerate(doc.nodes()): + if node.get("mesh") == mesh_idx: + set_node_trs_identity(doc, node_idx) + diff --git a/mesh_processor/utils.py b/mesh_processor/utils.py new file mode 100644 index 0000000..e941418 --- /dev/null +++ b/mesh_processor/utils.py @@ -0,0 +1,63 @@ +import json +import warnings +from typing import Any, Optional + +import numpy as np + + +def pad_to_4_bytes(data: bytes, pad_byte: bytes = b"\x20") -> bytes: + """Pad given bytes to the next 4-byte boundary. + + JSON chunks must be padded with spaces (0x20), + binary chunks must be padded with zeros (0x00). + """ + if not data: + return data + remainder = len(data) % 4 + if remainder == 0: + return data + return data + pad_byte * (4 - remainder) + + +def ensure_little_endian(array: np.ndarray, dtype: Optional[np.dtype] = None) -> np.ndarray: + """Return an array with little-endian dtype (copy if needed).""" + if dtype is not None: + array = array.astype(dtype, copy=False) + if array.dtype.byteorder in (">", "|"): + return array.byteswap().newbyteorder("<") + if array.dtype.byteorder == "<": + return array + # native without explicit endianness + return array.astype(array.dtype.newbyteorder("<"), copy=False) + + +def as_uint8_buffer(array: np.ndarray) -> np.ndarray: + """View array bytes as uint8 vector.""" + if not isinstance(array, np.ndarray): + raise TypeError("as_uint8_buffer expects a numpy array") + return np.frombuffer(array.tobytes(), dtype=np.uint8) + + +def warn(message: str) -> None: + warnings.warn(message) + + +class NpEncoder(json.JSONEncoder): + """JSON encoder that converts numpy scalars/arrays to native Python types. + + Important: cast numpy integers to int, not to string. + """ + + def default(self, obj: Any) -> Any: + # scalars + if isinstance(obj, (np.integer,)): + return int(obj) + if isinstance(obj, (np.floating,)): + return float(obj) + if isinstance(obj, (np.bool_,)): + return bool(obj) + # arrays + if isinstance(obj, np.ndarray): + return obj.tolist() + return super().default(obj) + diff --git a/mesh_processor/validate.py b/mesh_processor/validate.py new file mode 100644 index 0000000..b991d4c --- /dev/null +++ b/mesh_processor/validate.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +from typing import Dict, List, Tuple + +import numpy as np + +from .io_gltf import GltfDocument + + +def validate_before_save(doc: GltfDocument) -> List[str]: + errors: List[str] = [] + # buffers length and 4-byte alignment will be handled by writer, but sanity-check + for i, buf in enumerate(doc.buffers()): + expect = int(doc.binary(i).nbytes) + if buf.get("byteLength", -1) != expect: + errors.append(f"Buffer {i} byteLength mismatch: json={buf.get('byteLength')} actual={expect}") + # bufferViews within ranges + for i, bv in enumerate(doc.bufferViews()): + b = int(bv.get("buffer", 0)) + off = int(bv.get("byteOffset", 0)) + ln = int(bv.get("byteLength", 0)) + if off + ln > int(doc.binary(b).nbytes): + errors.append(f"bufferView {i} out of range: {off}+{ln}>{doc.binary(b).nbytes}") + if off % 4 != 0: + errors.append(f"bufferView {i} is not 4-byte aligned: offset={off}") + # accessors consistency + for i, acc in enumerate(doc.accessors()): + if "bufferView" not in acc: + continue + bv = doc.bufferViews()[acc["bufferView"]] + component_type = acc.get("componentType") + acc_type = acc.get("type") + if component_type not in (5120, 5121, 5122, 5123, 5125, 5126): + errors.append(f"accessor {i} invalid componentType {component_type}") + if acc_type not in ("SCALAR", "VEC2", "VEC3", "VEC4", "MAT2", "MAT3", "MAT4"): + errors.append(f"accessor {i} invalid type {acc_type}") + # position min/max + for i, acc in enumerate(doc.accessors()): + if acc.get("type") == "VEC3": + if "min" not in acc or "max" not in acc: + errors.append(f"accessor {i} missing min/max for VEC3") + # primitives mode + for m_idx, mesh in enumerate(doc.meshes()): + for p_idx, prim in enumerate(mesh.get("primitives", [])): + if prim.get("mode", 4) != 4: + errors.append(f"mesh {m_idx} prim {p_idx} mode must be 4 (TRIANGLES)") + return errors +