diff --git a/Gen_3D_Modules/CRM/model/crm/model.py b/Gen_3D_Modules/CRM/model/crm/model.py index fa4ca39..27ba4c9 100644 --- a/Gen_3D_Modules/CRM/model/crm/model.py +++ b/Gen_3D_Modules/CRM/model/crm/model.py @@ -14,7 +14,7 @@ from CRM.util.renderer import Renderer from CRM.model.archs.mlp_head import SdfMlp, RgbMlp import xatlas -from mesh_processer.mesh import Mesh +from mesh_processor.mesh import Mesh class Dummy: diff --git a/Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/DifferentiableRenderer/__init__.py b/Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/DifferentiableRenderer/__init__.py index b7ca480..67a1765 100644 --- a/Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/DifferentiableRenderer/__init__.py +++ b/Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/DifferentiableRenderer/__init__.py @@ -6,6 +6,6 @@ except ImportError: try: from .mesh_inpaint_processor_fallback import meshVerticeInpaint print("[WARNING] Using Python fallback mesh_inpaint_processor (slower)") - print("[INFO] To use faster C++ version, run: compile_mesh_painter.bat") + print("[INFO] To use faster C++ version, run: compile_mesh_painter.bat or compile_mesh_painter.sh") except ImportError: print("[ERROR] Neither C++ nor Python mesh_inpaint_processor found!") diff --git a/Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/mesh_inpaint_wheel/mesh_inpaint_processor/__init__.py b/Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/mesh_inpaint_wheel/mesh_inpaint_processor/__init__.py index b7ca480..67a1765 100644 --- a/Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/mesh_inpaint_wheel/mesh_inpaint_processor/__init__.py +++ b/Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/mesh_inpaint_wheel/mesh_inpaint_processor/__init__.py @@ -6,6 +6,6 @@ except ImportError: try: from .mesh_inpaint_processor_fallback import meshVerticeInpaint print("[WARNING] Using Python fallback mesh_inpaint_processor (slower)") - print("[INFO] To use faster C++ version, run: compile_mesh_painter.bat") + print("[INFO] To use faster C++ version, run: compile_mesh_painter.bat or compile_mesh_painter.sh") except ImportError: print("[ERROR] Neither C++ nor Python mesh_inpaint_processor found!") diff --git a/Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/mesh_inpaint_wheel/setup.py b/Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/mesh_inpaint_wheel/setup.py index b72b316..ae470b1 100644 --- a/Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/mesh_inpaint_wheel/setup.py +++ b/Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/mesh_inpaint_wheel/setup.py @@ -1,7 +1,6 @@ from setuptools import setup, find_packages import platform -# Определяем платформу для правильного именования wheel if platform.system() == "Windows": plat_name = "win_amd64" elif platform.system() == "Darwin": @@ -13,14 +12,14 @@ setup( name="mesh_inpaint_processor", version="1.0.0", description="Fast C++ mesh inpainting processor with Python fallback", - long_description="Компилированный модуль для быстрой обработки инпейнтинга мешей", + long_description="Compiled module for fast mesh inpainting processing", author="Hunyuan3D Team", packages=find_packages(), package_data={ 'mesh_inpaint_processor': [ - '*.pyd', # Включаем .pyd файлы - '*.so', # На случай Linux версии - '*.py' # Python fallback файлы + '*.pyd', + '*.so', # For Linux version + '*.py' # Python fallback files ] }, include_package_data=True, @@ -41,7 +40,6 @@ setup( "Operating System :: POSIX :: Linux", "Operating System :: MacOS", ], - # Указываем что это binary distribution zip_safe=False, - has_ext_modules=lambda: True, # Указывает что есть нативные модули + has_ext_modules=lambda: True, ) \ No newline at end of file diff --git a/Gen_3D_Modules/LGM/nerf_marching_cubes_converter.py b/Gen_3D_Modules/LGM/nerf_marching_cubes_converter.py index 3d7e9bd..978ec85 100644 --- a/Gen_3D_Modules/LGM/nerf_marching_cubes_converter.py +++ b/Gen_3D_Modules/LGM/nerf_marching_cubes_converter.py @@ -15,8 +15,8 @@ from kiui.op import inverse_sigmoid from kiui.cam import orbit_camera, get_perspective from kiui.nn import MLP, trunc_exp -from mesh_processer.mesh import Mesh -from mesh_processer.mesh_utils import marching_cubes_density_to_mesh, color_func_to_albedo +from mesh_processor.mesh import Mesh +from mesh_processor.mesh_utils import marching_cubes_density_to_mesh, color_func_to_albedo from LGM.core.options import Options from LGM.core.gs import GaussianRenderer from LGM.core.utils import get_rays diff --git a/Gen_3D_Modules/MV_Adapter/mvadapter_node_utils.py b/Gen_3D_Modules/MV_Adapter/mvadapter_node_utils.py index 3a98854..d846ee2 100644 --- a/Gen_3D_Modules/MV_Adapter/mvadapter_node_utils.py +++ b/Gen_3D_Modules/MV_Adapter/mvadapter_node_utils.py @@ -540,7 +540,7 @@ def prepare_tg2mv_pipeline( if hasattr(pipe, 'enable_model_cpu_offload'): if not use_mmgp: pipe.enable_model_cpu_offload() - print("[INFO] Включен CPU offload для экономии VRAM") + print("[INFO] CPU offload for memory optimization") if use_mmgp and MMGP_AVAILABLE: torch.cuda.empty_cache() diff --git a/Gen_3D_Modules/TriplaneGaussian/models/renderer.py b/Gen_3D_Modules/TriplaneGaussian/models/renderer.py index beec388..c176a95 100644 --- a/Gen_3D_Modules/TriplaneGaussian/models/renderer.py +++ b/Gen_3D_Modules/TriplaneGaussian/models/renderer.py @@ -13,7 +13,7 @@ from ..utils.ops import scale_tensor from .networks import MLP from einops import rearrange, reduce -from mesh_processer.mesh_utils import construct_list_of_gs_attributes, write_gs_ply +from mesh_processor.mesh_utils import construct_list_of_gs_attributes, write_gs_ply inverse_sigmoid = lambda x: np.log(x / (1 - x)) diff --git a/Gen_3D_Modules/TripoSR/models/isosurface.py b/Gen_3D_Modules/TripoSR/models/isosurface.py index c4a9e31..e8acb3d 100644 --- a/Gen_3D_Modules/TripoSR/models/isosurface.py +++ b/Gen_3D_Modules/TripoSR/models/isosurface.py @@ -4,7 +4,7 @@ import torch import torch.nn as nn from mcubes import marching_cubes -from mesh_processer.mesh_utils import switch_vector_axis +from mesh_processor.mesh_utils import switch_vector_axis class IsosurfaceHelper(nn.Module): diff --git a/Gen_3D_Modules/Unique3D/scripts/project_mesh.py b/Gen_3D_Modules/Unique3D/scripts/project_mesh.py index 1782e07..a86ac0f 100644 --- a/Gen_3D_Modules/Unique3D/scripts/project_mesh.py +++ b/Gen_3D_Modules/Unique3D/scripts/project_mesh.py @@ -12,7 +12,7 @@ from pytorch3d.renderer import ( FoVOrthographicCameras, ) from pytorch3d.renderer import MeshRasterizer -from mesh_processer.mesh import Mesh +from mesh_processor.mesh import Mesh def get_camera(world_to_cam, fov_in_degrees=60, focal_length=1 / (2**0.5), cam_type='fov'): # pytorch3d expects transforms as row-vectors, so flip rotation: https://github.com/facebookresearch/pytorch3d/issues/1183 diff --git a/MVs_Algorithms/DiffRastMesh/diff_mesh_renderer.py b/MVs_Algorithms/DiffRastMesh/diff_mesh_renderer.py index 98a07da..d559a14 100644 --- a/MVs_Algorithms/DiffRastMesh/diff_mesh_renderer.py +++ b/MVs_Algorithms/DiffRastMesh/diff_mesh_renderer.py @@ -9,7 +9,7 @@ import nvdiffrast.torch as dr from kiui.op import inverse_sigmoid -from mesh_processer.mesh import safe_normalize +from mesh_processor.mesh import safe_normalize def scale_img_nhwc(x, size, mag='bilinear', min='bilinear'): assert (x.shape[1] >= size[0] and x.shape[2] >= size[1]) or (x.shape[1] < size[0] and x.shape[2] < size[1]), "Trying to magnify image in one dimension and minify in the other" diff --git a/MVs_Algorithms/FlexiCubes/flexicubes_trainer.py b/MVs_Algorithms/FlexiCubes/flexicubes_trainer.py index 41ba267..cccce84 100644 --- a/MVs_Algorithms/FlexiCubes/flexicubes_trainer.py +++ b/MVs_Algorithms/FlexiCubes/flexicubes_trainer.py @@ -13,7 +13,7 @@ from .util import SimpleMesh from .loss import sdf_reg_loss from shared_utils.camera_utils import OrbitCamera -from mesh_processer.mesh import Mesh +from mesh_processor.mesh import Mesh def lr_schedule(iter): diff --git a/MVs_Algorithms/GaussianSplatting/main_3DGS_renderer.py b/MVs_Algorithms/GaussianSplatting/main_3DGS_renderer.py index d05394e..0715575 100644 --- a/MVs_Algorithms/GaussianSplatting/main_3DGS_renderer.py +++ b/MVs_Algorithms/GaussianSplatting/main_3DGS_renderer.py @@ -15,8 +15,8 @@ from kornia.geometry.conversions import ( from kiui.op import inverse_sigmoid from shared_utils.sh_utils import eval_sh, SH2RGB, RGB2SH -from mesh_processer.mesh import Mesh, PointCloud -from mesh_processer.mesh_utils import construct_list_of_gs_attributes, write_gs_ply, read_gs_ply, K_nearest_neighbors_func +from mesh_processor.mesh import Mesh, PointCloud +from mesh_processor.mesh_utils import construct_list_of_gs_attributes, write_gs_ply, read_gs_ply, K_nearest_neighbors_func def get_expon_lr_func( lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000 diff --git a/__init__.py b/__init__.py index cc8c101..bf931fc 100644 --- a/__init__.py +++ b/__init__.py @@ -8,10 +8,12 @@ import logging ROOT_PATH = os.path.dirname(os.path.realpath(__file__)) MODULE_PATH = os.path.join(ROOT_PATH, "Gen_3D_Modules") MV_ALGO_PATH = os.path.join(ROOT_PATH, "MVs_Algorithms") +MESH_PROCESSOR_PATH = os.path.join(ROOT_PATH, "mesh_processor") -sys.path.append(ROOT_PATH) +sys.path.insert(0, ROOT_PATH) sys.path.append(MODULE_PATH) sys.path.append(MV_ALGO_PATH) +sys.path.append(MESH_PROCESSOR_PATH) import shutil import __main__ diff --git a/mesh_processer/__init__.py b/mesh_processer/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/mesh_processor/__init__.py b/mesh_processor/__init__.py new file mode 100644 index 0000000..2ba5108 --- /dev/null +++ b/mesh_processor/__init__.py @@ -0,0 +1,85 @@ +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, + 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", + "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/export_utils.py b/mesh_processor/export_utils.py new file mode 100644 index 0000000..b8612ed --- /dev/null +++ b/mesh_processor/export_utils.py @@ -0,0 +1,192 @@ +""" +Utilities for exporting from pipelines to our mesh objects +Direct functions WITHOUT trimesh dependency +""" + +import torch +import numpy as np +from .mesh import Mesh, FastMesh + + +def export_to_fastmesh(mesh_output, device=None): + """ + Export from pipeline to FastMesh + + Args: + mesh_output: object with mesh_v and mesh_f from pipeline + device: torch device + + Returns: + FastMesh object or list of FastMesh objects + """ + + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + if isinstance(mesh_output, list): + outputs = [] + for mesh in mesh_output: + if mesh is None: + outputs.append(None) + else: + fastmesh = _convert_single_mesh_to_fastmesh(mesh, device) + outputs.append(fastmesh) + return outputs + else: + return _convert_single_mesh_to_fastmesh(mesh_output, device) + + +def export_to_mesh(mesh_output, device=None): + """ + Export from pipeline to standard Mesh + + Args: + mesh_output: object with mesh_v and mesh_f from pipeline + device: torch device + + Returns: + Mesh object or list of Mesh objects + """ + + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + if isinstance(mesh_output, list): + outputs = [] + for mesh in mesh_output: + if mesh is None: + outputs.append(None) + else: + standard_mesh = _convert_single_mesh_to_mesh(mesh, device) + outputs.append(standard_mesh) + return outputs + else: + return _convert_single_mesh_to_mesh(mesh_output, device) + + +def _convert_single_mesh_to_fastmesh(mesh_obj, device): + """Convert one mesh object to FastMesh""" + + try: + # Extract vertices and faces + vertices = mesh_obj.mesh_v + faces = mesh_obj.mesh_f + + # Flip faces (as in original function) + faces = faces[:, ::-1] + + # Convert to numpy if needed + if hasattr(vertices, 'cpu'): + vertices = vertices.cpu().numpy() + if hasattr(faces, 'cpu'): + faces = faces.cpu().numpy() + + vertices = np.array(vertices, dtype=np.float32) + faces = np.array(faces, dtype=np.int32) + + print(f"[export_to_fastmesh] Convert: {vertices.shape[0]} vertices, {faces.shape[0]} faces") + + fastmesh = FastMesh(device=device) + fastmesh.v = torch.tensor(vertices, dtype=torch.float32, device=device) + fastmesh.f = torch.tensor(faces, dtype=torch.int32, device=device) + + fastmesh.auto_normal() + + fastmesh._create_empty_albedo() + + print(f"[export_to_fastmesh] FastMesh created successfully") + return fastmesh + + except Exception as e: + print(f"[export_to_fastmesh] Error: {e}") + return None + + +def _convert_single_mesh_to_mesh(mesh_obj, device): + """Convert one mesh object to standard Mesh""" + + try: + vertices = mesh_obj.mesh_v + faces = mesh_obj.mesh_f + + faces = faces[:, ::-1] + + if hasattr(vertices, 'cpu'): + vertices = vertices.cpu().numpy() + if hasattr(faces, 'cpu'): + faces = faces.cpu().numpy() + + vertices = np.array(vertices, dtype=np.float32) + faces = np.array(faces, dtype=np.int32) + + print(f"[export_to_mesh] Convert: {vertices.shape[0]} vertices, {faces.shape[0]} faces") + + mesh = Mesh(device=device) + mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device) + mesh.f = torch.tensor(faces, dtype=torch.int32, device=device) + + mesh.auto_normal() + + mesh._create_empty_albedo_fast() + + print(f"[export_to_mesh] Mesh created successfully") + return mesh + + except Exception as e: + print(f"[export_to_mesh] Error: {e}") + return None + + +def export_to_mesh_ultra_fast(mesh_output, device=None): + """ + Ultra-fast export to standard Mesh WITHOUT any slow operations + """ + + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + try: + # Process list or single object + if isinstance(mesh_output, list): + return [_convert_single_mesh_to_mesh_ultra_fast(mesh, device) for mesh in mesh_output] + else: + return _convert_single_mesh_to_mesh_ultra_fast(mesh_output, device) + + except Exception as e: + print(f"[export_to_mesh_ultra_fast] Error: {e}") + return None + + +def _convert_single_mesh_to_mesh_ultra_fast(mesh_obj, device): + """Convert one mesh object to standard Mesh Ultra-fast""" + + try: + # Extract vertices and faces + vertices = mesh_obj.mesh_v + faces = mesh_obj.mesh_f + + # Flip faces (as in original function) + faces = faces[:, ::-1] + + # Convert to numpy if needed + if hasattr(vertices, 'cpu'): + vertices = vertices.cpu().numpy() + if hasattr(faces, 'cpu'): + faces = faces.cpu().numpy() + + vertices = np.array(vertices, dtype=np.float32) + faces = np.array(faces, dtype=np.int32) + + print(f"[export_to_mesh_ultra_fast] Convert: {vertices.shape[0]} vertices, {faces.shape[0]} faces") + + # Create standard Mesh Minimal + mesh = Mesh(device=device) + mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device) + mesh.f = torch.tensor(faces, dtype=torch.int32, device=device) + + print(f"[export_to_mesh_ultra_fast] Mesh created instantly") + return mesh + + except Exception as e: + print(f"[export_to_mesh_ultra_fast] Error: {e}") + return None 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(" 3: + vertex_colors = vertex_colors[:, :3] # Remove alpha channel + print(f"[load_gltf] extracted vertex colors: {vertex_colors.shape}") + except Exception as e: + print(f"[load_gltf] failed to extract vertex colors: {e}") + + # Set UV coordinates + mesh.vt = ( + torch.tensor(texcoords, dtype=torch.float32, device=device) + if texcoords is not None + else None + ) + mesh.ft = mesh.f if texcoords is not None else None + + # Set normals + mesh.vn = ( + torch.tensor(normals, dtype=torch.float32, device=device) + if normals is not None + else None + ) + mesh.fn = mesh.f if normals is not None else None + + # Set vertex colors + mesh.vc = ( + torch.tensor(vertex_colors, dtype=torch.float32, device=device) + if vertex_colors is not None + else None + ) + + # Try to extract textures and materials + should_create_empty_albedo = True + + if doc.materials() and doc.textures() and doc.images(): + try: + # Take the first material + material = doc.materials()[0] + pbr = material.get("pbrMetallicRoughness", {}) + + # Base texture (albedo) + base_color_texture = pbr.get("baseColorTexture") + if base_color_texture: + texture_idx = base_color_texture.get("index", 0) + if texture_idx < len(doc.textures()): + texture = doc.textures()[texture_idx] + image_idx = texture.get("source", 0) + if image_idx < len(doc.images()): + image_info = doc.images()[image_idx] + + # Extract image from buffer + buffer_view_idx = image_info.get("bufferView") + if buffer_view_idx is not None: + buffer_view = doc.bufferViews()[buffer_view_idx] + buffer_idx = buffer_view.get("buffer", 0) + byte_offset = buffer_view.get("byteOffset", 0) + byte_length = buffer_view.get("byteLength", 0) + + binary_data = get_binary_data(doc, buffer_idx) + image_bytes = binary_data[byte_offset:byte_offset + byte_length] + + # Decode image + image_array = np.frombuffer(image_bytes, dtype=np.uint8) + texture_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR) + + if texture_image is not None: + texture_image = cv2.cvtColor(texture_image, cv2.COLOR_BGR2RGB) + texture_image = texture_image.astype(np.float32) / 255.0 + mesh.albedo = torch.tensor(texture_image, dtype=torch.float32, device=device).contiguous() + should_create_empty_albedo = False + print(f"[load_gltf] loaded albedo texture: {texture_image.shape}") + + # Metallic-Roughness texture + metallic_roughness_texture = pbr.get("metallicRoughnessTexture") + if metallic_roughness_texture and not should_create_empty_albedo: + texture_idx = metallic_roughness_texture.get("index", 0) + if texture_idx < len(doc.textures()): + # Similarly for metallic-roughness + # (simplified version for now) + pass + + except Exception as e: + print(f"[load_gltf] failed to extract textures: {e}") + + # Create empty texture if not able to load + if should_create_empty_albedo: + mesh.set_new_albedo(1024, 1024) + print(f"[load_gltf] created empty albedo texture with shape: {mesh.albedo.shape}") + + return mesh + + except Exception as e: + print(f"[load_gltf] failed to load {path}: {e}") + return None + def set_new_albedo(self, res_H, res_W): if self.albedo is None: texture = np.ones((res_H, res_W, 3), dtype=np.float32) * np.array([0.5, 0.5, 0.5]) self.albedo = torch.tensor(texture, dtype=torch.float32, device=self.device) else: self.albedo = prepare_torch_img(self.albedo.unsqueeze(0), res_H, res_W, self.device).squeeze(0).permute(1, 2, 0).contiguous() # (1, 3, H, W) -> (H, W, 3) + + def _create_empty_albedo_fast(self): + """Create empty albedo texture quickly (without prepare_torch_img)""" + if self.albedo is None: + texture = np.ones((1024, 1024, 3), dtype=np.float32) * np.array([0.5, 0.5, 0.5]) + self.albedo = torch.tensor(texture, dtype=torch.float32, device=self.device) + print(f"[Mesh] Fast empty texture: {self.albedo.shape}") # aabb def aabb(self): @@ -625,8 +833,30 @@ class Mesh: v_np = self.v.detach().cpu().numpy() f_np = self.f.detach().cpu().numpy() - _mesh = trimesh.Trimesh(vertices=v_np, faces=f_np) - _mesh.export(path) + # Simple PLY writing without trimesh + self._write_ply_simple(path, v_np, f_np) + + def _write_ply_simple(self, path, vertices, faces): + """Simple PLY writing without trimesh""" + with open(path, 'w') as f: + # PLY header + f.write("ply\n") + f.write("format ascii 1.0\n") + f.write(f"element vertex {len(vertices)}\n") + f.write("property float x\n") + f.write("property float y\n") + f.write("property float z\n") + f.write(f"element face {len(faces)}\n") + f.write("property list uchar int vertex_indices\n") + f.write("end_header\n") + + # Vertices + for v in vertices: + f.write(f"{v[0]} {v[1]} {v[2]}\n") + + # Faces + for face in faces: + f.write(f"3 {face[0]} {face[1]} {face[2]}\n") def write_glb(self, path): @@ -897,9 +1127,407 @@ class Mesh: shs = np.random.random((num_pts, 3)) / 255.0 normals = np.zeros((num_pts, 3)) pcd = PointCloud(points=xyz, colors=SH2RGB(shs), normals=normals) - return pcd + return pcd + + +class FastMesh: + """ + Fast implementation of mesh based on FastGLB without dependency on trimesh. + Uses direct loading of GLB/GLTF through mesh_processor. + """ + + def __init__( + self, + v: Optional[Tensor] = None, + f: Optional[Tensor] = None, + vn: Optional[Tensor] = None, + fn: Optional[Tensor] = None, + vt: Optional[Tensor] = None, + ft: Optional[Tensor] = None, + vc: Optional[Tensor] = None, + albedo: Optional[Tensor] = None, + metallicRoughness: Optional[Tensor] = None, + device: Optional[torch.device] = None, + ): + """Initialize FastMesh with the same parameters as Mesh""" + self.device = device if device is not None else torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.v = v + self.vn = vn + self.vt = vt + self.f = f + self.fn = fn + self.ft = ft + self.vc = vc + self.albedo = albedo + self.metallicRoughness = metallicRoughness + self.ori_center = 0 + self.ori_scale = 1 + + @classmethod + def load(cls, path, resize=True, renormal=True, retex=False, clean=False, bound=0.5, front_dir='+z', **kwargs): + """Loading mesh through FastGLB approach""" + device = kwargs.get('device') + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + if path.endswith(".obj"): + mesh = cls._load_obj_fast(path, device=device, **kwargs) + elif path.endswith((".glb", ".gltf")): + mesh = cls._load_gltf_fast(path, device=device, **kwargs) + else: + standard_mesh = Mesh.load(path, resize=False, renormal=False, retex=False, clean=False, **kwargs) + mesh = cls._from_standard_mesh(standard_mesh) + + if mesh is None: + return None + + if clean: + mesh._clean_mesh() + + print(f"[FastMesh loading] v: {mesh.v.shape}, f: {mesh.f.shape}") + + if resize: + mesh.auto_size(bound=bound) + + if renormal or mesh.vn is None: + mesh.auto_normal() + print(f"[FastMesh loading] vn: {mesh.vn.shape}, fn: {mesh.fn.shape}") + + if retex or (mesh.albedo is not None and mesh.vt is None): + mesh.auto_uv(cache_path=path) + print(f"[FastMesh loading] vt: {mesh.vt.shape}, ft: {mesh.ft.shape}") + + if front_dir != "+z": + mesh._apply_front_dir_rotation(front_dir) + + return mesh + + @classmethod + def _load_gltf_fast(cls, path, device=None, **kwargs): + """Fast loading of GLB/GLTF through mesh_processor""" + + try: + print(f"[FastMesh] Loading GLB/GLTF: {path}") + + # Load GLTF document + from .io_gltf import load_gltf_or_glb, get_binary_data + from .mesh_ops import get_all_meshes_triangles + from .accessors import access_data + + doc = load_gltf_or_glb(path) + print(f"[FastMesh] GLTF document: {len(doc.meshes())} meshes") + + # Extract geometry + vertices, faces, mesh_groups = get_all_meshes_triangles(doc, transform_to_global=True) + + if len(vertices) == 0 or len(faces) == 0: + print(f"[FastMesh] No geometry in {path}") + return None + + mesh = cls(device=device) + mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device) + mesh.f = torch.tensor(faces, dtype=torch.int32, device=device) + + # Extract additional data + cls._extract_gltf_attributes(mesh, doc, device) + cls._extract_gltf_textures(mesh, doc, device) + + print(f"[FastMesh] Loaded: {len(vertices)} vertices, {len(faces)} faces") + return mesh + + except Exception as e: + print(f"[FastMesh] Error loading GLB/GLTF {path}: {e}") + return None + + @classmethod + def _extract_gltf_attributes(cls, mesh, doc, device): + """Extract attributes from GLTF (UV, normals, colors)""" + + if not doc.meshes(): + return + + first_mesh = doc.meshes()[0] + for primitive in first_mesh.get("primitives", []): + attributes = primitive.get("attributes", {}) + + # UV coordinates + if "TEXCOORD_0" in attributes and mesh.vt is None: + try: + from .accessors import access_data + uv_accessor_idx = int(attributes["TEXCOORD_0"]) + uv_data = access_data(doc, uv_accessor_idx).astype(np.float32) + uv_data[:, 1] = 1.0 - uv_data[:, 1] # Flip V + mesh.vt = torch.tensor(uv_data, dtype=torch.float32, device=device) + mesh.ft = mesh.f + print(f"[FastMesh] UV: {mesh.vt.shape}") + except Exception as e: + print(f"[FastMesh] Error UV: {e}") + + # Normals + if "NORMAL" in attributes and mesh.vn is None: + try: + from .accessors import access_data + normal_accessor_idx = int(attributes["NORMAL"]) + normal_data = access_data(doc, normal_accessor_idx).astype(np.float32) + mesh.vn = torch.tensor(normal_data, dtype=torch.float32, device=device) + mesh.fn = mesh.f + print(f"[FastMesh] Normals: {mesh.vn.shape}") + except Exception as e: + print(f"[FastMesh] Error normals: {e}") + + # Vertex colors + if "COLOR_0" in attributes and mesh.vc is None: + try: + from .accessors import access_data + color_accessor_idx = int(attributes["COLOR_0"]) + color_data = access_data(doc, color_accessor_idx).astype(np.float32) + if color_data.shape[1] > 3: + color_data = color_data[:, :3] + mesh.vc = torch.tensor(color_data, dtype=torch.float32, device=device) + print(f"[FastMesh] Colors: {mesh.vc.shape}") + except Exception as e: + print(f"[FastMesh] Error colors: {e}") + + @classmethod + def _extract_gltf_textures(cls, mesh, doc, device): + """Extract textures from GLTF""" + + should_create_empty = True + + try: + if not (doc.materials() and doc.textures() and doc.images()): + mesh._create_empty_albedo() + return + + material = doc.materials()[0] + pbr = material.get("pbrMetallicRoughness", {}) + + # Albedo texture + base_color_texture = pbr.get("baseColorTexture") + if base_color_texture: + albedo_texture = cls._extract_texture_by_index(doc, base_color_texture.get("index", 0)) + if albedo_texture is not None: + albedo_float = albedo_texture.astype(np.float32) / 255.0 + mesh.albedo = torch.tensor(albedo_float, dtype=torch.float32, device=device).contiguous() + should_create_empty = False + print(f"[FastMesh] Albedo: {mesh.albedo.shape}") + + # Metallic-Roughness texture + mr_texture_info = pbr.get("metallicRoughnessTexture") + if mr_texture_info and not should_create_empty: + mr_texture = cls._extract_texture_by_index(doc, mr_texture_info.get("index", 0)) + if mr_texture is not None: + mr_float = mr_texture.astype(np.float32) / 255.0 + mesh.metallicRoughness = torch.tensor(mr_float, dtype=torch.float32, device=device).contiguous() + print(f"[FastMesh] MetallicRoughness: {mesh.metallicRoughness.shape}") + + except Exception as e: + print(f"[FastMesh] Error extracting textures: {e}") + + if should_create_empty: + mesh._create_empty_albedo() + + @classmethod + def _extract_texture_by_index(cls, doc, texture_idx): + """Extract texture by index from GLTF""" + + try: + from .io_gltf import get_binary_data + + if texture_idx >= len(doc.textures()): + return None + + texture = doc.textures()[texture_idx] + image_idx = texture.get("source", 0) + + if image_idx >= len(doc.images()): + return None + + image_info = doc.images()[image_idx] + buffer_view_idx = image_info.get("bufferView") + + if buffer_view_idx is None: + return None + + buffer_view = doc.bufferViews()[buffer_view_idx] + buffer_idx = buffer_view.get("buffer", 0) + byte_offset = buffer_view.get("byteOffset", 0) + byte_length = buffer_view.get("byteLength", 0) + + binary_data = get_binary_data(doc, buffer_idx) + image_bytes = binary_data[byte_offset:byte_offset + byte_length] + + # Decode image + image_array = np.frombuffer(image_bytes, dtype=np.uint8) + texture_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR) + + if texture_image is not None: + texture_image = cv2.cvtColor(texture_image, cv2.COLOR_BGR2RGB) + return texture_image + + except Exception as e: + print(f"[FastMesh] Error extracting texture {texture_idx}: {e}") + + return None + + @classmethod + def _load_obj_fast(cls, path, device=None, **kwargs): + """Fast loading of OBJ (fallback to standard method)""" + standard_mesh = Mesh.load_obj(path, device=device, **kwargs) + return cls._from_standard_mesh(standard_mesh) + + @classmethod + def _from_standard_mesh(cls, standard_mesh): + """Convert standard Mesh to FastMesh""" + if standard_mesh is None: + return None + + return cls( + v=standard_mesh.v, + f=standard_mesh.f, + vn=standard_mesh.vn, + fn=standard_mesh.fn, + vt=standard_mesh.vt, + ft=standard_mesh.ft, + vc=standard_mesh.vc, + albedo=standard_mesh.albedo, + metallicRoughness=standard_mesh.metallicRoughness, + device=standard_mesh.device + ) + + def _create_empty_albedo(self): + """Create empty albedo texture""" + texture = np.ones((1024, 1024, 3), dtype=np.float32) * np.array([0.5, 0.5, 0.5]) + self.albedo = torch.tensor(texture, dtype=torch.float32, device=self.device) + print(f"[FastMesh] Empty texture: {self.albedo.shape}") + + def _clean_mesh(self): + """Clean mesh (simplified version)""" + try: + from kiui.mesh_utils import clean_mesh + vertices = self.v.detach().cpu().numpy() + triangles = self.f.detach().cpu().numpy() + vertices, triangles = clean_mesh(vertices, triangles, remesh=False) + self.v = torch.from_numpy(vertices).contiguous().float().to(self.device) + self.f = torch.from_numpy(triangles).contiguous().int().to(self.device) + except Exception as e: + print(f"[FastMesh] Error cleaning: {e}") + + def _apply_front_dir_rotation(self, front_dir): + """Apply rotation for aligning front_dir with +z""" + # Use the same logic as in standard Mesh + if front_dir == "+z": + return + + # Define transformation matrix + if "-z" in front_dir: + T = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, -1]], device=self.device, dtype=torch.float32) + elif "+x" in front_dir: + T = torch.tensor([[0, 0, 1], [0, 1, 0], [1, 0, 0]], device=self.device, dtype=torch.float32) + elif "-x" in front_dir: + T = torch.tensor([[0, 0, -1], [0, 1, 0], [1, 0, 0]], device=self.device, dtype=torch.float32) + elif "+y" in front_dir: + T = torch.tensor([[1, 0, 0], [0, 0, 1], [0, 1, 0]], device=self.device, dtype=torch.float32) + elif "-y" in front_dir: + T = torch.tensor([[1, 0, 0], [0, 0, -1], [0, 1, 0]], device=self.device, dtype=torch.float32) + else: + T = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], device=self.device, dtype=torch.float32) + + # Additional rotations + if '1' in front_dir: + T @= torch.tensor([[0, -1, 0], [1, 0, 0], [0, 0, 1]], device=self.device, dtype=torch.float32) + elif '2' in front_dir: + T @= torch.tensor([[1, 0, 0], [0, -1, 0], [0, 0, 1]], device=self.device, dtype=torch.float32) + elif '3' in front_dir: + T @= torch.tensor([[0, 1, 0], [-1, 0, 0], [0, 0, 1]], device=self.device, dtype=torch.float32) + + # Apply transformation + if self.v is not None: + self.v @= T + if self.vn is not None: + self.vn @= T + + # Use methods from standard Mesh + def aabb(self): + """AABB mesh""" + return torch.min(self.v, dim=0).values, torch.max(self.v, dim=0).values + + @torch.no_grad() + def auto_size(self, bound=0.9): + """Automatic size change""" + vmin, vmax = self.aabb() + self.ori_center = (vmax + vmin) / 2 + self.ori_scale = 2 * bound / torch.max(vmax - vmin).item() + self.v = (self.v - self.ori_center) * self.ori_scale + + def auto_normal(self): + """Automatic calculation of normals""" + i0, i1, i2 = self.f[:, 0].long(), self.f[:, 1].long(), self.f[:, 2].long() + v0, v1, v2 = self.v[i0, :], self.v[i1, :], self.v[i2, :] + + face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1) + + vn = torch.zeros_like(self.v) + vn.scatter_add_(0, i0[:, None].repeat(1, 3), face_normals) + vn.scatter_add_(0, i1[:, None].repeat(1, 3), face_normals) + vn.scatter_add_(0, i2[:, None].repeat(1, 3), face_normals) + + vn = torch.where( + dot(vn, vn) > 1e-20, + vn, + torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device=vn.device), + ) + vn = safe_normalize(vn) + + self.vn = vn + self.fn = self.f + + def auto_uv(self, cache_path=None, vmap=True): + """Automatic calculation of UV (use method from standard Mesh)""" + # Temporarily delegate to standard Mesh + temp_mesh = Mesh( + v=self.v, f=self.f, vn=self.vn, fn=self.fn, + vt=self.vt, ft=self.ft, vc=self.vc, + albedo=self.albedo, metallicRoughness=self.metallicRoughness, + device=self.device + ) + temp_mesh.auto_uv(cache_path=cache_path, vmap=vmap) + + # Copy result back + self.vt = temp_mesh.vt + self.ft = temp_mesh.ft + if vmap: + self.v = temp_mesh.v + self.f = temp_mesh.f + self.vn = temp_mesh.vn + self.fn = temp_mesh.fn + if temp_mesh.vc is not None: + self.vc = temp_mesh.vc + + def to(self, device): + """Move to another device""" + self.device = device + for name in ["v", "f", "vn", "fn", "vt", "ft", "albedo", "vc", "metallicRoughness"]: + tensor = getattr(self, name) + if tensor is not None: + setattr(self, name, tensor.to(device)) + return self + + def write(self, path): + """Saving to file (use methods from standard Mesh)""" + # Create temporary standard Mesh for saving + temp_mesh = Mesh( + v=self.v, f=self.f, vn=self.vn, fn=self.fn, + vt=self.vt, ft=self.ft, vc=self.vc, + albedo=self.albedo, metallicRoughness=self.metallicRoughness, + device=self.device + ) + temp_mesh.ori_center = self.ori_center + temp_mesh.ori_scale = self.ori_scale + return temp_mesh.write(path) + + class PointCloud(NamedTuple): points: np.array colors: np.array diff --git a/mesh_processor/mesh_ops.py b/mesh_processor/mesh_ops.py new file mode 100644 index 0000000..63da0c3 --- /dev/null +++ b/mesh_processor/mesh_ops.py @@ -0,0 +1,161 @@ +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) + diff --git a/mesh_processer/mesh_utils.py b/mesh_processor/mesh_utils.py similarity index 93% rename from mesh_processer/mesh_utils.py rename to mesh_processor/mesh_utils.py index c3799df..24857c1 100644 --- a/mesh_processer/mesh_utils.py +++ b/mesh_processor/mesh_utils.py @@ -598,6 +598,66 @@ def K_nearest_neighbors_func( else: return nn[0, :, :, :], idx[0, :, :], dist[0, :, :] + +def reduce_mesh_faces_inplace(mesh_obj, max_faces: int = 40000, backend="pymeshlab", optimalplacement=True, verbose=True): + """ + Reduce faces for Mesh/FastMesh objects using decimate_mesh (in-place modification) + + Args: + mesh_obj: Mesh or FastMesh object (modified in-place) + max_faces: maximum number of faces + backend: algorithm backend ("pymeshlab" or "pyfqmr") + optimalplacement: optimal placement for decimation + verbose: print progress + + Returns: + Same mesh object (modified in-place) + """ + + faces_before = mesh_obj.f.shape[0] + + # Skip if already under limit + if faces_before <= max_faces: + if verbose: + print(f"Mesh already has {faces_before} faces (≤ {max_faces}) - skipping") + return mesh_obj + + try: + + # Convert to numpy + vertices_np = mesh_obj.v.detach().cpu().numpy().astype(np.float32) + faces_np = mesh_obj.f.detach().cpu().numpy().astype(np.int32) + + # Apply decimation + reduced_vertices, reduced_faces = decimate_mesh( + vertices_np, faces_np, + target=max_faces, + backend=backend, + optimalplacement=optimalplacement, + verbose=verbose + ) + + # Update existing mesh object (in-place) + device = mesh_obj.v.device + mesh_obj.v = torch.from_numpy(reduced_vertices).contiguous().float().to(device) + mesh_obj.f = torch.from_numpy(reduced_faces).contiguous().int().to(device) + + # Recalculate normals + mesh_obj.auto_normal() + + # Stats + if verbose: + faces_after = mesh_obj.f.shape[0] + reduction_ratio = ((faces_before - faces_after) / faces_before) * 100 + + return mesh_obj + + except Exception as e: + if verbose: + print(f"Face reduction error: {e}, keep original mesh") + return mesh_obj + + def interpolate_texture_map_attr(mesh, texture_size: int = 256, batch_size: int = 64, interpolate_color=True, interpolate_position=False): # Get UV coordinates and faces if mesh.vt is None: 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 + diff --git a/nodes.py b/nodes.py index e4b01fc..5f3f43b 100644 --- a/nodes.py +++ b/nodes.py @@ -37,8 +37,9 @@ from plyfile import PlyData import trimesh from PIL import Image -from .mesh_processer.mesh import Mesh -from .mesh_processer.mesh_utils import ( +from mesh_processor.export_utils import export_to_fastmesh, export_to_mesh +from mesh_processor.mesh import FastMesh, Mesh +from mesh_processor.mesh_utils import ( ply_to_points_cloud, get_target_axis_and_scale, switch_ply_axis_and_scale, @@ -48,6 +49,7 @@ from .mesh_processer.mesh_utils import ( color_func_to_albedo, interpolate_texture_map_attr, decimate_mesh, + reduce_mesh_faces_inplace ) from FlexiCubes.flexicubes_trainer import FlexiCubesTrainer @@ -128,6 +130,7 @@ from Gen_3D_Modules.PartCrafter.partcrafter_src.pipelines.pipeline_partcrafter i from Gen_3D_Modules.PartCrafter.partcrafter_src.utils.data_utils import get_colored_mesh_composition from Gen_3D_Modules.PartCrafter.partcrafter_src.utils.render_utils import explode_mesh import zipfile +from typing import Union os.environ['SPCONV_ALGO'] = 'native' @@ -4569,7 +4572,6 @@ class Hunyuan3D_V2_ShapeGen_MV: return (Mesh.load_trimesh(given_mesh=mesh),) -#-------------------------------- class Load_StableGen_Trellis_Pipeline: CATEGORY = "Comfy3D/Algorithm" RETURN_TYPES = ("DIFFUSERS_PIPE",) @@ -4660,7 +4662,6 @@ class Load_StableGen_StableX_Pipeline: pipe = self.__class__._build_pipe(repo, use_fp16) return (pipe,) - class StableGen_Trellis_Image_To_3D: """ 3D generation pipeline using Trellis model. @@ -4766,7 +4767,6 @@ class StableGen_Trellis_Image_To_3D: except Exception as e: raise Exception(f"[StableGen_Trellis_Image_To_3D] 3D generation failed: {str(e)}") - class StableGen_StableX_Process_Image: """ Image processing pipeline using StableX model. @@ -4814,7 +4814,7 @@ class StableGen_StableX_Process_Image: except Exception as e: raise Exception(f"[StableGen_StableX_Process_Image] Image processing failed: {str(e)}") -# --- MV START ---- + class Load_MVAdapter_IG2MV_Pipeline: """Loader pipeline for MV-Adapter (Image to Multi-View)""" CATEGORY = "Comfy3D/Algorithm" @@ -4851,7 +4851,6 @@ class Load_MVAdapter_IG2MV_Pipeline: vae_model = None if vae_model == "None" else vae_model lora_model = None if not lora_model else lora_model - # Подготавливаем параметры для пайплайна pipeline_kwargs = { "base_model": base_model, "vae_model": vae_model, @@ -4870,6 +4869,7 @@ class Load_MVAdapter_IG2MV_Pipeline: print("MV-Adapter IG2MV pipeline loaded successfully") return (pipe,) + class MVAdapter_IG2MV: """Generate multi-view images from single image and 3D mesh""" CATEGORY = "Comfy3D/Algorithm" @@ -4931,6 +4931,7 @@ class MVAdapter_IG2MV: return_images = pils_to_torch_imgs(images, device=DEVICE_STR) return (return_images,) + class Load_MVAdapter_TG2MV_Pipeline: """Loader pipeline for MV-Adapter Text-Guided to Multi-View (TG2MV)""" CATEGORY = "Comfy3D/Algorithm" @@ -5044,6 +5045,7 @@ class MVAdapter_TG2MV: return_images = pils_to_torch_imgs(images, device=DEVICE_STR) return (return_images,) + class Load_MVAdapter_Texture_Pipeline: """Load texture projection pipeline for MVAdapter""" CATEGORY = "Comfy3D/Algorithm" @@ -5206,6 +5208,7 @@ class MVAdapter_Texture_Projection: os.remove(temp_grid_path) raise e + class Load_Hunyuan3D_21_ShapeGen_Pipeline: """Load Hunyuan3D-2.1 Shape Generation Pipeline""" @@ -5274,6 +5277,7 @@ class Load_Hunyuan3D_21_ShapeGen_Pipeline: return (pipeline,) + class Load_Hunyuan3D_21_TexGen_Pipeline: """Load Hunyuan3D-2.1 Texture Generation Pipeline""" @@ -5394,8 +5398,9 @@ class Load_Hunyuan3D_21_TexGen_Pipeline: print(f"[TexGen-Loader] Cached new pipeline {cache_key}") return (pipeline,) + class Hunyuan3D_21_ShapeGen: - """Hunyuan3D-2.1 Shape Generation with automatic pipeline cleanup""" + """Hunyuan3D-2.1 Shape Generation """ CATEGORY = "Comfy3D/Algorithm/Hunyuan3D-2.1" RETURN_TYPES = ("MESH", "IMAGE") @@ -5414,20 +5419,20 @@ class Hunyuan3D_21_ShapeGen: "octree_resolution": ("INT", {"default": 256, "min": 64, "max": 512}), "remove_background": ("BOOLEAN", {"default": True}), "auto_cleanup": ("BOOLEAN", {"default": True}), + "max_faces": ("INT", {"default": 40000, "min": 1000, "max": 1000000}), } } @torch.no_grad() def generate(self, shapegen_pipe, image, seed, steps, guidance_scale, octree_resolution, remove_background, auto_cleanup): - pil_image = torch_imgs_to_pils(image)[0].convert("RGBA") - - if remove_background or pil_image.mode == "RGB": - rmbg_worker = BackgroundRemover_2_1() - pil_image = rmbg_worker(pil_image.convert('RGB')) - del rmbg_worker - - generator = torch.Generator(device=shapegen_pipe.device) - generator = generator.manual_seed(int(seed)) + pil_image = torch_imgs_to_pils(image)[0] + if remove_background: + bg_remover = BackgroundRemover_2_1() + if bg_remover is not None: + pil_image = bg_remover(pil_image) + del bg_remover + + generator = torch.Generator(device=shapegen_pipe.device).manual_seed(seed) outputs = shapegen_pipe( image=pil_image, @@ -5439,36 +5444,23 @@ class Hunyuan3D_21_ShapeGen: output_type='mesh' ) - mesh = export_to_trimesh_2_1(outputs)[0] + mesh_objects = export_to_fastmesh(outputs) + mesh_out = mesh_objects[0] if isinstance(mesh_objects, list) else mesh_objects - face_reduce_worker = FaceReducer_2_1() - mesh = face_reduce_worker(mesh) - del face_reduce_worker + if mesh_out is None: + raise Exception("Failed to create mesh from pipeline") + + reduce_mesh_faces_inplace(mesh_out, max_faces=max_faces, verbose=True) - # Auto cleanup pipeline if enabled if auto_cleanup: - try: - shapegen_pipe.to('cpu') - if hasattr(shapegen_pipe, 'unet'): - del shapegen_pipe.unet - if hasattr(shapegen_pipe, 'vae'): - del shapegen_pipe.vae - if hasattr(shapegen_pipe, 'scheduler'): - del shapegen_pipe.scheduler - del outputs - torch.cuda.empty_cache() - gc.collect() - print("Shape pipeline cleaned up") - except Exception as e: - print(f"Error during pipeline cleanup: {e}") - - mesh_out = Mesh.load_trimesh(given_mesh=mesh) - mesh_out.auto_normal() + torch.cuda.empty_cache() + gc.collect() processed_image_tensor = pils_to_torch_imgs([pil_image]) return (mesh_out, processed_image_tensor) + class Hunyuan3D_21_TexGen: """Hunyuan3D-2.1 Texture Generation""" @@ -5587,7 +5579,6 @@ class Hunyuan3D_21_TexGen: except: pass -# --------------------- PARTCRAFTER --------------------- class Load_PartCrafter_Pipeline: """Load PartCrafter Pipeline""" @@ -5790,7 +5781,7 @@ class PartCrafter_Generate: print(f"GLB mesh path for Preview_3DMesh: {relative_scene_path}") return (zip_path, relative_scene_path, processed_image_tensor) -#------ partcrafter scene --------------------- + class Load_PartCrafter_Scene_Pipeline: """Load PartCrafter Scene Pipeline""" @@ -5994,4 +5985,3 @@ class PartCrafter_Generate: print(f"GLB mesh path for Preview_3DMesh: {relative_scene_path}") return (zip_path, relative_scene_path, processed_image_tensor) -