From ae3156305dfdebd9ad91bba53cf20968caa6d55e Mon Sep 17 00:00:00 2001 From: DenisKochetov Date: Mon, 1 Sep 2025 13:04:58 +0200 Subject: [PATCH 1/3] renamed mesh_processor --- Gen_3D_Modules/CRM/model/crm/model.py | 2 +- .../hy3dpaint/mesh_inpaint_wheel/setup.py | 12 +- .../LGM/nerf_marching_cubes_converter.py | 4 +- .../MV_Adapter/mvadapter_node_utils.py | 2 +- .../TriplaneGaussian/models/renderer.py | 2 +- Gen_3D_Modules/TripoSR/models/isosurface.py | 2 +- .../Unique3D/scripts/project_mesh.py | 2 +- .../DiffRastMesh/diff_mesh_renderer.py | 2 +- .../FlexiCubes/flexicubes_trainer.py | 2 +- .../GaussianSplatting/main_3DGS_renderer.py | 4 +- __init__.py | 4 +- mesh_processer/__init__.py | 0 mesh_processer/mesh.py | 906 ------------------ mesh_processer/mesh_utils.py | 681 ------------- 14 files changed, 19 insertions(+), 1606 deletions(-) delete mode 100644 mesh_processer/__init__.py delete mode 100644 mesh_processer/mesh.py delete mode 100644 mesh_processer/mesh_utils.py 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/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_processer/mesh.py b/mesh_processer/mesh.py deleted file mode 100644 index 9bcc41b..0000000 --- a/mesh_processer/mesh.py +++ /dev/null @@ -1,906 +0,0 @@ -import os -import cv2 -from typing import NamedTuple -import torch -import trimesh -import numpy as np - -from kiui.op import safe_normalize, dot -from kiui.typing import * - -from shared_utils.sh_utils import SH2RGB -from shared_utils.image_utils import prepare_torch_img - - -class Mesh: - """ - A torch-native trimesh class, with support for ``ply/obj/glb`` formats. - - Note: - This class only supports one mesh with a single texture image (an albedo texture and a metallic-roughness texture). - """ - 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, # vertex color - albedo: Optional[Tensor] = None, - metallicRoughness: Optional[Tensor] = None, - device: Optional[torch.device] = None, - ): - """Init a mesh directly using all attributes. - - Args: - v (Optional[Tensor]): vertices, float [N, 3]. Defaults to None. - f (Optional[Tensor]): faces, int [M, 3]. Defaults to None. - vn (Optional[Tensor]): vertex normals, float [N, 3]. Defaults to None. - fn (Optional[Tensor]): faces for normals, int [M, 3]. Defaults to None. - vt (Optional[Tensor]): vertex uv coordinates, float [N, 2]. Defaults to None. - ft (Optional[Tensor]): faces for uvs, int [M, 3]. Defaults to None. - vc (Optional[Tensor]): vertex colors, float [N, 3]. Defaults to None. - albedo (Optional[Tensor]): albedo texture, float [H, W, 3], RGB format. Defaults to None. - metallicRoughness (Optional[Tensor]): metallic-roughness texture, float [H, W, 3], metallic(Blue) = metallicRoughness[..., 2], roughness(Green) = metallicRoughness[..., 1]. Defaults to None. - device (Optional[torch.device]): torch device. Defaults to None. - """ - self.device = device - self.v = v - self.vn = vn - self.vt = vt - self.f = f - self.fn = fn - self.ft = ft - # will first see if there is vertex color to use - self.vc = vc - # only support a single albedo image, (H, W, 3) - self.albedo = albedo - # pbr extension, metallic(Blue) = metallicRoughness[..., 2], roughness(Green) = metallicRoughness[..., 1] - # ref: https://registry.khronos.org/glTF/specs/2.0/glTF-2.0.html - 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): - """load mesh from path. - - Args: - path (str): path to mesh file, supports ply, obj, glb. - resize (bool, optional): auto resize the mesh using ``bound`` into [-bound, bound]^3. Defaults to True. - renormal (bool, optional): re-calc the vertex normals. Defaults to True. - retex (bool, optional): re-calc the uv coordinates, will overwrite the existing uv coordinates. Defaults to False. - clean (bool, optional): perform mesh cleaning at load (e.g., merge close vertices). Defaults to False. - bound (float, optional): bound to resize. Defaults to 0.9. - front_dir (str, optional): front-view direction of the mesh, should be [+-][xyz][ 123]. Defaults to '+z'. - device (torch.device, optional): torch device. Defaults to None. - - Note: - a ``device`` keyword argument can be provided to specify the torch device. - If it's not provided, we will try to use ``'cuda'`` as the device if it's available. - - Returns: - Mesh: the loaded Mesh object. - """ - # obj supports face uv - if path.endswith(".obj"): - mesh = cls.load_obj(path, **kwargs) - # trimesh only supports vertex uv, but can load more formats - else: - mesh = cls.load_trimesh(path, **kwargs) - - # clean - if clean: - from kiui.mesh_utils import clean_mesh - vertices = mesh.v.detach().cpu().numpy() - triangles = mesh.f.detach().cpu().numpy() - vertices, triangles = clean_mesh(vertices, triangles, remesh=False) - mesh.v = torch.from_numpy(vertices).contiguous().float().to(mesh.device) - mesh.f = torch.from_numpy(triangles).contiguous().int().to(mesh.device) - - print(f"[Mesh loading] v: {mesh.v.shape}, f: {mesh.f.shape}") - # auto-normalize - if resize: - mesh.auto_size(bound=bound) - # auto-fix normal - if renormal or mesh.vn is None: - mesh.auto_normal() - print(f"[Mesh loading] vn: {mesh.vn.shape}, fn: {mesh.fn.shape}") - # auto-fix texcoords - if retex or (mesh.albedo is not None and mesh.vt is None): - mesh.auto_uv(cache_path=path) - print(f"[Mesh loading] vt: {mesh.vt.shape}, ft: {mesh.ft.shape}") - - # rotate front dir to +z - if front_dir != "+z": - # axis switch - if "-z" in front_dir: - T = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, -1]], device=mesh.device, dtype=torch.float32) - elif "+x" in front_dir: - T = torch.tensor([[0, 0, 1], [0, 1, 0], [1, 0, 0]], device=mesh.device, dtype=torch.float32) - elif "-x" in front_dir: - T = torch.tensor([[0, 0, -1], [0, 1, 0], [1, 0, 0]], device=mesh.device, dtype=torch.float32) - elif "+y" in front_dir: - T = torch.tensor([[1, 0, 0], [0, 0, 1], [0, 1, 0]], device=mesh.device, dtype=torch.float32) - elif "-y" in front_dir: - T = torch.tensor([[1, 0, 0], [0, 0, -1], [0, 1, 0]], device=mesh.device, dtype=torch.float32) - else: - T = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], device=mesh.device, dtype=torch.float32) - # rotation (how many 90 degrees) - if '1' in front_dir: - T @= torch.tensor([[0, -1, 0], [1, 0, 0], [0, 0, 1]], device=mesh.device, dtype=torch.float32) - elif '2' in front_dir: - T @= torch.tensor([[1, 0, 0], [0, -1, 0], [0, 0, 1]], device=mesh.device, dtype=torch.float32) - elif '3' in front_dir: - T @= torch.tensor([[0, 1, 0], [-1, 0, 0], [0, 0, 1]], device=mesh.device, dtype=torch.float32) - mesh.v @= T - mesh.vn @= T - - return mesh - - # load from obj file - @classmethod - def load_obj(cls, path, albedo_path=None, device=None): - """load an ``obj`` mesh. - - Args: - path (str): path to mesh. - albedo_path (str, optional): path to the albedo texture image, will overwrite the existing texture path if specified in mtl. Defaults to None. - device (torch.device, optional): torch device. Defaults to None. - - Note: - We will try to read `mtl` path from `obj`, else we assume the file name is the same as `obj` but with `mtl` extension. - The `usemtl` statement is ignored, and we only use the last material path in `mtl` file. - - Returns: - Mesh: the loaded Mesh object. - """ - assert os.path.splitext(path)[-1] == ".obj" - - mesh = cls() - - # device - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - mesh.device = device - - # load obj - with open(path, "r") as f: - lines = f.readlines() - - def parse_f_v(fv): - # pass in a vertex term of a face, return {v, vt, vn} (-1 if not provided) - # supported forms: - # f v1 v2 v3 - # f v1/vt1 v2/vt2 v3/vt3 - # f v1/vt1/vn1 v2/vt2/vn2 v3/vt3/vn3 - # f v1//vn1 v2//vn2 v3//vn3 - xs = [int(x) - 1 if x != "" else -1 for x in fv.split("/")] - xs.extend([-1] * (3 - len(xs))) - return xs[0], xs[1], xs[2] - - vertices, texcoords, normals = [], [], [] - faces, tfaces, nfaces = [], [], [] - mtl_path = None - - for line in lines: - split_line = line.split() - # empty line - if len(split_line) == 0: - continue - prefix = split_line[0].lower() - # mtllib - if prefix == "mtllib": - mtl_path = split_line[1] - # usemtl - elif prefix == "usemtl": - pass # ignored - # v/vn/vt - elif prefix == "v": - vertices.append([float(v) for v in split_line[1:]]) - elif prefix == "vn": - normals.append([float(v) for v in split_line[1:]]) - elif prefix == "vt": - val = [float(v) for v in split_line[1:]] - texcoords.append([val[0], 1.0 - val[1]]) - elif prefix == "f": - vs = split_line[1:] - nv = len(vs) - v0, t0, n0 = parse_f_v(vs[0]) - for i in range(nv - 2): # triangulate (assume vertices are ordered) - v1, t1, n1 = parse_f_v(vs[i + 1]) - v2, t2, n2 = parse_f_v(vs[i + 2]) - faces.append([v0, v1, v2]) - tfaces.append([t0, t1, t2]) - nfaces.append([n0, n1, n2]) - - mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device) - mesh.vt = ( - torch.tensor(texcoords, dtype=torch.float32, device=device) - if len(texcoords) > 0 - else None - ) - mesh.vn = ( - torch.tensor(normals, dtype=torch.float32, device=device) - if len(normals) > 0 - else None - ) - - mesh.f = torch.tensor(faces, dtype=torch.int32, device=device) - mesh.ft = ( - torch.tensor(tfaces, dtype=torch.int32, device=device) - if len(texcoords) > 0 - else None - ) - mesh.fn = ( - torch.tensor(nfaces, dtype=torch.int32, device=device) - if len(normals) > 0 - else None - ) - - # see if there is vertex color - use_vertex_color = False - if mesh.v.shape[1] == 6: - use_vertex_color = True - mesh.vc = mesh.v[:, 3:] - mesh.v = mesh.v[:, :3] - print(f"[load_obj] use vertex color: {mesh.vc.shape}") - - # try to load texture image - if not use_vertex_color: - # try to retrieve mtl file - mtl_path_candidates = [] - if mtl_path is not None: - mtl_path_candidates.append(mtl_path) - mtl_path_candidates.append(os.path.join(os.path.dirname(path), mtl_path)) - mtl_path_candidates.append(path.replace(".obj", ".mtl")) - - mtl_path = None - for candidate in mtl_path_candidates: - if os.path.exists(candidate): - mtl_path = candidate - break - - # if albedo_path is not provided, try retrieve it from mtl - metallic_path = None - roughness_path = None - if mtl_path is not None and albedo_path is None: - with open(mtl_path, "r") as f: - lines = f.readlines() - - for line in lines: - split_line = line.split() - # empty line - if len(split_line) == 0: - continue - prefix = split_line[0] - - if "map_Kd" in prefix: - # assume relative path! - albedo_path = os.path.join(os.path.dirname(path), split_line[1]) - print(f"[load_obj] use texture from: {albedo_path}") - elif "map_Pm" in prefix: - metallic_path = os.path.join(os.path.dirname(path), split_line[1]) - elif "map_Pr" in prefix: - roughness_path = os.path.join(os.path.dirname(path), split_line[1]) - - # still not found albedo_path, or the path doesn't exist - if albedo_path is None or not os.path.exists(albedo_path): - # init an empty texture - print(f"[load_obj] init empty albedo!") - # albedo = np.random.rand(1024, 1024, 3).astype(np.float32) - albedo = np.ones((1024, 1024, 3), dtype=np.float32) * np.array([0.5, 0.5, 0.5]) # default color - else: - albedo = cv2.imread(albedo_path, cv2.IMREAD_UNCHANGED) - albedo = cv2.cvtColor(albedo, cv2.COLOR_BGR2RGB) - albedo = albedo.astype(np.float32) / 255 - print(f"[load_obj] load texture: {albedo.shape}") - - mesh.albedo = torch.tensor(albedo, dtype=torch.float32, device=device) - - # try to load metallic and roughness - if metallic_path is not None and roughness_path is not None: - print(f"[load_obj] load metallicRoughness from: {metallic_path}, {roughness_path}") - metallic = cv2.imread(metallic_path, cv2.IMREAD_UNCHANGED) - metallic = metallic.astype(np.float32) / 255 - roughness = cv2.imread(roughness_path, cv2.IMREAD_UNCHANGED) - roughness = roughness.astype(np.float32) / 255 - metallicRoughness = np.stack([np.zeros_like(metallic), roughness, metallic], axis=-1) - - mesh.metallicRoughness = torch.tensor(metallicRoughness, dtype=torch.float32, device=device).contiguous() - - return mesh - - @classmethod - def load_trimesh(cls, path=None, given_mesh=None, device=None): - """load a mesh using ``trimesh.load()``. - - Can load various formats like ``glb`` and serves as a fallback. - - Note: - We will try to merge all meshes if the glb contains more than one, - but **this may cause the texture to lose**, since we only support one texture image! - - Args: - path (str): path to the mesh file. - given_mesh (trimesh): trimesh instance - device (torch.device, optional): torch device. Defaults to None. - - Returns: - Mesh: the loaded Mesh object. - """ - mesh = cls() - - # device - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - mesh.device = device - - # use trimesh to load ply/glb, assume only has one single RootMesh... - if given_mesh is not None: - _mesh = given_mesh - elif path is not None: - _data = trimesh.load(path) - if isinstance(_data, trimesh.Scene): - if len(_data.geometry) == 1: - _mesh = list(_data.geometry.values())[0] - else: - print(f"[load_trimesh] concatenating {len(_data.geometry)} meshes.") - _concat = [] - # loop the scene graph and apply transform to each mesh - scene_graph = _data.graph.to_flattened() # dict {name: {transform: 4x4 mat, geometry: str}} - for k, v in scene_graph.items(): - name = v['geometry'] - if name in _data.geometry and isinstance(_data.geometry[name], trimesh.Trimesh): - transform = v['transform'] - _concat.append(_data.geometry[name].apply_transform(transform)) - _mesh = trimesh.util.concatenate(_concat) - else: - _mesh = _data - else: - print(f"[load_trimesh] failed to load mesh, either path or given_mesh must be given") - return None - - should_create_empty_albedo = False - if _mesh.visual.kind == 'vertex': - vertex_colors = _mesh.visual.vertex_colors - vertex_colors = np.array(vertex_colors[..., :3]).astype(np.float32) / 255 - mesh.vc = torch.tensor(vertex_colors, dtype=torch.float32, device=device) - print(f"[load_trimesh] use vertex color: {mesh.vc.shape}") - elif _mesh.visual.kind == 'texture': - _material = _mesh.visual.material - if isinstance(_material, trimesh.visual.material.PBRMaterial): - texture = np.array(_material.baseColorTexture).astype(np.float32) / 255 - # load metallicRoughness if present - if _material.metallicRoughnessTexture is not None: - metallicRoughness = np.array(_material.metallicRoughnessTexture).astype(np.float32) / 255 - mesh.metallicRoughness = torch.tensor(metallicRoughness, dtype=torch.float32, device=device).contiguous() - elif isinstance(_material, trimesh.visual.material.SimpleMaterial): - texture = np.array(_material.to_pbr().baseColorTexture).astype(np.float32) / 255 - else: - raise NotImplementedError(f"material type {type(_material)} not supported!") - - if texture.ndim ==3: - mesh.albedo = torch.tensor(texture[..., :3], dtype=torch.float32, device=device).contiguous() - print(f"[load_trimesh] loaded albedo texture: {texture.shape}") - else: - should_create_empty_albedo = True - else: - should_create_empty_albedo = True - - if should_create_empty_albedo: - mesh.set_new_albedo(1024, 1024) - print(f"[load_trimesh] failed to load albedo texture, create empty abbedo texture with shape: {mesh.albedo.shape}") - - vertices = _mesh.vertices - - try: - texcoords = _mesh.visual.uv - texcoords[:, 1] = 1 - texcoords[:, 1] - except Exception as e: - texcoords = None - - try: - normals = _mesh.vertex_normals - except Exception as e: - normals = None - - # trimesh only support vertex uv... - faces = tfaces = nfaces = _mesh.faces - - mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device) - mesh.vt = ( - torch.tensor(texcoords, dtype=torch.float32, device=device) - if texcoords is not None - else None - ) - mesh.vn = ( - torch.tensor(normals, dtype=torch.float32, device=device) - if normals is not None - else None - ) - - mesh.f = torch.tensor(faces, dtype=torch.int32, device=device) - mesh.ft = ( - torch.tensor(tfaces, dtype=torch.int32, device=device) - if texcoords is not None - else None - ) - mesh.fn = ( - torch.tensor(nfaces, dtype=torch.int32, device=device) - if normals is not None - else None - ) - - return mesh - - 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) - - # aabb - def aabb(self): - """get the axis-aligned bounding box of the mesh. - - Returns: - Tuple[torch.Tensor]: the min xyz and max xyz of the mesh. - """ - return torch.min(self.v, dim=0).values, torch.max(self.v, dim=0).values - - # unit size - @torch.no_grad() - def auto_size(self, bound=0.9): - """auto resize the mesh. - - Args: - bound (float, optional): resizing into ``[-bound, bound]^3``. Defaults to 0.9. - """ - 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): - """auto calculate the vertex 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) - - # Splat face normals to vertices - 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) - - # Normalize, replace zero (degenerated) normals with some default value - 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): - """auto calculate the uv coordinates. - - Args: - cache_path (str, optional): path to save/load the uv cache as a npz file, this can avoid calculating uv every time when loading the same mesh, which is time-consuming. Defaults to None. - vmap (bool, optional): remap vertices based on uv coordinates, so each v correspond to a unique vt (necessary for formats like gltf). - Usually this will duplicate the vertices on the edge of uv atlas. Defaults to True. - """ - # try to load cache - if cache_path is not None: - cache_path = os.path.splitext(cache_path)[0] + "_uv.npz" - if cache_path is not None and os.path.exists(cache_path): - data = np.load(cache_path) - vt_np, ft_np, vmapping = data["vt"], data["ft"], data["vmapping"] - else: - import xatlas - - v_np = self.v.detach().cpu().numpy() - f_np = self.f.detach().int().cpu().numpy() - atlas = xatlas.Atlas() - atlas.add_mesh(v_np, f_np) - chart_options = xatlas.ChartOptions() - # chart_options.max_iterations = 4 - atlas.generate(chart_options=chart_options) - vmapping, ft_np, vt_np = atlas[0] # [N], [M, 3], [N, 2] - - # save to cache - if cache_path is not None: - np.savez(cache_path, vt=vt_np, ft=ft_np, vmapping=vmapping) - - vt = torch.from_numpy(vt_np.astype(np.float32)).to(self.device) - ft = torch.from_numpy(ft_np.astype(np.int32)).to(self.device) - self.vt = vt - self.ft = ft - - if vmap: - # remap v/f to vt/ft, so each v correspond to a unique vt. (necessary for gltf) - vmapping = torch.from_numpy(vmapping.astype(np.int64)).long().to(self.device) - self.align_v_to_vt(vmapping) - - def align_v_to_vt(self, vmapping=None): - """ remap v/f and vn/fn to vt/ft. - - Args: - vmapping (np.ndarray, optional): the mapping relationship from f to ft. Defaults to None. - """ - if vmapping is None: - vt2v_mapping = self.get_default_vt_to_v_mapping() - else: - vt2v_mapping = vmapping - - self.v = self.v[vt2v_mapping] - if self.vc is not None: - self.vc = self.vc[vt2v_mapping] - self.f = self.ft - # assume fn == f - if self.vn is not None: - if vmapping is None: - vt2vn_mapping = self.get_default_vt_to_vn_mapping() - else: - vt2vn_mapping = vmapping - self.vn = self.vn[vt2vn_mapping] - self.fn = self.ft - - def get_default_vt_to_v_mapping(self): - """map from ft to f - - Returns: - vmapping: tensor array - """ - ft = self.ft.view(-1).long() - f = self.f.view(-1).long() - vmapping = torch.zeros(self.vt.shape[0], dtype=torch.long, device=self.device) - vmapping[ft] = f # scatter, randomly choose one if index is not unique - return vmapping - - def get_default_vt_to_vn_mapping(self): - """map from ft to f - - Returns: - vmapping: tensor array - """ - ft = self.ft.view(-1).long() - fn = self.fn.view(-1).long() - vmapping = torch.zeros(self.vt.shape[0], dtype=torch.long, device=self.device) - vmapping[ft] = fn # scatter, randomly choose one if index is not unique - return vmapping - - def to(self, device): - """move all tensor attributes to device. - - Args: - device (torch.device): target device. - - Returns: - Mesh: self. - """ - 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): - """write the mesh to a path. - - Args: - path (str): path to write, supports ply, obj and glb. - """ - if path.endswith(".ply"): - self.write_ply(path) - elif path.endswith(".obj"): - self.write_obj(path) - elif path.endswith(".glb") or path.endswith(".gltf"): - self.write_glb(path) - else: - raise NotImplementedError(f"format {path} not supported!") - - def write_ply(self, path): - """write the mesh in ply format. Only for geometry! - - Args: - path (str): path to write. - """ - - if self.albedo is not None: - print(f'[WARN] ply format does not support exporting texture, will ignore!') - - 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) - - - def write_glb(self, path): - """write the mesh in glb/gltf format. - This will create a scene with a single mesh. - - Args: - path (str): path to write. - """ - - # assert self.v.shape[0] == self.vn.shape[0] and self.v.shape[0] == self.vt.shape[0] - if self.vt is not None and self.v.shape[0] != self.vt.shape[0]: - self.align_v_to_vt() - - import pygltflib - - f_np = self.f.detach().cpu().numpy().astype(np.uint32) - f_np_blob = f_np.flatten().tobytes() - - v_np = self.v.detach().cpu().numpy().astype(np.float32) - v_np_blob = v_np.tobytes() - - blob = f_np_blob + v_np_blob - byteOffset = len(blob) - - # base mesh - gltf = pygltflib.GLTF2( - scene=0, - scenes=[pygltflib.Scene(nodes=[0])], - nodes=[pygltflib.Node(mesh=0)], - meshes=[pygltflib.Mesh(primitives=[pygltflib.Primitive( - # indices to accessors (0 is triangles) - attributes=pygltflib.Attributes( - POSITION=1, - ), - indices=0, - )])], - buffers=[ - pygltflib.Buffer(byteLength=len(f_np_blob) + len(v_np_blob)) - ], - # buffer view (based on dtype) - bufferViews=[ - # triangles; as flatten (element) array - pygltflib.BufferView( - buffer=0, - byteLength=len(f_np_blob), - target=pygltflib.ELEMENT_ARRAY_BUFFER, # GL_ELEMENT_ARRAY_BUFFER (34963) - ), - # positions; as vec3 array - pygltflib.BufferView( - buffer=0, - byteOffset=len(f_np_blob), - byteLength=len(v_np_blob), - byteStride=12, # vec3 - target=pygltflib.ARRAY_BUFFER, # GL_ARRAY_BUFFER (34962) - ), - ], - accessors=[ - # 0 = triangles - pygltflib.Accessor( - bufferView=0, - componentType=pygltflib.UNSIGNED_INT, # GL_UNSIGNED_INT (5125) - count=f_np.size, - type=pygltflib.SCALAR, - max=[int(f_np.max())], - min=[int(f_np.min())], - ), - # 1 = positions - pygltflib.Accessor( - bufferView=1, - componentType=pygltflib.FLOAT, # GL_FLOAT (5126) - count=len(v_np), - type=pygltflib.VEC3, - max=v_np.max(axis=0).tolist(), - min=v_np.min(axis=0).tolist(), - ), - ], - ) - - # append texture info - if self.vt is not None and self.albedo is not None: - - vt_np = self.vt.detach().cpu().numpy().astype(np.float32) - vt_np_blob = vt_np.tobytes() - - albedo = self.albedo.detach().cpu().numpy() - albedo = (albedo * 255).astype(np.uint8) - albedo = cv2.cvtColor(albedo, cv2.COLOR_RGB2BGR) - albedo_blob = cv2.imencode('.png', albedo)[1].tobytes() - - # update primitive - gltf.meshes[0].primitives[0].attributes.TEXCOORD_0 = 2 - gltf.meshes[0].primitives[0].material = 0 - - # update materials - gltf.materials.append(pygltflib.Material( - pbrMetallicRoughness=pygltflib.PbrMetallicRoughness( - baseColorTexture=pygltflib.TextureInfo(index=0, texCoord=0), - metallicFactor=0.0, - roughnessFactor=1.0, - ), - alphaMode=pygltflib.OPAQUE, - alphaCutoff=None, - doubleSided=True, - )) - - gltf.textures.append(pygltflib.Texture(sampler=0, source=0)) - gltf.samplers.append(pygltflib.Sampler(magFilter=pygltflib.LINEAR, minFilter=pygltflib.LINEAR_MIPMAP_LINEAR, wrapS=pygltflib.REPEAT, wrapT=pygltflib.REPEAT)) - gltf.images.append(pygltflib.Image(bufferView=3, mimeType="image/png")) - - # update buffers - gltf.bufferViews.append( - # index = 2, texcoords; as vec2 array - pygltflib.BufferView( - buffer=0, - byteOffset=byteOffset, - byteLength=len(vt_np_blob), - byteStride=8, # vec2 - target=pygltflib.ARRAY_BUFFER, - ) - ) - - gltf.accessors.append( - # 2 = texcoords - pygltflib.Accessor( - bufferView=2, - componentType=pygltflib.FLOAT, - count=len(vt_np), - type=pygltflib.VEC2, - max=vt_np.max(axis=0).tolist(), - min=vt_np.min(axis=0).tolist(), - ) - ) - - blob += vt_np_blob - byteOffset += len(vt_np_blob) - - gltf.bufferViews.append( - # index = 3, albedo texture; as none target - pygltflib.BufferView( - buffer=0, - byteOffset=byteOffset, - byteLength=len(albedo_blob), - ) - ) - - blob += albedo_blob - byteOffset += len(albedo_blob) - - gltf.buffers[0].byteLength = byteOffset - - # append metllic roughness - if self.metallicRoughness is not None: - metallicRoughness = self.metallicRoughness.detach().cpu().numpy() - metallicRoughness = (metallicRoughness * 255).astype(np.uint8) - metallicRoughness = cv2.cvtColor(metallicRoughness, cv2.COLOR_RGB2BGR) - metallicRoughness_blob = cv2.imencode('.png', metallicRoughness)[1].tobytes() - - # update texture definition - gltf.materials[0].pbrMetallicRoughness.metallicFactor = 1.0 - gltf.materials[0].pbrMetallicRoughness.roughnessFactor = 1.0 - gltf.materials[0].pbrMetallicRoughness.metallicRoughnessTexture = pygltflib.TextureInfo(index=1, texCoord=0) - - gltf.textures.append(pygltflib.Texture(sampler=1, source=1)) - gltf.samplers.append(pygltflib.Sampler(magFilter=pygltflib.LINEAR, minFilter=pygltflib.LINEAR_MIPMAP_LINEAR, wrapS=pygltflib.REPEAT, wrapT=pygltflib.REPEAT)) - gltf.images.append(pygltflib.Image(bufferView=4, mimeType="image/png")) - - # update buffers - gltf.bufferViews.append( - # index = 4, metallicRoughness texture; as none target - pygltflib.BufferView( - buffer=0, - byteOffset=byteOffset, - byteLength=len(metallicRoughness_blob), - ) - ) - - blob += metallicRoughness_blob - byteOffset += len(metallicRoughness_blob) - - gltf.buffers[0].byteLength = byteOffset - - - # set actual data - gltf.set_binary_blob(blob) - - # glb = b"".join(gltf.save_to_bytes()) - gltf.save(path) - - - def write_obj(self, path): - """write the mesh in obj format. Will also write the texture and mtl files. - - Args: - path (str): path to write. - """ - - mtl_path = path.replace(".obj", ".mtl") - albedo_path = path.replace(".obj", "_albedo.png") - metallic_path = path.replace(".obj", "_metallic.png") - roughness_path = path.replace(".obj", "_roughness.png") - - v_np = self.v.detach().cpu().numpy() - vt_np = self.vt.detach().cpu().numpy() if self.vt is not None else None - vn_np = self.vn.detach().cpu().numpy() if self.vn is not None else None - vc_np = self.vc.detach().cpu().numpy() if self.vc is not None else None - f_np = self.f.detach().cpu().numpy() - ft_np = self.ft.detach().cpu().numpy() if self.ft is not None else None - fn_np = self.fn.detach().cpu().numpy() if self.fn is not None else None - - with open(path, "w") as fp: - fp.write(f"mtllib {os.path.basename(mtl_path)} \n") - - if vc_np is None: - for v in v_np: - fp.write(f"v {v[0]} {v[1]} {v[2]} \n") - else: - for v_i in range(v_np.shape[0]): - v = v_np[v_i] - vc = vc_np[v_i] - fp.write(f"v {v[0]} {v[1]} {v[2]} {vc[0]} {vc[1]} {vc[2]}\n") - - if vt_np is not None: - for v in vt_np: - fp.write(f"vt {v[0]} {1 - v[1]} \n") - - if vn_np is not None: - for v in vn_np: - fp.write(f"vn {v[0]} {v[1]} {v[2]} \n") - - fp.write(f"usemtl defaultMat \n") - for i in range(len(f_np)): - fp.write( - f'f {f_np[i, 0] + 1}/{ft_np[i, 0] + 1 if ft_np is not None else ""}/{fn_np[i, 0] + 1 if fn_np is not None else ""} \ - {f_np[i, 1] + 1}/{ft_np[i, 1] + 1 if ft_np is not None else ""}/{fn_np[i, 1] + 1 if fn_np is not None else ""} \ - {f_np[i, 2] + 1}/{ft_np[i, 2] + 1 if ft_np is not None else ""}/{fn_np[i, 2] + 1 if fn_np is not None else ""} \n' - ) - - with open(mtl_path, "w") as fp: - fp.write(f"newmtl defaultMat \n") - fp.write(f"Ka 1 1 1 \n") - fp.write(f"Kd 1 1 1 \n") - fp.write(f"Ks 1 1 1 \n") - #fp.write(f"Tr 1 \n") # will cause the mesh materials in three.js completely transparent - fp.write(f"illum 1 \n") - fp.write(f"Ns 10 \n") - if self.albedo is not None: - fp.write(f"map_Kd {os.path.basename(albedo_path)} \n") - if self.metallicRoughness is not None: - # ref: https://en.wikipedia.org/wiki/Wavefront_.obj_file#Physically-based_Rendering - fp.write(f"map_Pm {os.path.basename(metallic_path)} \n") - fp.write(f"map_Pr {os.path.basename(roughness_path)} \n") - - if self.albedo is not None: - albedo = self.albedo.detach().cpu().numpy() - albedo = (albedo * 255).astype(np.uint8) - cv2.imwrite(albedo_path, cv2.cvtColor(albedo, cv2.COLOR_RGB2BGR)) - - if self.metallicRoughness is not None: - metallicRoughness = self.metallicRoughness.detach().cpu().numpy() - metallicRoughness = (metallicRoughness * 255).astype(np.uint8) - cv2.imwrite(metallic_path, metallicRoughness[..., 2]) - cv2.imwrite(roughness_path, metallicRoughness[..., 1]) - - def convert_to_pointcloud(self): - xyz = self.v.detach().cpu().numpy() - num_pts = self.v.shape[0] - 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 - - -class PointCloud(NamedTuple): - points: np.array - colors: np.array - normals: np.array \ No newline at end of file diff --git a/mesh_processer/mesh_utils.py b/mesh_processer/mesh_utils.py deleted file mode 100644 index c3799df..0000000 --- a/mesh_processer/mesh_utils.py +++ /dev/null @@ -1,681 +0,0 @@ -import torch -import numpy as np -from kornia.geometry.conversions import ( - quaternion_to_axis_angle, - axis_angle_to_quaternion, -) -import pymeshlab as pml -from plyfile import PlyData, PlyElement - -from typing import Optional - -pytorch3d_capable = True -try: - import pytorch3d - from pytorch3d.ops import knn_points -except ImportError: - pytorch3d_capable = False - -from .mesh import PointCloud -from shared_utils.sh_utils import SH2RGB, RGB2SH - -def _base_face_areas(face_vertices_0, face_vertices_1, face_vertices_2): - """Base function to compute the face areas.""" - x1, x2, x3 = torch.split(face_vertices_0 - face_vertices_1, 1, dim=-1) - y1, y2, y3 = torch.split(face_vertices_1 - face_vertices_2, 1, dim=-1) - - a = (x2 * y3 - x3 * y2) ** 2 - b = (x3 * y1 - x1 * y3) ** 2 - c = (x1 * y2 - x2 * y1) ** 2 - areas = torch.sqrt(a + b + c) * 0.5 - - return areas - - -def _base_sample_points_selected_faces(face_vertices, face_features=None): - """Base function to sample points over selected faces, sample one point per face. - The coordinates of the face vertices are interpolated to generate new samples. - - Args: - face_vertices (tuple of torch.Tensor): - Coordinates of vertices, corresponding to selected faces to sample from. - A tuple of 3 entries corresponding to each of the face vertices. - Each entry is a torch.Tensor of shape :math:`(\\text{batch_size}, \\text{num_samples}, 3)`. - face_features (tuple of torch.Tensor, Optional): - Features of face vertices, corresponding to selected faces to sample from. - A tuple of 3 entries corresponding to each of the face vertices. - Each entry is a torch.Tensor of shape - :math:`(\\text{batch_size}, \\text{num_samples}, \\text{feature_dim})`. - - Returns: - (torch.Tensor, torch.Tensor): - Sampled point coordinates of shape :math:`(\\text{batch_size}, \\text{num_samples}, 3)`. - Sampled points interpolated features of shape - :math:`(\\text{batch_size}, \\text{num_samples}, \\text{feature_dim})`. - If `face_vertices_features` arg is not specified, the returned interpolated features are None. - """ - - face_vertices0, face_vertices1, face_vertices2 = face_vertices - - sampling_shape = tuple(int(d) for d in face_vertices0.shape[:-1]) + (1,) - # u is proximity to middle point between v1 and v2 against v0. - # v is proximity to v2 against v1. - # - # The probability density for u should be f_U(u) = 2u. - # However, torch.rand use a uniform (f_X(x) = x) distribution, - # so using torch.sqrt we make a change of variable to have the desired density - # f_Y(y) = f_X(y ^ 2) * |d(y ^ 2) / dy| = 2y - u = torch.sqrt(torch.rand(sampling_shape, - device=face_vertices0.device, - dtype=face_vertices0.dtype)) - - v = torch.rand(sampling_shape, - device=face_vertices0.device, - dtype=face_vertices0.dtype) - w0 = 1 - u - w1 = u * (1 - v) - w2 = u * v - - points = w0 * face_vertices0 + w1 * face_vertices1 + w2 * face_vertices2 - - features = None - if face_features is not None: - face_features0, face_features1, face_features2 = face_features - features = w0 * face_features0 + w1 * face_features1 + \ - w2 * face_features2 - - return points, features - -# Modified from https://github.com/NVIDIAGameWorks/kaolin/blob/master/kaolin/ops/mesh/trianglemesh.py#L159 -def sample_points(vertices, faces, num_samples, areas=None, face_features=None): - r"""Uniformly sample points over the surface of triangle meshes. - - First, face on which the point is sampled is randomly selected, - with the probability of selection being proportional to the area of the face. - then the coordinate on the face is uniformly sampled. - - If ``face_features`` is defined for the mesh faces, - the sampled points will be returned with interpolated features as well, - otherwise, no feature interpolation will occur. - - Args: - vertices (torch.Tensor): - The vertices of the meshes, of shape - :math:`(\text{batch_size}, \text{num_vertices}, 3)`. - faces (torch.LongTensor): - The faces of the mesh, of shape :math:`(\text{num_faces}, 3)`. - num_samples (int): - The number of point sampled per mesh. - Also the number of faces sampled per mesh, and then sample a single point per face. - areas (torch.Tensor, optional): - The areas of each face, of shape :math:`(\text{batch_size}, \text{num_faces})`, - can be preprocessed, for fast on-the-fly sampling, - will be computed if None (default). - face_features (torch.Tensor, optional): - Per-vertex-per-face features, matching ``faces`` order, - of shape :math:`(\text{batch_size}, \text{num_faces}, 3, \text{feature_dim})`. - For example: - - 1. Texture uv coordinates would be of shape - :math:`(\text{batch_size}, \text{num_faces}, 3, 2)`. - 2. RGB color values would be of shape - :math:`(\text{batch_size}, \text{num_faces}, 3, 3)`. - - When specified, it is used to interpolate the features for new sampled points. - - See also: - :func:`~kaolin.ops.mesh.index_vertices_by_faces` for conversion of features defined per vertex - and need to be converted to per-vertex-per-face shape of :math:`(\text{num_faces}, 3)`. - - Returns: - (torch.Tensor, torch.LongTensor, (optional) torch.Tensor): - the pointclouds of shape :math:`(\text{batch_size}, \text{num_samples}, 3)`, - and the indexes of the faces selected, - of shape :math:`(\text{batch_size}, \text{num_samples})`. - - If ``face_features`` arg is specified, then the interpolated features of sampled points of shape - :math:`(\text{batch_size}, \text{num_samples}, \text{feature_dim})` are also returned. - """ - if faces.shape[-1] != 3: - raise NotImplementedError("sample_points is only implemented for triangle meshes") - faces_0, faces_1, faces_2 = torch.split(faces, 1, dim=1) # (num_faces, 3) -> tuple of (num_faces,) - face_v_0 = torch.index_select(vertices, 1, faces_0.reshape(-1)) # (batch_size, num_faces, 3) - face_v_1 = torch.index_select(vertices, 1, faces_1.reshape(-1)) # (batch_size, num_faces, 3) - face_v_2 = torch.index_select(vertices, 1, faces_2.reshape(-1)) # (batch_size, num_faces, 3) - - if areas is None: - areas = _base_face_areas(face_v_0, face_v_1, face_v_2).squeeze(-1) - face_dist = torch.distributions.Categorical(areas) - face_choices = face_dist.sample([num_samples]).transpose(0, 1) - _face_choices = face_choices.unsqueeze(-1).repeat(1, 1, 3) - v0 = torch.gather(face_v_0, 1, _face_choices) # (batch_size, num_samples, 3) - v1 = torch.gather(face_v_1, 1, _face_choices) # (batch_size, num_samples, 3) - v2 = torch.gather(face_v_2, 1, _face_choices) # (batch_size, num_samples, 3) - face_vertices_choices = (v0, v1, v2) - - # UV coordinates are available, make sure to calculate them for sampled points as well - face_features_choices = None - if face_features is not None: - feat_dim = face_features.shape[-1] - # (num_faces, 3) -> tuple of (num_faces,) - _face_choices = face_choices[..., None, None].repeat(1, 1, 3, feat_dim) - face_features_choices = torch.gather(face_features, 1, _face_choices) - face_features_choices = tuple( - tmp_feat.squeeze(2) for tmp_feat in torch.split(face_features_choices, 1, dim=2)) - - points, point_features = _base_sample_points_selected_faces( - face_vertices_choices, face_features_choices) - - if point_features is not None: - return points, face_choices, point_features - else: - return points, face_choices - -def poisson_mesh_reconstruction(points, normals=None): - # points/normals: [N, 3] np.ndarray - - import open3d as o3d - - pcd = o3d.geometry.PointCloud() - pcd.points = o3d.utility.Vector3dVector(points) - - # outlier removal - pcd, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=10) - - # normals - if normals is None: - pcd.estimate_normals() - else: - pcd.normals = o3d.utility.Vector3dVector(normals[ind]) - - # visualize - o3d.visualization.draw_geometries([pcd], point_show_normal=False) - - mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson( - pcd, depth=9 - ) - vertices_to_remove = densities < np.quantile(densities, 0.1) - mesh.remove_vertices_by_mask(vertices_to_remove) - - # visualize - o3d.visualization.draw_geometries([mesh]) - - vertices = np.asarray(mesh.vertices) - triangles = np.asarray(mesh.triangles) - - print( - f"[INFO] poisson mesh reconstruction: {points.shape} --> {vertices.shape} / {triangles.shape}" - ) - - return vertices, triangles - - -def decimate_mesh( - verts, faces, target=5e4, backend="pymeshlab", remesh=False, optimalplacement=True, verbose=True -): - """ perform mesh decimation. - - Args: - verts (np.ndarray): mesh vertices, float [N, 3] - faces (np.ndarray): mesh faces, int [M, 3] - target (int): targeted number of faces - backend (str, optional): algorithm backend, can be "pymeshlab" or "pyfqmr". Defaults to "pymeshlab". - remesh (bool, optional): whether to remesh after decimation. Defaults to False. - optimalplacement (bool, optional): For flat mesh, use False to prevent spikes. Defaults to True. - verbose (bool, optional): whether to print the decimation process. Defaults to True. - - Returns: - Tuple[np.ndarray]: vertices and faces after decimation. - """ - - _ori_vert_shape = verts.shape - _ori_face_shape = faces.shape - - if backend == "pyfqmr": - import pyfqmr - - solver = pyfqmr.Simplify() - solver.setMesh(verts, faces) - solver.simplify_mesh(target_count=target, preserve_border=False, verbose=False) - verts, faces, normals = solver.getMesh() - else: - m = pml.Mesh(verts, faces) - ms = pml.MeshSet() - ms.add_mesh(m, "mesh") # will copy! - - ms.meshing_merge_close_vertices() - - # filters - # ms.meshing_decimation_clustering(threshold=pml.PercentageValue(1)) - ms.meshing_decimation_quadric_edge_collapse( - targetfacenum=int(target), optimalplacement=optimalplacement - ) - - if remesh: - # ms.apply_coord_taubin_smoothing() - ms.meshing_isotropic_explicit_remeshing( - iterations=3, targetlen=pml.PercentageValue(1) - ) - - # extract mesh - m = ms.current_mesh() - m.compact() - verts = m.vertex_matrix() - faces = m.face_matrix() - - if verbose: - print(f"[INFO] mesh decimation: {_ori_vert_shape} --> {verts.shape}, {_ori_face_shape} --> {faces.shape}") - - return verts, faces - - -def clean_mesh( - verts, - faces, - v_pct=1, - min_f=64, - min_d=20, - repair=True, - remesh=True, - remesh_size=0.01, -): - # verts: [N, 3] - # faces: [N, 3] - - _ori_vert_shape = verts.shape - _ori_face_shape = faces.shape - - m = pml.Mesh(verts, faces) - ms = pml.MeshSet() - ms.add_mesh(m, "mesh") # will copy! - - # filters - ms.meshing_remove_unreferenced_vertices() # verts not refed by any faces - - if v_pct > 0: - ms.meshing_merge_close_vertices( - threshold=pml.Percentage(v_pct) - ) # 1/10000 of bounding box diagonal - - ms.meshing_remove_duplicate_faces() # faces defined by the same verts - ms.meshing_remove_null_faces() # faces with area == 0 - - if min_d > 0: - ms.meshing_remove_connected_component_by_diameter( - mincomponentdiag=pml.Percentage(min_d) - ) - - if min_f > 0: - ms.meshing_remove_connected_component_by_face_number(mincomponentsize=min_f) - - if repair: - # ms.meshing_remove_t_vertices(method=0, threshold=40, repeat=True) - ms.meshing_repair_non_manifold_edges(method=0) - ms.meshing_repair_non_manifold_vertices(vertdispratio=0) - - if remesh: - # ms.apply_coord_taubin_smoothing() - ms.meshing_isotropic_explicit_remeshing( - iterations=3, targetlen=pml.AbsoluteValue(remesh_size) - ) - - # extract mesh - m = ms.current_mesh() - verts = m.vertex_matrix() - faces = m.face_matrix() - - print( - f"[INFO] mesh cleaning: {_ori_vert_shape} --> {verts.shape}, {_ori_face_shape} --> {faces.shape}" - ) - - return verts, faces - -def construct_list_of_gs_attributes(features_dc, features_rest, scaling, rotation): - l = ['x', 'y', 'z', 'nx', 'ny', 'nz'] - for i in range(features_dc.shape[1]*features_dc.shape[2]): - l.append('f_dc_{}'.format(i)) - for i in range(features_rest.shape[1]*features_rest.shape[2]): - l.append('f_rest_{}'.format(i)) - l.append('opacity') - for i in range(scaling.shape[1]): - l.append('scale_{}'.format(i)) - for i in range(rotation.shape[1]): - l.append('rot_{}'.format(i)) - return l - -def calculate_max_sh_degree_from_gs_ply(plydata): - extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")] - #assert len(extra_f_names)!=3*(max_sh_degree + 1) ** 2 - 3: - max_sh_degree = int(((len(extra_f_names) + 3) / 3) ** 0.5 - 1) - return max_sh_degree, extra_f_names - -def write_gs_ply(xyz, normals, f_dc, f_rest, opacities, scale, rotation, list_of_attributes): - dtype_full = [(attribute, 'f4') for attribute in list_of_attributes] - - elements = np.empty(xyz.shape[0], dtype=dtype_full) - attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1) - elements[:] = list(map(tuple, attributes)) - el = PlyElement.describe(elements, 'vertex') - return PlyData([el]) - -def read_gs_ply(plydata): - xyz = np.stack((np.asarray(plydata.elements[0]["x"]), - np.asarray(plydata.elements[0]["y"]), - np.asarray(plydata.elements[0]["z"])), axis=1) - opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis] - - features_dc = np.zeros((xyz.shape[0], 3, 1)) - features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"]) - features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"]) - features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"]) - - max_sh_degree, extra_f_names = calculate_max_sh_degree_from_gs_ply(plydata) - - features_extra = np.zeros((xyz.shape[0], len(extra_f_names))) - for idx, attr_name in enumerate(extra_f_names): - features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name]) - # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC) - features_extra = features_extra.reshape((features_extra.shape[0], 3, (max_sh_degree + 1) ** 2 - 1)) - - scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")] - scales = np.zeros((xyz.shape[0], len(scale_names))) - for idx, attr_name in enumerate(scale_names): - scales[:, idx] = np.asarray(plydata.elements[0][attr_name]) - - rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")] - rots = np.zeros((xyz.shape[0], len(rot_names))) - for idx, attr_name in enumerate(rot_names): - rots[:, idx] = np.asarray(plydata.elements[0][attr_name]) - - return xyz, features_dc, features_extra, opacities, scales, rots - -def ply_to_points_cloud(plydata): - xyz, features_dc, features_extra, opacities, scales, rots = read_gs_ply(plydata) - - features_dc = np.transpose(features_dc, (0, 2, 1)) # equivalent of torch.transpose(features_dc, 1, 2) - features_extra = np.transpose(features_extra, (0, 2, 1)) - shs = np.concatenate((features_dc, features_extra), axis=1) - normals = np.zeros_like(xyz) - pcd = PointCloud(points=xyz, colors=SH2RGB(shs), normals=normals) - return pcd - - -def get_target_axis_and_scale(axis_string, scale_value=1.0): - """ - Coordinate system inverts when: - 1. Any of the axis inverts - 2. Two of the axises switch - - If coordinate system inverts twice in a row then it will not be inverted - """ - axis_names = ["x", "y", "z"] - - target_axis, target_scale, coordinate_invert_count = [], [], 0 - axis_switch_count = 0 - for i in range(len(axis_names)): - s = axis_string[i] - if s[0] == "-": - target_scale.append(-scale_value) - coordinate_invert_count += 1 - else: - target_scale.append(scale_value) - - new_axis_i = axis_names.index(s[1]) - if new_axis_i != i: - axis_switch_count += 1 - target_axis.append(new_axis_i) - - if axis_switch_count == 2: - coordinate_invert_count += 1 - - return target_axis, target_scale, coordinate_invert_count - -def switch_vector_axis(vector3s, target_axis): - """ - Example: - vector3s = torch.tensor([[1, 2, 3], [3, 2, 1], [2, 3, 1]]) # shape (N, 3) - - target_axis = (2, 0, 1) # or [2, 0, 1] - vector3s[:, [0, 1, 2]] = vector3s[:, target_axis] - - # Result: tensor([[3, 1, 2], [1, 3, 2], [1, 2, 3]]) - """ - vector3s[:, [0, 1, 2]] = vector3s[:, target_axis] - return vector3s - -def switch_ply_axis_and_scale(plydata, target_axis, target_scale, coordinate_invert_count): - """ - Args: - target_axis (array): shape (3) - target_scale (array): shape (3) - """ - xyz, features_dc, features_extra, opacities, scales, rots = read_gs_ply(plydata) - normals = np.zeros_like(xyz) - features_dc_2d = features_dc.reshape(features_dc.shape[0], features_dc.shape[1]*features_dc.shape[2]) - features_extra_2d = features_extra.reshape(features_extra.shape[0], features_extra.shape[1]*features_extra.shape[2]) - - target_scale = torch.tensor(target_scale).float().cuda() - xyz = switch_vector_axis(torch.tensor(xyz).float().cuda() * target_scale, target_axis).detach().cpu().numpy() - scales = switch_vector_axis(torch.tensor(scales).float().cuda(), target_axis).detach().cpu().numpy() - - # change rotation representation from quaternion (w, x, y, z) to axis angle vector (x, y, z) to make swich axis easier - rots_axis_angle = quaternion_to_axis_angle(torch.tensor(rots).float().cuda()) - rots_axis_angle = switch_vector_axis(rots_axis_angle * target_scale, target_axis) - """ - Since axis–angle vector is composed of axis (unit vector/direction) and clockwise radians angle (vector magnitude), - so in order to invert the sign of angle when coordinate system inverts, we also need to invert the direction of axis–angle vector - """ - if coordinate_invert_count % 2 != 0: - rots_axis_angle = -rots_axis_angle - rots = axis_angle_to_quaternion(rots_axis_angle).detach().cpu().numpy() - - return write_gs_ply(xyz, normals, features_dc_2d, features_extra_2d, opacities, scales, rots, construct_list_of_gs_attributes(features_dc, features_extra, scales, rots)) - -def switch_mesh_axis_and_scale(mesh, target_axis, target_scale, flip_normal=False): - """ - Args: - target_axis (array): shape (3) - target_scale (array): shape (3) - """ - target_scale = torch.tensor(target_scale).float().cuda() - mesh.v = switch_vector_axis(mesh.v * target_scale, target_axis) - mesh.vn = switch_vector_axis(mesh.vn * target_scale, target_axis) - if flip_normal: - mesh.vn *= -1 - return mesh - - -def marching_cubes_density_to_mesh(get_density_func, grid_size=256, S=128, density_thresh=10, decimate_target=5e4): - from mcubes import marching_cubes - from kiui.mesh_utils import clean_mesh, decimate_mesh - - - sigmas = np.zeros([grid_size, grid_size, grid_size], dtype=np.float32) - - X = torch.linspace(-1, 1, grid_size).split(S) - Y = torch.linspace(-1, 1, grid_size).split(S) - Z = torch.linspace(-1, 1, grid_size).split(S) - - for xi, xs in enumerate(X): - for yi, ys in enumerate(Y): - for zi, zs in enumerate(Z): - xx, yy, zz = torch.meshgrid(xs, ys, zs, indexing='ij') - pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1) # [S, 3] - val = get_density_func(pts) - sigmas[xi * S: xi * S + len(xs), yi * S: yi * S + len(ys), zi * S: zi * S + len(zs)] = val.reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy() # [S, 1] --> [x, y, z] - - print(f'[INFO] marching cubes thresh: {density_thresh} ({sigmas.min()} ~ {sigmas.max()})') - - vertices, triangles = marching_cubes(sigmas, density_thresh) - vertices = vertices / (grid_size - 1.0) * 2 - 1 - - # clean - vertices = vertices.astype(np.float32) - triangles = triangles.astype(np.int32) - vertices, triangles = clean_mesh(vertices, triangles, remesh=True, remesh_size=0.01) - if triangles.shape[0] > decimate_target: - vertices, triangles = decimate_mesh(vertices, triangles, decimate_target, optimalplacement=False) - - return vertices, triangles - -def color_func_to_albedo(mesh, get_rgb_func, texture_resolution=1024, padding=2, batch_size=640000, device="cuda", force_cuda_rast=False, glctx=None): - import nvdiffrast.torch as dr - from kiui.op import uv_padding - - if glctx is None: - if force_cuda_rast: - glctx = dr.RasterizeCudaContext() - else: - glctx = dr.RasterizeGLContext() - - # render uv maps - h = w = texture_resolution - if mesh.vt is None: - mesh.auto_uv() - - uv = mesh.vt * 2.0 - 1.0 # uvs to range [-1, 1] - uv = torch.cat((uv, torch.zeros_like(uv[..., :1]), torch.ones_like(uv[..., :1])), dim=-1) # [N, 4] - - rast, _ = dr.rasterize(glctx, uv.unsqueeze(0), mesh.ft, (h, w)) # [1, h, w, 4] - xyzs, _ = dr.interpolate(mesh.v.unsqueeze(0), rast, mesh.f) # [1, h, w, 3] - mask, _ = dr.interpolate(torch.ones_like(mesh.v[:, :1]).unsqueeze(0), rast, mesh.f) # [1, h, w, 1] - - # masked query - xyzs = xyzs.view(-1, 3) - mask = (mask > 0).view(-1) - - albedo = torch.zeros(h * w, 3, device=device, dtype=torch.float32) - - if mask.any(): - print(f"[INFO] querying texture...") - - xyzs = xyzs[mask] # [M, 3] - - # batched inference to avoid OOM - batch = [] - head = 0 - while head < xyzs.shape[0]: - tail = min(head + batch_size, xyzs.shape[0]) - batch.append(get_rgb_func(xyzs[head:tail]).float()) - head += batch_size - - albedo[mask] = torch.cat(batch, dim=0) - - albedo = albedo.view(h, w, -1) - mask = mask.view(h, w) - albedo = uv_padding(albedo, mask, padding) - - return albedo - -@torch.no_grad() -def K_nearest_neighbors_func( - points: torch.Tensor, - K: int, - query: Optional[torch.Tensor] = None, - return_dist=False, - skip_first_index=True, -): - if not pytorch3d_capable: - raise ImportError("pytorch3d is not installed, which is required for KNN") - - # query/points: Tensor of shape (N, P1/P2, D) giving a batch of N point clouds, each containing up to P1/P2 points of dimension D - if query is None: - query = points - dist, idx, nn = knn_points(query[None, ...], points[None, ...], K=K, return_nn=True) - - # idx: Tensor of shape (N, P1, K) - # nn: Tensor of shape (N, P1, K, D) - - # take the index 1 since index 0 is the point itself - if skip_first_index: - if not return_dist: - return nn[0, :, 1:, :], idx[0, :, 1:] - else: - return nn[0, :, 1:, :], idx[0, :, 1:], dist[0, :, 1:] - else: - if not return_dist: - return nn[0, :, :, :], idx[0, :, :] - else: - return nn[0, :, :, :], idx[0, :, :], dist[0, :, :] - -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: - mesh.auto_uv() - - # Get Faces on UV - texture_size_minus_one = texture_size - 1 - faces = mesh.ft - verts_uvs = mesh.vt * texture_size_minus_one - verts_uvs_idx = verts_uvs.to(torch.long) - - vmapping = mesh.get_default_vt_to_v_mapping() - # Get vertex colors - interpolate_color = interpolate_color and mesh.vc is not None - texture_map = mesh.albedo - if interpolate_color: - verts_colors = mesh.vc[vmapping] - # Create a blank texture map - texture_map = torch.zeros((texture_size, texture_size, 3), device=verts_colors.device) - - # Get vertex positions - position_map = None - if interpolate_position: - verts_positions = mesh.v[vmapping] - position_map = torch.zeros((texture_size, texture_size, 3), device=verts_colors.device) - - # Create a grid of UV coordinates - grid_x, grid_y = torch.meshgrid(torch.linspace(0, texture_size_minus_one, texture_size, dtype=verts_uvs.dtype, device=verts_colors.device), torch.linspace(0, texture_size_minus_one, texture_size, dtype=verts_uvs.dtype, device=verts_colors.device)) - grid = torch.stack([grid_x, grid_y], dim=-1) - - batch_idx_range = torch.arange(0, batch_size**2, 1, device=verts_colors.device, dtype=torch.long) - - all_face_verts = verts_uvs_idx[faces] - for uv_i in range(0, texture_size_minus_one, batch_size): - for uv_j in range(0, texture_size_minus_one, batch_size): - # Divide faces into squre batchs - faces_mask = (uv_i <= all_face_verts[:, :, 0]) & (all_face_verts[:, :, 0] <= uv_i + batch_size) & (uv_j <= all_face_verts[:, :, 1]) & (all_face_verts[:, :, 1] <= uv_j + batch_size) - faces_mask, _ = torch.max(faces_mask, dim=1) - faces_mask = faces_mask.nonzero().squeeze(-1) - batch_faces = faces[faces_mask] - - uvs = verts_uvs[batch_faces] - if interpolate_color: - colors = verts_colors[batch_faces] - if interpolate_position: - positions = verts_positions[batch_faces] - - v0, v1, v2 = uvs[:, 0], uvs[:, 1], uvs[:, 2] - v0_0, v0_1 = v0[:, 0], v0[:, 1] - v1_0, v1_1 = v1[:, 0], v1[:, 1] - v2_0, v2_1 = v2[:, 0], v2[:, 1] - - sub_grid = grid[uv_i:uv_i + batch_size, uv_j:uv_j + batch_size].reshape(-1, 2).unsqueeze(1) - - # Compute barycentric coordinates for the batch of faces - denom = (v1_1 - v2_1) * (v0_0 - v2_0) + (v2_0 - v1_0) * (v0_1 - v2_1) - a = ((v1_1 - v2_1) * (sub_grid[:, :, 0] - v2_0) + (v2_0 - v1_0) * (sub_grid[:, :, 1] - v2_1)) / denom - b = ((v2_1 - v0_1) * (sub_grid[:, :, 0] - v2_0) + (v0_0 - v2_0) * (sub_grid[:, :, 1] - v2_1)) / denom - c = 1 - a - b - - # Mask for points inside the triangle - mask = (a >= 0) & (b >= 0) & (c >= 0) - mask, mask_idx = torch.max(mask, dim=1) - - # Get texture coordinates - sub_grid = sub_grid.squeeze(1) - uv_coords = sub_grid[mask].to(torch.long) - - # Interpolate colors - a, b, c = a[batch_idx_range, mask_idx][mask], b[batch_idx_range, mask_idx][mask], c[batch_idx_range, mask_idx][mask] - - if interpolate_color: - colors = colors[mask_idx][mask] - interpolated_colors = (a[:, None] * colors[:, 0] + (b[:, None] * colors[:, 1]) + (c[:, None] * colors[:, 2])) - texture_map[uv_coords[:, 1], uv_coords[:, 0]] = interpolated_colors - if interpolate_position: - positions = positions[mask_idx][mask] - interpolated_positions = (a[:, None] * positions[:, 0] + (b[:, None] * positions[:, 1]) + (c[:, None] * positions[:, 2])) - position_map[uv_coords[:, 1], uv_coords[:, 0]] = interpolated_positions - - return texture_map, position_map \ No newline at end of file From b4aa1f8bf33af328d188963c9323dffc897743ed Mon Sep 17 00:00:00 2001 From: DenisKochetov Date: Tue, 2 Sep 2025 23:14:05 +0200 Subject: [PATCH 2/3] add functions for fastmesh --- mesh_processor/__init__.py | 89 +++++++++ mesh_processor/accessors.py | 206 ++++++++++++++++++++ mesh_processor/gltf_ops.py | 42 +++++ mesh_processor/io_gltf.py | 186 +++++++++++++++++++ mesh_processor/mesh_ops.py | 161 ++++++++++++++++ mesh_processor/rigging_ops.py | 340 ++++++++++++++++++++++++++++++++++ mesh_processor/transforms.py | 178 ++++++++++++++++++ mesh_processor/utils.py | 63 +++++++ mesh_processor/validate.py | 48 +++++ 9 files changed, 1313 insertions(+) create mode 100644 mesh_processor/__init__.py create mode 100644 mesh_processor/accessors.py create mode 100644 mesh_processor/gltf_ops.py create mode 100644 mesh_processor/io_gltf.py create mode 100644 mesh_processor/mesh_ops.py create mode 100644 mesh_processor/rigging_ops.py create mode 100644 mesh_processor/transforms.py create mode 100644 mesh_processor/utils.py create mode 100644 mesh_processor/validate.py 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 + From a295aa69592a0d6f24bfbdd3746eca559d49c97e Mon Sep 17 00:00:00 2001 From: DenisKochetov Date: Wed, 3 Sep 2025 02:21:37 +0200 Subject: [PATCH 3/3] integrated fastmesh to hun2.1 --- .../DifferentiableRenderer/__init__.py | 2 +- .../mesh_inpaint_processor/__init__.py | 2 +- mesh_processor/__init__.py | 4 - mesh_processor/export_utils.py | 192 +++ mesh_processor/mesh.py | 1534 +++++++++++++++++ mesh_processor/mesh_utils.py | 741 ++++++++ nodes.py | 74 +- 7 files changed, 2501 insertions(+), 48 deletions(-) create mode 100644 mesh_processor/export_utils.py create mode 100644 mesh_processor/mesh.py create mode 100644 mesh_processor/mesh_utils.py 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/mesh_processor/__init__.py b/mesh_processor/__init__.py index 7250bc4..2ba5108 100644 --- a/mesh_processor/__init__.py +++ b/mesh_processor/__init__.py @@ -35,8 +35,6 @@ from .validate import validate_before_save from .mesh import ( Mesh, FastMesh, - TrimeshLike, - TrimeshLikeLoader, PointCloud, ) @@ -67,8 +65,6 @@ __all__ = [ "validate_before_save", "Mesh", "FastMesh", - "TrimeshLike", - "TrimeshLikeLoader", "PointCloud", ] 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/mesh.py b/mesh_processor/mesh.py new file mode 100644 index 0000000..c1181f3 --- /dev/null +++ b/mesh_processor/mesh.py @@ -0,0 +1,1534 @@ +import os +import cv2 +from typing import NamedTuple, Optional +import torch +from torch import Tensor +import trimesh +import numpy as np + +from kiui.op import safe_normalize, dot +from kiui.typing import * + +from shared_utils.sh_utils import SH2RGB +from shared_utils.image_utils import prepare_torch_img + +# new loader +from .io_gltf import load_gltf_or_glb, get_binary_data +from .mesh_ops import get_all_meshes_triangles +from .accessors import access_data + + +class Mesh: + """ + A torch-native trimesh class, with support for ``ply/obj/glb`` formats. + + Note: + This class only supports one mesh with a single texture image (an albedo texture and a metallic-roughness texture). + """ + 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, # vertex color + albedo: Optional[Tensor] = None, + metallicRoughness: Optional[Tensor] = None, + device: Optional[torch.device] = None, + ): + """Init a mesh directly using all attributes. + + Args: + v (Optional[Tensor]): vertices, float [N, 3]. Defaults to None. + f (Optional[Tensor]): faces, int [M, 3]. Defaults to None. + vn (Optional[Tensor]): vertex normals, float [N, 3]. Defaults to None. + fn (Optional[Tensor]): faces for normals, int [M, 3]. Defaults to None. + vt (Optional[Tensor]): vertex uv coordinates, float [N, 2]. Defaults to None. + ft (Optional[Tensor]): faces for uvs, int [M, 3]. Defaults to None. + vc (Optional[Tensor]): vertex colors, float [N, 3]. Defaults to None. + albedo (Optional[Tensor]): albedo texture, float [H, W, 3], RGB format. Defaults to None. + metallicRoughness (Optional[Tensor]): metallic-roughness texture, float [H, W, 3], metallic(Blue) = metallicRoughness[..., 2], roughness(Green) = metallicRoughness[..., 1]. Defaults to None. + device (Optional[torch.device]): torch device. Defaults to None. + """ + self.device = device + self.v = v + self.vn = vn + self.vt = vt + self.f = f + self.fn = fn + self.ft = ft + # will first see if there is vertex color to use + self.vc = vc + # only support a single albedo image, (H, W, 3) + self.albedo = albedo + # pbr extension, metallic(Blue) = metallicRoughness[..., 2], roughness(Green) = metallicRoughness[..., 1] + # ref: https://registry.khronos.org/glTF/specs/2.0/glTF-2.0.html + 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', use_new_gltf_loader=False, **kwargs): + """load mesh from path. + + Args: + path (str): path to mesh file, supports ply, obj, glb. + resize (bool, optional): auto resize the mesh using ``bound`` into [-bound, bound]^3. Defaults to True. + renormal (bool, optional): re-calc the vertex normals. Defaults to True. + retex (bool, optional): re-calc the uv coordinates, will overwrite the existing uv coordinates. Defaults to False. + clean (bool, optional): perform mesh cleaning at load (e.g., merge close vertices). Defaults to False. + bound (float, optional): bound to resize. Defaults to 0.9. + front_dir (str, optional): front-view direction of the mesh, should be [+-][xyz][ 123]. Defaults to '+z'. + use_new_gltf_loader (bool, optional): use new mesh_processor-based GLTF loader instead of trimesh. Defaults to False. + device (torch.device, optional): torch device. Defaults to None. + + Note: + a ``device`` keyword argument can be provided to specify the torch device. + If it's not provided, we will try to use ``'cuda'`` as the device if it's available. + + Returns: + Mesh: the loaded Mesh object. + """ + # obj supports face uv + if path.endswith(".obj"): + mesh = cls.load_obj(path, **kwargs) + # Select loader for GLB/GLTF + elif use_new_gltf_loader and (path.endswith(".glb") or path.endswith(".gltf")): + mesh = cls.load_gltf(path, **kwargs) + # trimesh only supports vertex uv, but can load more formats + else: + mesh = cls.load_trimesh(path, **kwargs) + + # clean + if clean: + from kiui.mesh_utils import clean_mesh + vertices = mesh.v.detach().cpu().numpy() + triangles = mesh.f.detach().cpu().numpy() + vertices, triangles = clean_mesh(vertices, triangles, remesh=False) + mesh.v = torch.from_numpy(vertices).contiguous().float().to(mesh.device) + mesh.f = torch.from_numpy(triangles).contiguous().int().to(mesh.device) + + print(f"[Mesh loading] v: {mesh.v.shape}, f: {mesh.f.shape}") + # auto-normalize + if resize: + mesh.auto_size(bound=bound) + # auto-fix normal + if renormal or mesh.vn is None: + mesh.auto_normal() + print(f"[Mesh loading] vn: {mesh.vn.shape}, fn: {mesh.fn.shape}") + # auto-fix texcoords + if retex or (mesh.albedo is not None and mesh.vt is None): + mesh.auto_uv(cache_path=path) + print(f"[Mesh loading] vt: {mesh.vt.shape}, ft: {mesh.ft.shape}") + + # rotate front dir to +z + if front_dir != "+z": + # axis switch + if "-z" in front_dir: + T = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, -1]], device=mesh.device, dtype=torch.float32) + elif "+x" in front_dir: + T = torch.tensor([[0, 0, 1], [0, 1, 0], [1, 0, 0]], device=mesh.device, dtype=torch.float32) + elif "-x" in front_dir: + T = torch.tensor([[0, 0, -1], [0, 1, 0], [1, 0, 0]], device=mesh.device, dtype=torch.float32) + elif "+y" in front_dir: + T = torch.tensor([[1, 0, 0], [0, 0, 1], [0, 1, 0]], device=mesh.device, dtype=torch.float32) + elif "-y" in front_dir: + T = torch.tensor([[1, 0, 0], [0, 0, -1], [0, 1, 0]], device=mesh.device, dtype=torch.float32) + else: + T = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], device=mesh.device, dtype=torch.float32) + # rotation (how many 90 degrees) + if '1' in front_dir: + T @= torch.tensor([[0, -1, 0], [1, 0, 0], [0, 0, 1]], device=mesh.device, dtype=torch.float32) + elif '2' in front_dir: + T @= torch.tensor([[1, 0, 0], [0, -1, 0], [0, 0, 1]], device=mesh.device, dtype=torch.float32) + elif '3' in front_dir: + T @= torch.tensor([[0, 1, 0], [-1, 0, 0], [0, 0, 1]], device=mesh.device, dtype=torch.float32) + mesh.v @= T + mesh.vn @= T + + return mesh + + # load from obj file + @classmethod + def load_obj(cls, path, albedo_path=None, device=None): + """load an ``obj`` mesh. + + Args: + path (str): path to mesh. + albedo_path (str, optional): path to the albedo texture image, will overwrite the existing texture path if specified in mtl. Defaults to None. + device (torch.device, optional): torch device. Defaults to None. + + Note: + We will try to read `mtl` path from `obj`, else we assume the file name is the same as `obj` but with `mtl` extension. + The `usemtl` statement is ignored, and we only use the last material path in `mtl` file. + + Returns: + Mesh: the loaded Mesh object. + """ + assert os.path.splitext(path)[-1] == ".obj" + + mesh = cls() + + # device + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + mesh.device = device + + # load obj + with open(path, "r") as f: + lines = f.readlines() + + def parse_f_v(fv): + # pass in a vertex term of a face, return {v, vt, vn} (-1 if not provided) + # supported forms: + # f v1 v2 v3 + # f v1/vt1 v2/vt2 v3/vt3 + # f v1/vt1/vn1 v2/vt2/vn2 v3/vt3/vn3 + # f v1//vn1 v2//vn2 v3//vn3 + xs = [int(x) - 1 if x != "" else -1 for x in fv.split("/")] + xs.extend([-1] * (3 - len(xs))) + return xs[0], xs[1], xs[2] + + vertices, texcoords, normals = [], [], [] + faces, tfaces, nfaces = [], [], [] + mtl_path = None + + for line in lines: + split_line = line.split() + # empty line + if len(split_line) == 0: + continue + prefix = split_line[0].lower() + # mtllib + if prefix == "mtllib": + mtl_path = split_line[1] + # usemtl + elif prefix == "usemtl": + pass # ignored + # v/vn/vt + elif prefix == "v": + vertices.append([float(v) for v in split_line[1:]]) + elif prefix == "vn": + normals.append([float(v) for v in split_line[1:]]) + elif prefix == "vt": + val = [float(v) for v in split_line[1:]] + texcoords.append([val[0], 1.0 - val[1]]) + elif prefix == "f": + vs = split_line[1:] + nv = len(vs) + v0, t0, n0 = parse_f_v(vs[0]) + for i in range(nv - 2): # triangulate (assume vertices are ordered) + v1, t1, n1 = parse_f_v(vs[i + 1]) + v2, t2, n2 = parse_f_v(vs[i + 2]) + faces.append([v0, v1, v2]) + tfaces.append([t0, t1, t2]) + nfaces.append([n0, n1, n2]) + + mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device) + mesh.vt = ( + torch.tensor(texcoords, dtype=torch.float32, device=device) + if len(texcoords) > 0 + else None + ) + mesh.vn = ( + torch.tensor(normals, dtype=torch.float32, device=device) + if len(normals) > 0 + else None + ) + + mesh.f = torch.tensor(faces, dtype=torch.int32, device=device) + mesh.ft = ( + torch.tensor(tfaces, dtype=torch.int32, device=device) + if len(texcoords) > 0 + else None + ) + mesh.fn = ( + torch.tensor(nfaces, dtype=torch.int32, device=device) + if len(normals) > 0 + else None + ) + + # see if there is vertex color + use_vertex_color = False + if mesh.v.shape[1] == 6: + use_vertex_color = True + mesh.vc = mesh.v[:, 3:] + mesh.v = mesh.v[:, :3] + print(f"[load_obj] use vertex color: {mesh.vc.shape}") + + # try to load texture image + if not use_vertex_color: + # try to retrieve mtl file + mtl_path_candidates = [] + if mtl_path is not None: + mtl_path_candidates.append(mtl_path) + mtl_path_candidates.append(os.path.join(os.path.dirname(path), mtl_path)) + mtl_path_candidates.append(path.replace(".obj", ".mtl")) + + mtl_path = None + for candidate in mtl_path_candidates: + if os.path.exists(candidate): + mtl_path = candidate + break + + # if albedo_path is not provided, try retrieve it from mtl + metallic_path = None + roughness_path = None + if mtl_path is not None and albedo_path is None: + with open(mtl_path, "r") as f: + lines = f.readlines() + + for line in lines: + split_line = line.split() + # empty line + if len(split_line) == 0: + continue + prefix = split_line[0] + + if "map_Kd" in prefix: + # assume relative path! + albedo_path = os.path.join(os.path.dirname(path), split_line[1]) + print(f"[load_obj] use texture from: {albedo_path}") + elif "map_Pm" in prefix: + metallic_path = os.path.join(os.path.dirname(path), split_line[1]) + elif "map_Pr" in prefix: + roughness_path = os.path.join(os.path.dirname(path), split_line[1]) + + # still not found albedo_path, or the path doesn't exist + if albedo_path is None or not os.path.exists(albedo_path): + # init an empty texture + print(f"[load_obj] init empty albedo!") + # albedo = np.random.rand(1024, 1024, 3).astype(np.float32) + albedo = np.ones((1024, 1024, 3), dtype=np.float32) * np.array([0.5, 0.5, 0.5]) # default color + else: + albedo = cv2.imread(albedo_path, cv2.IMREAD_UNCHANGED) + albedo = cv2.cvtColor(albedo, cv2.COLOR_BGR2RGB) + albedo = albedo.astype(np.float32) / 255 + print(f"[load_obj] load texture: {albedo.shape}") + + mesh.albedo = torch.tensor(albedo, dtype=torch.float32, device=device) + + # try to load metallic and roughness + if metallic_path is not None and roughness_path is not None: + print(f"[load_obj] load metallicRoughness from: {metallic_path}, {roughness_path}") + metallic = cv2.imread(metallic_path, cv2.IMREAD_UNCHANGED) + metallic = metallic.astype(np.float32) / 255 + roughness = cv2.imread(roughness_path, cv2.IMREAD_UNCHANGED) + roughness = roughness.astype(np.float32) / 255 + metallicRoughness = np.stack([np.zeros_like(metallic), roughness, metallic], axis=-1) + + mesh.metallicRoughness = torch.tensor(metallicRoughness, dtype=torch.float32, device=device).contiguous() + + return mesh + + @classmethod + def load_trimesh(cls, path=None, given_mesh=None, device=None): + """load a mesh using ``trimesh.load()``. + + Can load various formats like ``glb`` and serves as a fallback. + + Note: + We will try to merge all meshes if the glb contains more than one, + but **this may cause the texture to lose**, since we only support one texture image! + + Args: + path (str): path to the mesh file. + given_mesh (trimesh): trimesh instance + device (torch.device, optional): torch device. Defaults to None. + + Returns: + Mesh: the loaded Mesh object. + """ + mesh = cls() + + # device + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + mesh.device = device + + # use trimesh to load ply/glb, assume only has one single RootMesh... + if given_mesh is not None: + _mesh = given_mesh + elif path is not None: + _data = trimesh.load(path) + if isinstance(_data, trimesh.Scene): + if len(_data.geometry) == 1: + _mesh = list(_data.geometry.values())[0] + else: + print(f"[load_trimesh] concatenating {len(_data.geometry)} meshes.") + _concat = [] + # loop the scene graph and apply transform to each mesh + scene_graph = _data.graph.to_flattened() # dict {name: {transform: 4x4 mat, geometry: str}} + for k, v in scene_graph.items(): + name = v['geometry'] + if name in _data.geometry and isinstance(_data.geometry[name], trimesh.Trimesh): + transform = v['transform'] + _concat.append(_data.geometry[name].apply_transform(transform)) + _mesh = trimesh.util.concatenate(_concat) + else: + _mesh = _data + else: + print(f"[load_trimesh] failed to load mesh, either path or given_mesh must be given") + return None + + should_create_empty_albedo = False + if _mesh.visual.kind == 'vertex': + vertex_colors = _mesh.visual.vertex_colors + vertex_colors = np.array(vertex_colors[..., :3]).astype(np.float32) / 255 + mesh.vc = torch.tensor(vertex_colors, dtype=torch.float32, device=device) + print(f"[load_trimesh] use vertex color: {mesh.vc.shape}") + elif _mesh.visual.kind == 'texture': + _material = _mesh.visual.material + if isinstance(_material, trimesh.visual.material.PBRMaterial): + texture = np.array(_material.baseColorTexture).astype(np.float32) / 255 + # load metallicRoughness if present + if _material.metallicRoughnessTexture is not None: + metallicRoughness = np.array(_material.metallicRoughnessTexture).astype(np.float32) / 255 + mesh.metallicRoughness = torch.tensor(metallicRoughness, dtype=torch.float32, device=device).contiguous() + elif isinstance(_material, trimesh.visual.material.SimpleMaterial): + texture = np.array(_material.to_pbr().baseColorTexture).astype(np.float32) / 255 + else: + raise NotImplementedError(f"material type {type(_material)} not supported!") + + if texture.ndim ==3: + mesh.albedo = torch.tensor(texture[..., :3], dtype=torch.float32, device=device).contiguous() + print(f"[load_trimesh] loaded albedo texture: {texture.shape}") + else: + should_create_empty_albedo = True + else: + should_create_empty_albedo = True + + if should_create_empty_albedo: + mesh.set_new_albedo(1024, 1024) + print(f"[load_trimesh] failed to load albedo texture, create empty abbedo texture with shape: {mesh.albedo.shape}") + + vertices = _mesh.vertices + + try: + texcoords = _mesh.visual.uv + texcoords[:, 1] = 1 - texcoords[:, 1] + except Exception as e: + texcoords = None + + try: + normals = _mesh.vertex_normals + except Exception as e: + normals = None + + # trimesh only support vertex uv... + faces = tfaces = nfaces = _mesh.faces + + mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device) + mesh.vt = ( + torch.tensor(texcoords, dtype=torch.float32, device=device) + if texcoords is not None + else None + ) + mesh.vn = ( + torch.tensor(normals, dtype=torch.float32, device=device) + if normals is not None + else None + ) + + mesh.f = torch.tensor(faces, dtype=torch.int32, device=device) + mesh.ft = ( + torch.tensor(tfaces, dtype=torch.int32, device=device) + if texcoords is not None + else None + ) + mesh.fn = ( + torch.tensor(nfaces, dtype=torch.int32, device=device) + if normals is not None + else None + ) + + return mesh + + @classmethod + def from_trimesh(cls, trimesh_obj, device=None, use_new_converter=False): + """Create Mesh object from trimesh object. + + Args: + trimesh_obj: trimesh.Trimesh or trimesh.Scene + device: torch device + use_new_converter: use new converter (experimental) + + Returns: + Mesh: Mesh object + """ + if use_new_converter: + # New way - try to extract data directly without full dependency on trimesh + return cls._from_trimesh_new(trimesh_obj, device) + else: + # Old way - use load_trimesh + return cls.load_trimesh(given_mesh=trimesh_obj, device=device) + + + @classmethod + def load_gltf(cls, path=None, device=None): + """Load mesh with using mesh_processor (new implementation). + + Alternative to load_trimesh() for GLB/GLTF files without dependency on trimesh. + + Args: + path (str): path to GLB/GLTF file. + device (torch.device, optional): torch device. Defaults to None. + + Returns: + Mesh: loaded Mesh object. + """ + mesh = cls() + + # device + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + mesh.device = device + + if path is None: + print(f"[load_gltf] path is required") + return None + + try: + # Load GLTF document + doc = load_gltf_or_glb(path) + print(f"[load_gltf] loaded GLTF document with {len(doc.meshes())} meshes") + + # Extract geometry of all meshes + vertices, faces, mesh_groups = get_all_meshes_triangles(doc, transform_to_global=True) + + if len(vertices) == 0 or len(faces) == 0: + print(f"[load_gltf] no valid geometry found") + return None + + print(f"[load_gltf] extracted {len(vertices)} vertices, {len(faces)} faces") + + # Convert to torch tensors + mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device) + mesh.f = torch.tensor(faces, dtype=torch.int32, device=device) + + # Try to extract UV coordinates and normals from the first mesh + texcoords = None + normals = None + vertex_colors = None + + if doc.meshes(): + first_mesh = doc.meshes()[0] + for primitive in first_mesh.get("primitives", []): + attributes = primitive.get("attributes", {}) + + # UV coordinates + if "TEXCOORD_0" in attributes and texcoords is None: + try: + texcoord_accessor_idx = int(attributes["TEXCOORD_0"]) + texcoords = access_data(doc, texcoord_accessor_idx).astype(np.float32) + # Flip V coordinate (as in trimesh) + texcoords[:, 1] = 1.0 - texcoords[:, 1] + print(f"[load_gltf] extracted UV coordinates: {texcoords.shape}") + except Exception as e: + print(f"[load_gltf] failed to extract UV coordinates: {e}") + + # Normals + if "NORMAL" in attributes and normals is None: + try: + normal_accessor_idx = int(attributes["NORMAL"]) + normals = access_data(doc, normal_accessor_idx).astype(np.float32) + print(f"[load_gltf] extracted normals: {normals.shape}") + except Exception as e: + print(f"[load_gltf] failed to extract normals: {e}") + + # Vertex colors + if "COLOR_0" in attributes and vertex_colors is None: + try: + color_accessor_idx = int(attributes["COLOR_0"]) + vertex_colors = access_data(doc, color_accessor_idx).astype(np.float32) + if vertex_colors.shape[1] > 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): + """get the axis-aligned bounding box of the mesh. + + Returns: + Tuple[torch.Tensor]: the min xyz and max xyz of the mesh. + """ + return torch.min(self.v, dim=0).values, torch.max(self.v, dim=0).values + + # unit size + @torch.no_grad() + def auto_size(self, bound=0.9): + """auto resize the mesh. + + Args: + bound (float, optional): resizing into ``[-bound, bound]^3``. Defaults to 0.9. + """ + 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): + """auto calculate the vertex 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) + + # Splat face normals to vertices + 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) + + # Normalize, replace zero (degenerated) normals with some default value + 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): + """auto calculate the uv coordinates. + + Args: + cache_path (str, optional): path to save/load the uv cache as a npz file, this can avoid calculating uv every time when loading the same mesh, which is time-consuming. Defaults to None. + vmap (bool, optional): remap vertices based on uv coordinates, so each v correspond to a unique vt (necessary for formats like gltf). + Usually this will duplicate the vertices on the edge of uv atlas. Defaults to True. + """ + # try to load cache + if cache_path is not None: + cache_path = os.path.splitext(cache_path)[0] + "_uv.npz" + if cache_path is not None and os.path.exists(cache_path): + data = np.load(cache_path) + vt_np, ft_np, vmapping = data["vt"], data["ft"], data["vmapping"] + else: + import xatlas + + v_np = self.v.detach().cpu().numpy() + f_np = self.f.detach().int().cpu().numpy() + atlas = xatlas.Atlas() + atlas.add_mesh(v_np, f_np) + chart_options = xatlas.ChartOptions() + # chart_options.max_iterations = 4 + atlas.generate(chart_options=chart_options) + vmapping, ft_np, vt_np = atlas[0] # [N], [M, 3], [N, 2] + + # save to cache + if cache_path is not None: + np.savez(cache_path, vt=vt_np, ft=ft_np, vmapping=vmapping) + + vt = torch.from_numpy(vt_np.astype(np.float32)).to(self.device) + ft = torch.from_numpy(ft_np.astype(np.int32)).to(self.device) + self.vt = vt + self.ft = ft + + if vmap: + # remap v/f to vt/ft, so each v correspond to a unique vt. (necessary for gltf) + vmapping = torch.from_numpy(vmapping.astype(np.int64)).long().to(self.device) + self.align_v_to_vt(vmapping) + + def align_v_to_vt(self, vmapping=None): + """ remap v/f and vn/fn to vt/ft. + + Args: + vmapping (np.ndarray, optional): the mapping relationship from f to ft. Defaults to None. + """ + if vmapping is None: + vt2v_mapping = self.get_default_vt_to_v_mapping() + else: + vt2v_mapping = vmapping + + self.v = self.v[vt2v_mapping] + if self.vc is not None: + self.vc = self.vc[vt2v_mapping] + self.f = self.ft + # assume fn == f + if self.vn is not None: + if vmapping is None: + vt2vn_mapping = self.get_default_vt_to_vn_mapping() + else: + vt2vn_mapping = vmapping + self.vn = self.vn[vt2vn_mapping] + self.fn = self.ft + + def get_default_vt_to_v_mapping(self): + """map from ft to f + + Returns: + vmapping: tensor array + """ + ft = self.ft.view(-1).long() + f = self.f.view(-1).long() + vmapping = torch.zeros(self.vt.shape[0], dtype=torch.long, device=self.device) + vmapping[ft] = f # scatter, randomly choose one if index is not unique + return vmapping + + def get_default_vt_to_vn_mapping(self): + """map from ft to f + + Returns: + vmapping: tensor array + """ + ft = self.ft.view(-1).long() + fn = self.fn.view(-1).long() + vmapping = torch.zeros(self.vt.shape[0], dtype=torch.long, device=self.device) + vmapping[ft] = fn # scatter, randomly choose one if index is not unique + return vmapping + + def to(self, device): + """move all tensor attributes to device. + + Args: + device (torch.device): target device. + + Returns: + Mesh: self. + """ + 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): + """write the mesh to a path. + + Args: + path (str): path to write, supports ply, obj and glb. + """ + if path.endswith(".ply"): + self.write_ply(path) + elif path.endswith(".obj"): + self.write_obj(path) + elif path.endswith(".glb") or path.endswith(".gltf"): + self.write_glb(path) + else: + raise NotImplementedError(f"format {path} not supported!") + + def write_ply(self, path): + """write the mesh in ply format. Only for geometry! + + Args: + path (str): path to write. + """ + + if self.albedo is not None: + print(f'[WARN] ply format does not support exporting texture, will ignore!') + + v_np = self.v.detach().cpu().numpy() + f_np = self.f.detach().cpu().numpy() + + # 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): + """write the mesh in glb/gltf format. + This will create a scene with a single mesh. + + Args: + path (str): path to write. + """ + + # assert self.v.shape[0] == self.vn.shape[0] and self.v.shape[0] == self.vt.shape[0] + if self.vt is not None and self.v.shape[0] != self.vt.shape[0]: + self.align_v_to_vt() + + import pygltflib + + f_np = self.f.detach().cpu().numpy().astype(np.uint32) + f_np_blob = f_np.flatten().tobytes() + + v_np = self.v.detach().cpu().numpy().astype(np.float32) + v_np_blob = v_np.tobytes() + + blob = f_np_blob + v_np_blob + byteOffset = len(blob) + + # base mesh + gltf = pygltflib.GLTF2( + scene=0, + scenes=[pygltflib.Scene(nodes=[0])], + nodes=[pygltflib.Node(mesh=0)], + meshes=[pygltflib.Mesh(primitives=[pygltflib.Primitive( + # indices to accessors (0 is triangles) + attributes=pygltflib.Attributes( + POSITION=1, + ), + indices=0, + )])], + buffers=[ + pygltflib.Buffer(byteLength=len(f_np_blob) + len(v_np_blob)) + ], + # buffer view (based on dtype) + bufferViews=[ + # triangles; as flatten (element) array + pygltflib.BufferView( + buffer=0, + byteLength=len(f_np_blob), + target=pygltflib.ELEMENT_ARRAY_BUFFER, # GL_ELEMENT_ARRAY_BUFFER (34963) + ), + # positions; as vec3 array + pygltflib.BufferView( + buffer=0, + byteOffset=len(f_np_blob), + byteLength=len(v_np_blob), + byteStride=12, # vec3 + target=pygltflib.ARRAY_BUFFER, # GL_ARRAY_BUFFER (34962) + ), + ], + accessors=[ + # 0 = triangles + pygltflib.Accessor( + bufferView=0, + componentType=pygltflib.UNSIGNED_INT, # GL_UNSIGNED_INT (5125) + count=f_np.size, + type=pygltflib.SCALAR, + max=[int(f_np.max())], + min=[int(f_np.min())], + ), + # 1 = positions + pygltflib.Accessor( + bufferView=1, + componentType=pygltflib.FLOAT, # GL_FLOAT (5126) + count=len(v_np), + type=pygltflib.VEC3, + max=v_np.max(axis=0).tolist(), + min=v_np.min(axis=0).tolist(), + ), + ], + ) + + # append texture info + if self.vt is not None and self.albedo is not None: + + vt_np = self.vt.detach().cpu().numpy().astype(np.float32) + vt_np_blob = vt_np.tobytes() + + albedo = self.albedo.detach().cpu().numpy() + albedo = (albedo * 255).astype(np.uint8) + albedo = cv2.cvtColor(albedo, cv2.COLOR_RGB2BGR) + albedo_blob = cv2.imencode('.png', albedo)[1].tobytes() + + # update primitive + gltf.meshes[0].primitives[0].attributes.TEXCOORD_0 = 2 + gltf.meshes[0].primitives[0].material = 0 + + # update materials + gltf.materials.append(pygltflib.Material( + pbrMetallicRoughness=pygltflib.PbrMetallicRoughness( + baseColorTexture=pygltflib.TextureInfo(index=0, texCoord=0), + metallicFactor=0.0, + roughnessFactor=1.0, + ), + alphaMode=pygltflib.OPAQUE, + alphaCutoff=None, + doubleSided=True, + )) + + gltf.textures.append(pygltflib.Texture(sampler=0, source=0)) + gltf.samplers.append(pygltflib.Sampler(magFilter=pygltflib.LINEAR, minFilter=pygltflib.LINEAR_MIPMAP_LINEAR, wrapS=pygltflib.REPEAT, wrapT=pygltflib.REPEAT)) + gltf.images.append(pygltflib.Image(bufferView=3, mimeType="image/png")) + + # update buffers + gltf.bufferViews.append( + # index = 2, texcoords; as vec2 array + pygltflib.BufferView( + buffer=0, + byteOffset=byteOffset, + byteLength=len(vt_np_blob), + byteStride=8, # vec2 + target=pygltflib.ARRAY_BUFFER, + ) + ) + + gltf.accessors.append( + # 2 = texcoords + pygltflib.Accessor( + bufferView=2, + componentType=pygltflib.FLOAT, + count=len(vt_np), + type=pygltflib.VEC2, + max=vt_np.max(axis=0).tolist(), + min=vt_np.min(axis=0).tolist(), + ) + ) + + blob += vt_np_blob + byteOffset += len(vt_np_blob) + + gltf.bufferViews.append( + # index = 3, albedo texture; as none target + pygltflib.BufferView( + buffer=0, + byteOffset=byteOffset, + byteLength=len(albedo_blob), + ) + ) + + blob += albedo_blob + byteOffset += len(albedo_blob) + + gltf.buffers[0].byteLength = byteOffset + + # append metllic roughness + if self.metallicRoughness is not None: + metallicRoughness = self.metallicRoughness.detach().cpu().numpy() + metallicRoughness = (metallicRoughness * 255).astype(np.uint8) + metallicRoughness = cv2.cvtColor(metallicRoughness, cv2.COLOR_RGB2BGR) + metallicRoughness_blob = cv2.imencode('.png', metallicRoughness)[1].tobytes() + + # update texture definition + gltf.materials[0].pbrMetallicRoughness.metallicFactor = 1.0 + gltf.materials[0].pbrMetallicRoughness.roughnessFactor = 1.0 + gltf.materials[0].pbrMetallicRoughness.metallicRoughnessTexture = pygltflib.TextureInfo(index=1, texCoord=0) + + gltf.textures.append(pygltflib.Texture(sampler=1, source=1)) + gltf.samplers.append(pygltflib.Sampler(magFilter=pygltflib.LINEAR, minFilter=pygltflib.LINEAR_MIPMAP_LINEAR, wrapS=pygltflib.REPEAT, wrapT=pygltflib.REPEAT)) + gltf.images.append(pygltflib.Image(bufferView=4, mimeType="image/png")) + + # update buffers + gltf.bufferViews.append( + # index = 4, metallicRoughness texture; as none target + pygltflib.BufferView( + buffer=0, + byteOffset=byteOffset, + byteLength=len(metallicRoughness_blob), + ) + ) + + blob += metallicRoughness_blob + byteOffset += len(metallicRoughness_blob) + + gltf.buffers[0].byteLength = byteOffset + + + # set actual data + gltf.set_binary_blob(blob) + + # glb = b"".join(gltf.save_to_bytes()) + gltf.save(path) + + + def write_obj(self, path): + """write the mesh in obj format. Will also write the texture and mtl files. + + Args: + path (str): path to write. + """ + + mtl_path = path.replace(".obj", ".mtl") + albedo_path = path.replace(".obj", "_albedo.png") + metallic_path = path.replace(".obj", "_metallic.png") + roughness_path = path.replace(".obj", "_roughness.png") + + v_np = self.v.detach().cpu().numpy() + vt_np = self.vt.detach().cpu().numpy() if self.vt is not None else None + vn_np = self.vn.detach().cpu().numpy() if self.vn is not None else None + vc_np = self.vc.detach().cpu().numpy() if self.vc is not None else None + f_np = self.f.detach().cpu().numpy() + ft_np = self.ft.detach().cpu().numpy() if self.ft is not None else None + fn_np = self.fn.detach().cpu().numpy() if self.fn is not None else None + + with open(path, "w") as fp: + fp.write(f"mtllib {os.path.basename(mtl_path)} \n") + + if vc_np is None: + for v in v_np: + fp.write(f"v {v[0]} {v[1]} {v[2]} \n") + else: + for v_i in range(v_np.shape[0]): + v = v_np[v_i] + vc = vc_np[v_i] + fp.write(f"v {v[0]} {v[1]} {v[2]} {vc[0]} {vc[1]} {vc[2]}\n") + + if vt_np is not None: + for v in vt_np: + fp.write(f"vt {v[0]} {1 - v[1]} \n") + + if vn_np is not None: + for v in vn_np: + fp.write(f"vn {v[0]} {v[1]} {v[2]} \n") + + fp.write(f"usemtl defaultMat \n") + for i in range(len(f_np)): + fp.write( + f'f {f_np[i, 0] + 1}/{ft_np[i, 0] + 1 if ft_np is not None else ""}/{fn_np[i, 0] + 1 if fn_np is not None else ""} \ + {f_np[i, 1] + 1}/{ft_np[i, 1] + 1 if ft_np is not None else ""}/{fn_np[i, 1] + 1 if fn_np is not None else ""} \ + {f_np[i, 2] + 1}/{ft_np[i, 2] + 1 if ft_np is not None else ""}/{fn_np[i, 2] + 1 if fn_np is not None else ""} \n' + ) + + with open(mtl_path, "w") as fp: + fp.write(f"newmtl defaultMat \n") + fp.write(f"Ka 1 1 1 \n") + fp.write(f"Kd 1 1 1 \n") + fp.write(f"Ks 1 1 1 \n") + #fp.write(f"Tr 1 \n") # will cause the mesh materials in three.js completely transparent + fp.write(f"illum 1 \n") + fp.write(f"Ns 10 \n") + if self.albedo is not None: + fp.write(f"map_Kd {os.path.basename(albedo_path)} \n") + if self.metallicRoughness is not None: + # ref: https://en.wikipedia.org/wiki/Wavefront_.obj_file#Physically-based_Rendering + fp.write(f"map_Pm {os.path.basename(metallic_path)} \n") + fp.write(f"map_Pr {os.path.basename(roughness_path)} \n") + + if self.albedo is not None: + albedo = self.albedo.detach().cpu().numpy() + albedo = (albedo * 255).astype(np.uint8) + cv2.imwrite(albedo_path, cv2.cvtColor(albedo, cv2.COLOR_RGB2BGR)) + + if self.metallicRoughness is not None: + metallicRoughness = self.metallicRoughness.detach().cpu().numpy() + metallicRoughness = (metallicRoughness * 255).astype(np.uint8) + cv2.imwrite(metallic_path, metallicRoughness[..., 2]) + cv2.imwrite(roughness_path, metallicRoughness[..., 1]) + + def convert_to_pointcloud(self): + xyz = self.v.detach().cpu().numpy() + num_pts = self.v.shape[0] + 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 + + +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 + normals: np.array \ No newline at end of file diff --git a/mesh_processor/mesh_utils.py b/mesh_processor/mesh_utils.py new file mode 100644 index 0000000..24857c1 --- /dev/null +++ b/mesh_processor/mesh_utils.py @@ -0,0 +1,741 @@ +import torch +import numpy as np +from kornia.geometry.conversions import ( + quaternion_to_axis_angle, + axis_angle_to_quaternion, +) +import pymeshlab as pml +from plyfile import PlyData, PlyElement + +from typing import Optional + +pytorch3d_capable = True +try: + import pytorch3d + from pytorch3d.ops import knn_points +except ImportError: + pytorch3d_capable = False + +from .mesh import PointCloud +from shared_utils.sh_utils import SH2RGB, RGB2SH + +def _base_face_areas(face_vertices_0, face_vertices_1, face_vertices_2): + """Base function to compute the face areas.""" + x1, x2, x3 = torch.split(face_vertices_0 - face_vertices_1, 1, dim=-1) + y1, y2, y3 = torch.split(face_vertices_1 - face_vertices_2, 1, dim=-1) + + a = (x2 * y3 - x3 * y2) ** 2 + b = (x3 * y1 - x1 * y3) ** 2 + c = (x1 * y2 - x2 * y1) ** 2 + areas = torch.sqrt(a + b + c) * 0.5 + + return areas + + +def _base_sample_points_selected_faces(face_vertices, face_features=None): + """Base function to sample points over selected faces, sample one point per face. + The coordinates of the face vertices are interpolated to generate new samples. + + Args: + face_vertices (tuple of torch.Tensor): + Coordinates of vertices, corresponding to selected faces to sample from. + A tuple of 3 entries corresponding to each of the face vertices. + Each entry is a torch.Tensor of shape :math:`(\\text{batch_size}, \\text{num_samples}, 3)`. + face_features (tuple of torch.Tensor, Optional): + Features of face vertices, corresponding to selected faces to sample from. + A tuple of 3 entries corresponding to each of the face vertices. + Each entry is a torch.Tensor of shape + :math:`(\\text{batch_size}, \\text{num_samples}, \\text{feature_dim})`. + + Returns: + (torch.Tensor, torch.Tensor): + Sampled point coordinates of shape :math:`(\\text{batch_size}, \\text{num_samples}, 3)`. + Sampled points interpolated features of shape + :math:`(\\text{batch_size}, \\text{num_samples}, \\text{feature_dim})`. + If `face_vertices_features` arg is not specified, the returned interpolated features are None. + """ + + face_vertices0, face_vertices1, face_vertices2 = face_vertices + + sampling_shape = tuple(int(d) for d in face_vertices0.shape[:-1]) + (1,) + # u is proximity to middle point between v1 and v2 against v0. + # v is proximity to v2 against v1. + # + # The probability density for u should be f_U(u) = 2u. + # However, torch.rand use a uniform (f_X(x) = x) distribution, + # so using torch.sqrt we make a change of variable to have the desired density + # f_Y(y) = f_X(y ^ 2) * |d(y ^ 2) / dy| = 2y + u = torch.sqrt(torch.rand(sampling_shape, + device=face_vertices0.device, + dtype=face_vertices0.dtype)) + + v = torch.rand(sampling_shape, + device=face_vertices0.device, + dtype=face_vertices0.dtype) + w0 = 1 - u + w1 = u * (1 - v) + w2 = u * v + + points = w0 * face_vertices0 + w1 * face_vertices1 + w2 * face_vertices2 + + features = None + if face_features is not None: + face_features0, face_features1, face_features2 = face_features + features = w0 * face_features0 + w1 * face_features1 + \ + w2 * face_features2 + + return points, features + +# Modified from https://github.com/NVIDIAGameWorks/kaolin/blob/master/kaolin/ops/mesh/trianglemesh.py#L159 +def sample_points(vertices, faces, num_samples, areas=None, face_features=None): + r"""Uniformly sample points over the surface of triangle meshes. + + First, face on which the point is sampled is randomly selected, + with the probability of selection being proportional to the area of the face. + then the coordinate on the face is uniformly sampled. + + If ``face_features`` is defined for the mesh faces, + the sampled points will be returned with interpolated features as well, + otherwise, no feature interpolation will occur. + + Args: + vertices (torch.Tensor): + The vertices of the meshes, of shape + :math:`(\text{batch_size}, \text{num_vertices}, 3)`. + faces (torch.LongTensor): + The faces of the mesh, of shape :math:`(\text{num_faces}, 3)`. + num_samples (int): + The number of point sampled per mesh. + Also the number of faces sampled per mesh, and then sample a single point per face. + areas (torch.Tensor, optional): + The areas of each face, of shape :math:`(\text{batch_size}, \text{num_faces})`, + can be preprocessed, for fast on-the-fly sampling, + will be computed if None (default). + face_features (torch.Tensor, optional): + Per-vertex-per-face features, matching ``faces`` order, + of shape :math:`(\text{batch_size}, \text{num_faces}, 3, \text{feature_dim})`. + For example: + + 1. Texture uv coordinates would be of shape + :math:`(\text{batch_size}, \text{num_faces}, 3, 2)`. + 2. RGB color values would be of shape + :math:`(\text{batch_size}, \text{num_faces}, 3, 3)`. + + When specified, it is used to interpolate the features for new sampled points. + + See also: + :func:`~kaolin.ops.mesh.index_vertices_by_faces` for conversion of features defined per vertex + and need to be converted to per-vertex-per-face shape of :math:`(\text{num_faces}, 3)`. + + Returns: + (torch.Tensor, torch.LongTensor, (optional) torch.Tensor): + the pointclouds of shape :math:`(\text{batch_size}, \text{num_samples}, 3)`, + and the indexes of the faces selected, + of shape :math:`(\text{batch_size}, \text{num_samples})`. + + If ``face_features`` arg is specified, then the interpolated features of sampled points of shape + :math:`(\text{batch_size}, \text{num_samples}, \text{feature_dim})` are also returned. + """ + if faces.shape[-1] != 3: + raise NotImplementedError("sample_points is only implemented for triangle meshes") + faces_0, faces_1, faces_2 = torch.split(faces, 1, dim=1) # (num_faces, 3) -> tuple of (num_faces,) + face_v_0 = torch.index_select(vertices, 1, faces_0.reshape(-1)) # (batch_size, num_faces, 3) + face_v_1 = torch.index_select(vertices, 1, faces_1.reshape(-1)) # (batch_size, num_faces, 3) + face_v_2 = torch.index_select(vertices, 1, faces_2.reshape(-1)) # (batch_size, num_faces, 3) + + if areas is None: + areas = _base_face_areas(face_v_0, face_v_1, face_v_2).squeeze(-1) + face_dist = torch.distributions.Categorical(areas) + face_choices = face_dist.sample([num_samples]).transpose(0, 1) + _face_choices = face_choices.unsqueeze(-1).repeat(1, 1, 3) + v0 = torch.gather(face_v_0, 1, _face_choices) # (batch_size, num_samples, 3) + v1 = torch.gather(face_v_1, 1, _face_choices) # (batch_size, num_samples, 3) + v2 = torch.gather(face_v_2, 1, _face_choices) # (batch_size, num_samples, 3) + face_vertices_choices = (v0, v1, v2) + + # UV coordinates are available, make sure to calculate them for sampled points as well + face_features_choices = None + if face_features is not None: + feat_dim = face_features.shape[-1] + # (num_faces, 3) -> tuple of (num_faces,) + _face_choices = face_choices[..., None, None].repeat(1, 1, 3, feat_dim) + face_features_choices = torch.gather(face_features, 1, _face_choices) + face_features_choices = tuple( + tmp_feat.squeeze(2) for tmp_feat in torch.split(face_features_choices, 1, dim=2)) + + points, point_features = _base_sample_points_selected_faces( + face_vertices_choices, face_features_choices) + + if point_features is not None: + return points, face_choices, point_features + else: + return points, face_choices + +def poisson_mesh_reconstruction(points, normals=None): + # points/normals: [N, 3] np.ndarray + + import open3d as o3d + + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(points) + + # outlier removal + pcd, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=10) + + # normals + if normals is None: + pcd.estimate_normals() + else: + pcd.normals = o3d.utility.Vector3dVector(normals[ind]) + + # visualize + o3d.visualization.draw_geometries([pcd], point_show_normal=False) + + mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson( + pcd, depth=9 + ) + vertices_to_remove = densities < np.quantile(densities, 0.1) + mesh.remove_vertices_by_mask(vertices_to_remove) + + # visualize + o3d.visualization.draw_geometries([mesh]) + + vertices = np.asarray(mesh.vertices) + triangles = np.asarray(mesh.triangles) + + print( + f"[INFO] poisson mesh reconstruction: {points.shape} --> {vertices.shape} / {triangles.shape}" + ) + + return vertices, triangles + + +def decimate_mesh( + verts, faces, target=5e4, backend="pymeshlab", remesh=False, optimalplacement=True, verbose=True +): + """ perform mesh decimation. + + Args: + verts (np.ndarray): mesh vertices, float [N, 3] + faces (np.ndarray): mesh faces, int [M, 3] + target (int): targeted number of faces + backend (str, optional): algorithm backend, can be "pymeshlab" or "pyfqmr". Defaults to "pymeshlab". + remesh (bool, optional): whether to remesh after decimation. Defaults to False. + optimalplacement (bool, optional): For flat mesh, use False to prevent spikes. Defaults to True. + verbose (bool, optional): whether to print the decimation process. Defaults to True. + + Returns: + Tuple[np.ndarray]: vertices and faces after decimation. + """ + + _ori_vert_shape = verts.shape + _ori_face_shape = faces.shape + + if backend == "pyfqmr": + import pyfqmr + + solver = pyfqmr.Simplify() + solver.setMesh(verts, faces) + solver.simplify_mesh(target_count=target, preserve_border=False, verbose=False) + verts, faces, normals = solver.getMesh() + else: + m = pml.Mesh(verts, faces) + ms = pml.MeshSet() + ms.add_mesh(m, "mesh") # will copy! + + ms.meshing_merge_close_vertices() + + # filters + # ms.meshing_decimation_clustering(threshold=pml.PercentageValue(1)) + ms.meshing_decimation_quadric_edge_collapse( + targetfacenum=int(target), optimalplacement=optimalplacement + ) + + if remesh: + # ms.apply_coord_taubin_smoothing() + ms.meshing_isotropic_explicit_remeshing( + iterations=3, targetlen=pml.PercentageValue(1) + ) + + # extract mesh + m = ms.current_mesh() + m.compact() + verts = m.vertex_matrix() + faces = m.face_matrix() + + if verbose: + print(f"[INFO] mesh decimation: {_ori_vert_shape} --> {verts.shape}, {_ori_face_shape} --> {faces.shape}") + + return verts, faces + + +def clean_mesh( + verts, + faces, + v_pct=1, + min_f=64, + min_d=20, + repair=True, + remesh=True, + remesh_size=0.01, +): + # verts: [N, 3] + # faces: [N, 3] + + _ori_vert_shape = verts.shape + _ori_face_shape = faces.shape + + m = pml.Mesh(verts, faces) + ms = pml.MeshSet() + ms.add_mesh(m, "mesh") # will copy! + + # filters + ms.meshing_remove_unreferenced_vertices() # verts not refed by any faces + + if v_pct > 0: + ms.meshing_merge_close_vertices( + threshold=pml.Percentage(v_pct) + ) # 1/10000 of bounding box diagonal + + ms.meshing_remove_duplicate_faces() # faces defined by the same verts + ms.meshing_remove_null_faces() # faces with area == 0 + + if min_d > 0: + ms.meshing_remove_connected_component_by_diameter( + mincomponentdiag=pml.Percentage(min_d) + ) + + if min_f > 0: + ms.meshing_remove_connected_component_by_face_number(mincomponentsize=min_f) + + if repair: + # ms.meshing_remove_t_vertices(method=0, threshold=40, repeat=True) + ms.meshing_repair_non_manifold_edges(method=0) + ms.meshing_repair_non_manifold_vertices(vertdispratio=0) + + if remesh: + # ms.apply_coord_taubin_smoothing() + ms.meshing_isotropic_explicit_remeshing( + iterations=3, targetlen=pml.AbsoluteValue(remesh_size) + ) + + # extract mesh + m = ms.current_mesh() + verts = m.vertex_matrix() + faces = m.face_matrix() + + print( + f"[INFO] mesh cleaning: {_ori_vert_shape} --> {verts.shape}, {_ori_face_shape} --> {faces.shape}" + ) + + return verts, faces + +def construct_list_of_gs_attributes(features_dc, features_rest, scaling, rotation): + l = ['x', 'y', 'z', 'nx', 'ny', 'nz'] + for i in range(features_dc.shape[1]*features_dc.shape[2]): + l.append('f_dc_{}'.format(i)) + for i in range(features_rest.shape[1]*features_rest.shape[2]): + l.append('f_rest_{}'.format(i)) + l.append('opacity') + for i in range(scaling.shape[1]): + l.append('scale_{}'.format(i)) + for i in range(rotation.shape[1]): + l.append('rot_{}'.format(i)) + return l + +def calculate_max_sh_degree_from_gs_ply(plydata): + extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")] + #assert len(extra_f_names)!=3*(max_sh_degree + 1) ** 2 - 3: + max_sh_degree = int(((len(extra_f_names) + 3) / 3) ** 0.5 - 1) + return max_sh_degree, extra_f_names + +def write_gs_ply(xyz, normals, f_dc, f_rest, opacities, scale, rotation, list_of_attributes): + dtype_full = [(attribute, 'f4') for attribute in list_of_attributes] + + elements = np.empty(xyz.shape[0], dtype=dtype_full) + attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1) + elements[:] = list(map(tuple, attributes)) + el = PlyElement.describe(elements, 'vertex') + return PlyData([el]) + +def read_gs_ply(plydata): + xyz = np.stack((np.asarray(plydata.elements[0]["x"]), + np.asarray(plydata.elements[0]["y"]), + np.asarray(plydata.elements[0]["z"])), axis=1) + opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis] + + features_dc = np.zeros((xyz.shape[0], 3, 1)) + features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"]) + features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"]) + features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"]) + + max_sh_degree, extra_f_names = calculate_max_sh_degree_from_gs_ply(plydata) + + features_extra = np.zeros((xyz.shape[0], len(extra_f_names))) + for idx, attr_name in enumerate(extra_f_names): + features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name]) + # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC) + features_extra = features_extra.reshape((features_extra.shape[0], 3, (max_sh_degree + 1) ** 2 - 1)) + + scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")] + scales = np.zeros((xyz.shape[0], len(scale_names))) + for idx, attr_name in enumerate(scale_names): + scales[:, idx] = np.asarray(plydata.elements[0][attr_name]) + + rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")] + rots = np.zeros((xyz.shape[0], len(rot_names))) + for idx, attr_name in enumerate(rot_names): + rots[:, idx] = np.asarray(plydata.elements[0][attr_name]) + + return xyz, features_dc, features_extra, opacities, scales, rots + +def ply_to_points_cloud(plydata): + xyz, features_dc, features_extra, opacities, scales, rots = read_gs_ply(plydata) + + features_dc = np.transpose(features_dc, (0, 2, 1)) # equivalent of torch.transpose(features_dc, 1, 2) + features_extra = np.transpose(features_extra, (0, 2, 1)) + shs = np.concatenate((features_dc, features_extra), axis=1) + normals = np.zeros_like(xyz) + pcd = PointCloud(points=xyz, colors=SH2RGB(shs), normals=normals) + return pcd + + +def get_target_axis_and_scale(axis_string, scale_value=1.0): + """ + Coordinate system inverts when: + 1. Any of the axis inverts + 2. Two of the axises switch + + If coordinate system inverts twice in a row then it will not be inverted + """ + axis_names = ["x", "y", "z"] + + target_axis, target_scale, coordinate_invert_count = [], [], 0 + axis_switch_count = 0 + for i in range(len(axis_names)): + s = axis_string[i] + if s[0] == "-": + target_scale.append(-scale_value) + coordinate_invert_count += 1 + else: + target_scale.append(scale_value) + + new_axis_i = axis_names.index(s[1]) + if new_axis_i != i: + axis_switch_count += 1 + target_axis.append(new_axis_i) + + if axis_switch_count == 2: + coordinate_invert_count += 1 + + return target_axis, target_scale, coordinate_invert_count + +def switch_vector_axis(vector3s, target_axis): + """ + Example: + vector3s = torch.tensor([[1, 2, 3], [3, 2, 1], [2, 3, 1]]) # shape (N, 3) + + target_axis = (2, 0, 1) # or [2, 0, 1] + vector3s[:, [0, 1, 2]] = vector3s[:, target_axis] + + # Result: tensor([[3, 1, 2], [1, 3, 2], [1, 2, 3]]) + """ + vector3s[:, [0, 1, 2]] = vector3s[:, target_axis] + return vector3s + +def switch_ply_axis_and_scale(plydata, target_axis, target_scale, coordinate_invert_count): + """ + Args: + target_axis (array): shape (3) + target_scale (array): shape (3) + """ + xyz, features_dc, features_extra, opacities, scales, rots = read_gs_ply(plydata) + normals = np.zeros_like(xyz) + features_dc_2d = features_dc.reshape(features_dc.shape[0], features_dc.shape[1]*features_dc.shape[2]) + features_extra_2d = features_extra.reshape(features_extra.shape[0], features_extra.shape[1]*features_extra.shape[2]) + + target_scale = torch.tensor(target_scale).float().cuda() + xyz = switch_vector_axis(torch.tensor(xyz).float().cuda() * target_scale, target_axis).detach().cpu().numpy() + scales = switch_vector_axis(torch.tensor(scales).float().cuda(), target_axis).detach().cpu().numpy() + + # change rotation representation from quaternion (w, x, y, z) to axis angle vector (x, y, z) to make swich axis easier + rots_axis_angle = quaternion_to_axis_angle(torch.tensor(rots).float().cuda()) + rots_axis_angle = switch_vector_axis(rots_axis_angle * target_scale, target_axis) + """ + Since axis–angle vector is composed of axis (unit vector/direction) and clockwise radians angle (vector magnitude), + so in order to invert the sign of angle when coordinate system inverts, we also need to invert the direction of axis–angle vector + """ + if coordinate_invert_count % 2 != 0: + rots_axis_angle = -rots_axis_angle + rots = axis_angle_to_quaternion(rots_axis_angle).detach().cpu().numpy() + + return write_gs_ply(xyz, normals, features_dc_2d, features_extra_2d, opacities, scales, rots, construct_list_of_gs_attributes(features_dc, features_extra, scales, rots)) + +def switch_mesh_axis_and_scale(mesh, target_axis, target_scale, flip_normal=False): + """ + Args: + target_axis (array): shape (3) + target_scale (array): shape (3) + """ + target_scale = torch.tensor(target_scale).float().cuda() + mesh.v = switch_vector_axis(mesh.v * target_scale, target_axis) + mesh.vn = switch_vector_axis(mesh.vn * target_scale, target_axis) + if flip_normal: + mesh.vn *= -1 + return mesh + + +def marching_cubes_density_to_mesh(get_density_func, grid_size=256, S=128, density_thresh=10, decimate_target=5e4): + from mcubes import marching_cubes + from kiui.mesh_utils import clean_mesh, decimate_mesh + + + sigmas = np.zeros([grid_size, grid_size, grid_size], dtype=np.float32) + + X = torch.linspace(-1, 1, grid_size).split(S) + Y = torch.linspace(-1, 1, grid_size).split(S) + Z = torch.linspace(-1, 1, grid_size).split(S) + + for xi, xs in enumerate(X): + for yi, ys in enumerate(Y): + for zi, zs in enumerate(Z): + xx, yy, zz = torch.meshgrid(xs, ys, zs, indexing='ij') + pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1) # [S, 3] + val = get_density_func(pts) + sigmas[xi * S: xi * S + len(xs), yi * S: yi * S + len(ys), zi * S: zi * S + len(zs)] = val.reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy() # [S, 1] --> [x, y, z] + + print(f'[INFO] marching cubes thresh: {density_thresh} ({sigmas.min()} ~ {sigmas.max()})') + + vertices, triangles = marching_cubes(sigmas, density_thresh) + vertices = vertices / (grid_size - 1.0) * 2 - 1 + + # clean + vertices = vertices.astype(np.float32) + triangles = triangles.astype(np.int32) + vertices, triangles = clean_mesh(vertices, triangles, remesh=True, remesh_size=0.01) + if triangles.shape[0] > decimate_target: + vertices, triangles = decimate_mesh(vertices, triangles, decimate_target, optimalplacement=False) + + return vertices, triangles + +def color_func_to_albedo(mesh, get_rgb_func, texture_resolution=1024, padding=2, batch_size=640000, device="cuda", force_cuda_rast=False, glctx=None): + import nvdiffrast.torch as dr + from kiui.op import uv_padding + + if glctx is None: + if force_cuda_rast: + glctx = dr.RasterizeCudaContext() + else: + glctx = dr.RasterizeGLContext() + + # render uv maps + h = w = texture_resolution + if mesh.vt is None: + mesh.auto_uv() + + uv = mesh.vt * 2.0 - 1.0 # uvs to range [-1, 1] + uv = torch.cat((uv, torch.zeros_like(uv[..., :1]), torch.ones_like(uv[..., :1])), dim=-1) # [N, 4] + + rast, _ = dr.rasterize(glctx, uv.unsqueeze(0), mesh.ft, (h, w)) # [1, h, w, 4] + xyzs, _ = dr.interpolate(mesh.v.unsqueeze(0), rast, mesh.f) # [1, h, w, 3] + mask, _ = dr.interpolate(torch.ones_like(mesh.v[:, :1]).unsqueeze(0), rast, mesh.f) # [1, h, w, 1] + + # masked query + xyzs = xyzs.view(-1, 3) + mask = (mask > 0).view(-1) + + albedo = torch.zeros(h * w, 3, device=device, dtype=torch.float32) + + if mask.any(): + print(f"[INFO] querying texture...") + + xyzs = xyzs[mask] # [M, 3] + + # batched inference to avoid OOM + batch = [] + head = 0 + while head < xyzs.shape[0]: + tail = min(head + batch_size, xyzs.shape[0]) + batch.append(get_rgb_func(xyzs[head:tail]).float()) + head += batch_size + + albedo[mask] = torch.cat(batch, dim=0) + + albedo = albedo.view(h, w, -1) + mask = mask.view(h, w) + albedo = uv_padding(albedo, mask, padding) + + return albedo + +@torch.no_grad() +def K_nearest_neighbors_func( + points: torch.Tensor, + K: int, + query: Optional[torch.Tensor] = None, + return_dist=False, + skip_first_index=True, +): + if not pytorch3d_capable: + raise ImportError("pytorch3d is not installed, which is required for KNN") + + # query/points: Tensor of shape (N, P1/P2, D) giving a batch of N point clouds, each containing up to P1/P2 points of dimension D + if query is None: + query = points + dist, idx, nn = knn_points(query[None, ...], points[None, ...], K=K, return_nn=True) + + # idx: Tensor of shape (N, P1, K) + # nn: Tensor of shape (N, P1, K, D) + + # take the index 1 since index 0 is the point itself + if skip_first_index: + if not return_dist: + return nn[0, :, 1:, :], idx[0, :, 1:] + else: + return nn[0, :, 1:, :], idx[0, :, 1:], dist[0, :, 1:] + else: + if not return_dist: + return nn[0, :, :, :], idx[0, :, :] + 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: + mesh.auto_uv() + + # Get Faces on UV + texture_size_minus_one = texture_size - 1 + faces = mesh.ft + verts_uvs = mesh.vt * texture_size_minus_one + verts_uvs_idx = verts_uvs.to(torch.long) + + vmapping = mesh.get_default_vt_to_v_mapping() + # Get vertex colors + interpolate_color = interpolate_color and mesh.vc is not None + texture_map = mesh.albedo + if interpolate_color: + verts_colors = mesh.vc[vmapping] + # Create a blank texture map + texture_map = torch.zeros((texture_size, texture_size, 3), device=verts_colors.device) + + # Get vertex positions + position_map = None + if interpolate_position: + verts_positions = mesh.v[vmapping] + position_map = torch.zeros((texture_size, texture_size, 3), device=verts_colors.device) + + # Create a grid of UV coordinates + grid_x, grid_y = torch.meshgrid(torch.linspace(0, texture_size_minus_one, texture_size, dtype=verts_uvs.dtype, device=verts_colors.device), torch.linspace(0, texture_size_minus_one, texture_size, dtype=verts_uvs.dtype, device=verts_colors.device)) + grid = torch.stack([grid_x, grid_y], dim=-1) + + batch_idx_range = torch.arange(0, batch_size**2, 1, device=verts_colors.device, dtype=torch.long) + + all_face_verts = verts_uvs_idx[faces] + for uv_i in range(0, texture_size_minus_one, batch_size): + for uv_j in range(0, texture_size_minus_one, batch_size): + # Divide faces into squre batchs + faces_mask = (uv_i <= all_face_verts[:, :, 0]) & (all_face_verts[:, :, 0] <= uv_i + batch_size) & (uv_j <= all_face_verts[:, :, 1]) & (all_face_verts[:, :, 1] <= uv_j + batch_size) + faces_mask, _ = torch.max(faces_mask, dim=1) + faces_mask = faces_mask.nonzero().squeeze(-1) + batch_faces = faces[faces_mask] + + uvs = verts_uvs[batch_faces] + if interpolate_color: + colors = verts_colors[batch_faces] + if interpolate_position: + positions = verts_positions[batch_faces] + + v0, v1, v2 = uvs[:, 0], uvs[:, 1], uvs[:, 2] + v0_0, v0_1 = v0[:, 0], v0[:, 1] + v1_0, v1_1 = v1[:, 0], v1[:, 1] + v2_0, v2_1 = v2[:, 0], v2[:, 1] + + sub_grid = grid[uv_i:uv_i + batch_size, uv_j:uv_j + batch_size].reshape(-1, 2).unsqueeze(1) + + # Compute barycentric coordinates for the batch of faces + denom = (v1_1 - v2_1) * (v0_0 - v2_0) + (v2_0 - v1_0) * (v0_1 - v2_1) + a = ((v1_1 - v2_1) * (sub_grid[:, :, 0] - v2_0) + (v2_0 - v1_0) * (sub_grid[:, :, 1] - v2_1)) / denom + b = ((v2_1 - v0_1) * (sub_grid[:, :, 0] - v2_0) + (v0_0 - v2_0) * (sub_grid[:, :, 1] - v2_1)) / denom + c = 1 - a - b + + # Mask for points inside the triangle + mask = (a >= 0) & (b >= 0) & (c >= 0) + mask, mask_idx = torch.max(mask, dim=1) + + # Get texture coordinates + sub_grid = sub_grid.squeeze(1) + uv_coords = sub_grid[mask].to(torch.long) + + # Interpolate colors + a, b, c = a[batch_idx_range, mask_idx][mask], b[batch_idx_range, mask_idx][mask], c[batch_idx_range, mask_idx][mask] + + if interpolate_color: + colors = colors[mask_idx][mask] + interpolated_colors = (a[:, None] * colors[:, 0] + (b[:, None] * colors[:, 1]) + (c[:, None] * colors[:, 2])) + texture_map[uv_coords[:, 1], uv_coords[:, 0]] = interpolated_colors + if interpolate_position: + positions = positions[mask_idx][mask] + interpolated_positions = (a[:, None] * positions[:, 0] + (b[:, None] * positions[:, 1]) + (c[:, None] * positions[:, 2])) + position_map[uv_coords[:, 1], uv_coords[:, 0]] = interpolated_positions + + return texture_map, position_map \ No newline at end of file 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) -