From a295aa69592a0d6f24bfbdd3746eca559d49c97e Mon Sep 17 00:00:00 2001 From: DenisKochetov Date: Wed, 3 Sep 2025 02:21:37 +0200 Subject: [PATCH] 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) -