From ae3156305dfdebd9ad91bba53cf20968caa6d55e Mon Sep 17 00:00:00 2001 From: DenisKochetov Date: Mon, 1 Sep 2025 13:04:58 +0200 Subject: [PATCH] renamed mesh_processor --- Gen_3D_Modules/CRM/model/crm/model.py | 2 +- .../hy3dpaint/mesh_inpaint_wheel/setup.py | 12 +- .../LGM/nerf_marching_cubes_converter.py | 4 +- .../MV_Adapter/mvadapter_node_utils.py | 2 +- .../TriplaneGaussian/models/renderer.py | 2 +- Gen_3D_Modules/TripoSR/models/isosurface.py | 2 +- .../Unique3D/scripts/project_mesh.py | 2 +- .../DiffRastMesh/diff_mesh_renderer.py | 2 +- .../FlexiCubes/flexicubes_trainer.py | 2 +- .../GaussianSplatting/main_3DGS_renderer.py | 4 +- __init__.py | 4 +- mesh_processer/__init__.py | 0 mesh_processer/mesh.py | 906 ------------------ mesh_processer/mesh_utils.py | 681 ------------- 14 files changed, 19 insertions(+), 1606 deletions(-) delete mode 100644 mesh_processer/__init__.py delete mode 100644 mesh_processer/mesh.py delete mode 100644 mesh_processer/mesh_utils.py diff --git a/Gen_3D_Modules/CRM/model/crm/model.py b/Gen_3D_Modules/CRM/model/crm/model.py index fa4ca39..27ba4c9 100644 --- a/Gen_3D_Modules/CRM/model/crm/model.py +++ b/Gen_3D_Modules/CRM/model/crm/model.py @@ -14,7 +14,7 @@ from CRM.util.renderer import Renderer from CRM.model.archs.mlp_head import SdfMlp, RgbMlp import xatlas -from mesh_processer.mesh import Mesh +from mesh_processor.mesh import Mesh class Dummy: diff --git a/Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/mesh_inpaint_wheel/setup.py b/Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/mesh_inpaint_wheel/setup.py index b72b316..ae470b1 100644 --- a/Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/mesh_inpaint_wheel/setup.py +++ b/Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/mesh_inpaint_wheel/setup.py @@ -1,7 +1,6 @@ from setuptools import setup, find_packages import platform -# Определяем платформу для правильного именования wheel if platform.system() == "Windows": plat_name = "win_amd64" elif platform.system() == "Darwin": @@ -13,14 +12,14 @@ setup( name="mesh_inpaint_processor", version="1.0.0", description="Fast C++ mesh inpainting processor with Python fallback", - long_description="Компилированный модуль для быстрой обработки инпейнтинга мешей", + long_description="Compiled module for fast mesh inpainting processing", author="Hunyuan3D Team", packages=find_packages(), package_data={ 'mesh_inpaint_processor': [ - '*.pyd', # Включаем .pyd файлы - '*.so', # На случай Linux версии - '*.py' # Python fallback файлы + '*.pyd', + '*.so', # For Linux version + '*.py' # Python fallback files ] }, include_package_data=True, @@ -41,7 +40,6 @@ setup( "Operating System :: POSIX :: Linux", "Operating System :: MacOS", ], - # Указываем что это binary distribution zip_safe=False, - has_ext_modules=lambda: True, # Указывает что есть нативные модули + has_ext_modules=lambda: True, ) \ No newline at end of file diff --git a/Gen_3D_Modules/LGM/nerf_marching_cubes_converter.py b/Gen_3D_Modules/LGM/nerf_marching_cubes_converter.py index 3d7e9bd..978ec85 100644 --- a/Gen_3D_Modules/LGM/nerf_marching_cubes_converter.py +++ b/Gen_3D_Modules/LGM/nerf_marching_cubes_converter.py @@ -15,8 +15,8 @@ from kiui.op import inverse_sigmoid from kiui.cam import orbit_camera, get_perspective from kiui.nn import MLP, trunc_exp -from mesh_processer.mesh import Mesh -from mesh_processer.mesh_utils import marching_cubes_density_to_mesh, color_func_to_albedo +from mesh_processor.mesh import Mesh +from mesh_processor.mesh_utils import marching_cubes_density_to_mesh, color_func_to_albedo from LGM.core.options import Options from LGM.core.gs import GaussianRenderer from LGM.core.utils import get_rays diff --git a/Gen_3D_Modules/MV_Adapter/mvadapter_node_utils.py b/Gen_3D_Modules/MV_Adapter/mvadapter_node_utils.py index 3a98854..d846ee2 100644 --- a/Gen_3D_Modules/MV_Adapter/mvadapter_node_utils.py +++ b/Gen_3D_Modules/MV_Adapter/mvadapter_node_utils.py @@ -540,7 +540,7 @@ def prepare_tg2mv_pipeline( if hasattr(pipe, 'enable_model_cpu_offload'): if not use_mmgp: pipe.enable_model_cpu_offload() - print("[INFO] Включен CPU offload для экономии VRAM") + print("[INFO] CPU offload for memory optimization") if use_mmgp and MMGP_AVAILABLE: torch.cuda.empty_cache() diff --git a/Gen_3D_Modules/TriplaneGaussian/models/renderer.py b/Gen_3D_Modules/TriplaneGaussian/models/renderer.py index beec388..c176a95 100644 --- a/Gen_3D_Modules/TriplaneGaussian/models/renderer.py +++ b/Gen_3D_Modules/TriplaneGaussian/models/renderer.py @@ -13,7 +13,7 @@ from ..utils.ops import scale_tensor from .networks import MLP from einops import rearrange, reduce -from mesh_processer.mesh_utils import construct_list_of_gs_attributes, write_gs_ply +from mesh_processor.mesh_utils import construct_list_of_gs_attributes, write_gs_ply inverse_sigmoid = lambda x: np.log(x / (1 - x)) diff --git a/Gen_3D_Modules/TripoSR/models/isosurface.py b/Gen_3D_Modules/TripoSR/models/isosurface.py index c4a9e31..e8acb3d 100644 --- a/Gen_3D_Modules/TripoSR/models/isosurface.py +++ b/Gen_3D_Modules/TripoSR/models/isosurface.py @@ -4,7 +4,7 @@ import torch import torch.nn as nn from mcubes import marching_cubes -from mesh_processer.mesh_utils import switch_vector_axis +from mesh_processor.mesh_utils import switch_vector_axis class IsosurfaceHelper(nn.Module): diff --git a/Gen_3D_Modules/Unique3D/scripts/project_mesh.py b/Gen_3D_Modules/Unique3D/scripts/project_mesh.py index 1782e07..a86ac0f 100644 --- a/Gen_3D_Modules/Unique3D/scripts/project_mesh.py +++ b/Gen_3D_Modules/Unique3D/scripts/project_mesh.py @@ -12,7 +12,7 @@ from pytorch3d.renderer import ( FoVOrthographicCameras, ) from pytorch3d.renderer import MeshRasterizer -from mesh_processer.mesh import Mesh +from mesh_processor.mesh import Mesh def get_camera(world_to_cam, fov_in_degrees=60, focal_length=1 / (2**0.5), cam_type='fov'): # pytorch3d expects transforms as row-vectors, so flip rotation: https://github.com/facebookresearch/pytorch3d/issues/1183 diff --git a/MVs_Algorithms/DiffRastMesh/diff_mesh_renderer.py b/MVs_Algorithms/DiffRastMesh/diff_mesh_renderer.py index 98a07da..d559a14 100644 --- a/MVs_Algorithms/DiffRastMesh/diff_mesh_renderer.py +++ b/MVs_Algorithms/DiffRastMesh/diff_mesh_renderer.py @@ -9,7 +9,7 @@ import nvdiffrast.torch as dr from kiui.op import inverse_sigmoid -from mesh_processer.mesh import safe_normalize +from mesh_processor.mesh import safe_normalize def scale_img_nhwc(x, size, mag='bilinear', min='bilinear'): assert (x.shape[1] >= size[0] and x.shape[2] >= size[1]) or (x.shape[1] < size[0] and x.shape[2] < size[1]), "Trying to magnify image in one dimension and minify in the other" diff --git a/MVs_Algorithms/FlexiCubes/flexicubes_trainer.py b/MVs_Algorithms/FlexiCubes/flexicubes_trainer.py index 41ba267..cccce84 100644 --- a/MVs_Algorithms/FlexiCubes/flexicubes_trainer.py +++ b/MVs_Algorithms/FlexiCubes/flexicubes_trainer.py @@ -13,7 +13,7 @@ from .util import SimpleMesh from .loss import sdf_reg_loss from shared_utils.camera_utils import OrbitCamera -from mesh_processer.mesh import Mesh +from mesh_processor.mesh import Mesh def lr_schedule(iter): diff --git a/MVs_Algorithms/GaussianSplatting/main_3DGS_renderer.py b/MVs_Algorithms/GaussianSplatting/main_3DGS_renderer.py index d05394e..0715575 100644 --- a/MVs_Algorithms/GaussianSplatting/main_3DGS_renderer.py +++ b/MVs_Algorithms/GaussianSplatting/main_3DGS_renderer.py @@ -15,8 +15,8 @@ from kornia.geometry.conversions import ( from kiui.op import inverse_sigmoid from shared_utils.sh_utils import eval_sh, SH2RGB, RGB2SH -from mesh_processer.mesh import Mesh, PointCloud -from mesh_processer.mesh_utils import construct_list_of_gs_attributes, write_gs_ply, read_gs_ply, K_nearest_neighbors_func +from mesh_processor.mesh import Mesh, PointCloud +from mesh_processor.mesh_utils import construct_list_of_gs_attributes, write_gs_ply, read_gs_ply, K_nearest_neighbors_func def get_expon_lr_func( lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000 diff --git a/__init__.py b/__init__.py index cc8c101..bf931fc 100644 --- a/__init__.py +++ b/__init__.py @@ -8,10 +8,12 @@ import logging ROOT_PATH = os.path.dirname(os.path.realpath(__file__)) MODULE_PATH = os.path.join(ROOT_PATH, "Gen_3D_Modules") MV_ALGO_PATH = os.path.join(ROOT_PATH, "MVs_Algorithms") +MESH_PROCESSOR_PATH = os.path.join(ROOT_PATH, "mesh_processor") -sys.path.append(ROOT_PATH) +sys.path.insert(0, ROOT_PATH) sys.path.append(MODULE_PATH) sys.path.append(MV_ALGO_PATH) +sys.path.append(MESH_PROCESSOR_PATH) import shutil import __main__ diff --git a/mesh_processer/__init__.py b/mesh_processer/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/mesh_processer/mesh.py b/mesh_processer/mesh.py deleted file mode 100644 index 9bcc41b..0000000 --- a/mesh_processer/mesh.py +++ /dev/null @@ -1,906 +0,0 @@ -import os -import cv2 -from typing import NamedTuple -import torch -import trimesh -import numpy as np - -from kiui.op import safe_normalize, dot -from kiui.typing import * - -from shared_utils.sh_utils import SH2RGB -from shared_utils.image_utils import prepare_torch_img - - -class Mesh: - """ - A torch-native trimesh class, with support for ``ply/obj/glb`` formats. - - Note: - This class only supports one mesh with a single texture image (an albedo texture and a metallic-roughness texture). - """ - def __init__( - self, - v: Optional[Tensor] = None, - f: Optional[Tensor] = None, - vn: Optional[Tensor] = None, - fn: Optional[Tensor] = None, - vt: Optional[Tensor] = None, - ft: Optional[Tensor] = None, - vc: Optional[Tensor] = None, # vertex color - albedo: Optional[Tensor] = None, - metallicRoughness: Optional[Tensor] = None, - device: Optional[torch.device] = None, - ): - """Init a mesh directly using all attributes. - - Args: - v (Optional[Tensor]): vertices, float [N, 3]. Defaults to None. - f (Optional[Tensor]): faces, int [M, 3]. Defaults to None. - vn (Optional[Tensor]): vertex normals, float [N, 3]. Defaults to None. - fn (Optional[Tensor]): faces for normals, int [M, 3]. Defaults to None. - vt (Optional[Tensor]): vertex uv coordinates, float [N, 2]. Defaults to None. - ft (Optional[Tensor]): faces for uvs, int [M, 3]. Defaults to None. - vc (Optional[Tensor]): vertex colors, float [N, 3]. Defaults to None. - albedo (Optional[Tensor]): albedo texture, float [H, W, 3], RGB format. Defaults to None. - metallicRoughness (Optional[Tensor]): metallic-roughness texture, float [H, W, 3], metallic(Blue) = metallicRoughness[..., 2], roughness(Green) = metallicRoughness[..., 1]. Defaults to None. - device (Optional[torch.device]): torch device. Defaults to None. - """ - self.device = device - self.v = v - self.vn = vn - self.vt = vt - self.f = f - self.fn = fn - self.ft = ft - # will first see if there is vertex color to use - self.vc = vc - # only support a single albedo image, (H, W, 3) - self.albedo = albedo - # pbr extension, metallic(Blue) = metallicRoughness[..., 2], roughness(Green) = metallicRoughness[..., 1] - # ref: https://registry.khronos.org/glTF/specs/2.0/glTF-2.0.html - self.metallicRoughness = metallicRoughness - - self.ori_center = 0 - self.ori_scale = 1 - - @classmethod - def load(cls, path, resize=True, renormal=True, retex=False, clean=False, bound=0.5, front_dir='+z', **kwargs): - """load mesh from path. - - Args: - path (str): path to mesh file, supports ply, obj, glb. - resize (bool, optional): auto resize the mesh using ``bound`` into [-bound, bound]^3. Defaults to True. - renormal (bool, optional): re-calc the vertex normals. Defaults to True. - retex (bool, optional): re-calc the uv coordinates, will overwrite the existing uv coordinates. Defaults to False. - clean (bool, optional): perform mesh cleaning at load (e.g., merge close vertices). Defaults to False. - bound (float, optional): bound to resize. Defaults to 0.9. - front_dir (str, optional): front-view direction of the mesh, should be [+-][xyz][ 123]. Defaults to '+z'. - device (torch.device, optional): torch device. Defaults to None. - - Note: - a ``device`` keyword argument can be provided to specify the torch device. - If it's not provided, we will try to use ``'cuda'`` as the device if it's available. - - Returns: - Mesh: the loaded Mesh object. - """ - # obj supports face uv - if path.endswith(".obj"): - mesh = cls.load_obj(path, **kwargs) - # trimesh only supports vertex uv, but can load more formats - else: - mesh = cls.load_trimesh(path, **kwargs) - - # clean - if clean: - from kiui.mesh_utils import clean_mesh - vertices = mesh.v.detach().cpu().numpy() - triangles = mesh.f.detach().cpu().numpy() - vertices, triangles = clean_mesh(vertices, triangles, remesh=False) - mesh.v = torch.from_numpy(vertices).contiguous().float().to(mesh.device) - mesh.f = torch.from_numpy(triangles).contiguous().int().to(mesh.device) - - print(f"[Mesh loading] v: {mesh.v.shape}, f: {mesh.f.shape}") - # auto-normalize - if resize: - mesh.auto_size(bound=bound) - # auto-fix normal - if renormal or mesh.vn is None: - mesh.auto_normal() - print(f"[Mesh loading] vn: {mesh.vn.shape}, fn: {mesh.fn.shape}") - # auto-fix texcoords - if retex or (mesh.albedo is not None and mesh.vt is None): - mesh.auto_uv(cache_path=path) - print(f"[Mesh loading] vt: {mesh.vt.shape}, ft: {mesh.ft.shape}") - - # rotate front dir to +z - if front_dir != "+z": - # axis switch - if "-z" in front_dir: - T = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, -1]], device=mesh.device, dtype=torch.float32) - elif "+x" in front_dir: - T = torch.tensor([[0, 0, 1], [0, 1, 0], [1, 0, 0]], device=mesh.device, dtype=torch.float32) - elif "-x" in front_dir: - T = torch.tensor([[0, 0, -1], [0, 1, 0], [1, 0, 0]], device=mesh.device, dtype=torch.float32) - elif "+y" in front_dir: - T = torch.tensor([[1, 0, 0], [0, 0, 1], [0, 1, 0]], device=mesh.device, dtype=torch.float32) - elif "-y" in front_dir: - T = torch.tensor([[1, 0, 0], [0, 0, -1], [0, 1, 0]], device=mesh.device, dtype=torch.float32) - else: - T = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], device=mesh.device, dtype=torch.float32) - # rotation (how many 90 degrees) - if '1' in front_dir: - T @= torch.tensor([[0, -1, 0], [1, 0, 0], [0, 0, 1]], device=mesh.device, dtype=torch.float32) - elif '2' in front_dir: - T @= torch.tensor([[1, 0, 0], [0, -1, 0], [0, 0, 1]], device=mesh.device, dtype=torch.float32) - elif '3' in front_dir: - T @= torch.tensor([[0, 1, 0], [-1, 0, 0], [0, 0, 1]], device=mesh.device, dtype=torch.float32) - mesh.v @= T - mesh.vn @= T - - return mesh - - # load from obj file - @classmethod - def load_obj(cls, path, albedo_path=None, device=None): - """load an ``obj`` mesh. - - Args: - path (str): path to mesh. - albedo_path (str, optional): path to the albedo texture image, will overwrite the existing texture path if specified in mtl. Defaults to None. - device (torch.device, optional): torch device. Defaults to None. - - Note: - We will try to read `mtl` path from `obj`, else we assume the file name is the same as `obj` but with `mtl` extension. - The `usemtl` statement is ignored, and we only use the last material path in `mtl` file. - - Returns: - Mesh: the loaded Mesh object. - """ - assert os.path.splitext(path)[-1] == ".obj" - - mesh = cls() - - # device - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - mesh.device = device - - # load obj - with open(path, "r") as f: - lines = f.readlines() - - def parse_f_v(fv): - # pass in a vertex term of a face, return {v, vt, vn} (-1 if not provided) - # supported forms: - # f v1 v2 v3 - # f v1/vt1 v2/vt2 v3/vt3 - # f v1/vt1/vn1 v2/vt2/vn2 v3/vt3/vn3 - # f v1//vn1 v2//vn2 v3//vn3 - xs = [int(x) - 1 if x != "" else -1 for x in fv.split("/")] - xs.extend([-1] * (3 - len(xs))) - return xs[0], xs[1], xs[2] - - vertices, texcoords, normals = [], [], [] - faces, tfaces, nfaces = [], [], [] - mtl_path = None - - for line in lines: - split_line = line.split() - # empty line - if len(split_line) == 0: - continue - prefix = split_line[0].lower() - # mtllib - if prefix == "mtllib": - mtl_path = split_line[1] - # usemtl - elif prefix == "usemtl": - pass # ignored - # v/vn/vt - elif prefix == "v": - vertices.append([float(v) for v in split_line[1:]]) - elif prefix == "vn": - normals.append([float(v) for v in split_line[1:]]) - elif prefix == "vt": - val = [float(v) for v in split_line[1:]] - texcoords.append([val[0], 1.0 - val[1]]) - elif prefix == "f": - vs = split_line[1:] - nv = len(vs) - v0, t0, n0 = parse_f_v(vs[0]) - for i in range(nv - 2): # triangulate (assume vertices are ordered) - v1, t1, n1 = parse_f_v(vs[i + 1]) - v2, t2, n2 = parse_f_v(vs[i + 2]) - faces.append([v0, v1, v2]) - tfaces.append([t0, t1, t2]) - nfaces.append([n0, n1, n2]) - - mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device) - mesh.vt = ( - torch.tensor(texcoords, dtype=torch.float32, device=device) - if len(texcoords) > 0 - else None - ) - mesh.vn = ( - torch.tensor(normals, dtype=torch.float32, device=device) - if len(normals) > 0 - else None - ) - - mesh.f = torch.tensor(faces, dtype=torch.int32, device=device) - mesh.ft = ( - torch.tensor(tfaces, dtype=torch.int32, device=device) - if len(texcoords) > 0 - else None - ) - mesh.fn = ( - torch.tensor(nfaces, dtype=torch.int32, device=device) - if len(normals) > 0 - else None - ) - - # see if there is vertex color - use_vertex_color = False - if mesh.v.shape[1] == 6: - use_vertex_color = True - mesh.vc = mesh.v[:, 3:] - mesh.v = mesh.v[:, :3] - print(f"[load_obj] use vertex color: {mesh.vc.shape}") - - # try to load texture image - if not use_vertex_color: - # try to retrieve mtl file - mtl_path_candidates = [] - if mtl_path is not None: - mtl_path_candidates.append(mtl_path) - mtl_path_candidates.append(os.path.join(os.path.dirname(path), mtl_path)) - mtl_path_candidates.append(path.replace(".obj", ".mtl")) - - mtl_path = None - for candidate in mtl_path_candidates: - if os.path.exists(candidate): - mtl_path = candidate - break - - # if albedo_path is not provided, try retrieve it from mtl - metallic_path = None - roughness_path = None - if mtl_path is not None and albedo_path is None: - with open(mtl_path, "r") as f: - lines = f.readlines() - - for line in lines: - split_line = line.split() - # empty line - if len(split_line) == 0: - continue - prefix = split_line[0] - - if "map_Kd" in prefix: - # assume relative path! - albedo_path = os.path.join(os.path.dirname(path), split_line[1]) - print(f"[load_obj] use texture from: {albedo_path}") - elif "map_Pm" in prefix: - metallic_path = os.path.join(os.path.dirname(path), split_line[1]) - elif "map_Pr" in prefix: - roughness_path = os.path.join(os.path.dirname(path), split_line[1]) - - # still not found albedo_path, or the path doesn't exist - if albedo_path is None or not os.path.exists(albedo_path): - # init an empty texture - print(f"[load_obj] init empty albedo!") - # albedo = np.random.rand(1024, 1024, 3).astype(np.float32) - albedo = np.ones((1024, 1024, 3), dtype=np.float32) * np.array([0.5, 0.5, 0.5]) # default color - else: - albedo = cv2.imread(albedo_path, cv2.IMREAD_UNCHANGED) - albedo = cv2.cvtColor(albedo, cv2.COLOR_BGR2RGB) - albedo = albedo.astype(np.float32) / 255 - print(f"[load_obj] load texture: {albedo.shape}") - - mesh.albedo = torch.tensor(albedo, dtype=torch.float32, device=device) - - # try to load metallic and roughness - if metallic_path is not None and roughness_path is not None: - print(f"[load_obj] load metallicRoughness from: {metallic_path}, {roughness_path}") - metallic = cv2.imread(metallic_path, cv2.IMREAD_UNCHANGED) - metallic = metallic.astype(np.float32) / 255 - roughness = cv2.imread(roughness_path, cv2.IMREAD_UNCHANGED) - roughness = roughness.astype(np.float32) / 255 - metallicRoughness = np.stack([np.zeros_like(metallic), roughness, metallic], axis=-1) - - mesh.metallicRoughness = torch.tensor(metallicRoughness, dtype=torch.float32, device=device).contiguous() - - return mesh - - @classmethod - def load_trimesh(cls, path=None, given_mesh=None, device=None): - """load a mesh using ``trimesh.load()``. - - Can load various formats like ``glb`` and serves as a fallback. - - Note: - We will try to merge all meshes if the glb contains more than one, - but **this may cause the texture to lose**, since we only support one texture image! - - Args: - path (str): path to the mesh file. - given_mesh (trimesh): trimesh instance - device (torch.device, optional): torch device. Defaults to None. - - Returns: - Mesh: the loaded Mesh object. - """ - mesh = cls() - - # device - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - mesh.device = device - - # use trimesh to load ply/glb, assume only has one single RootMesh... - if given_mesh is not None: - _mesh = given_mesh - elif path is not None: - _data = trimesh.load(path) - if isinstance(_data, trimesh.Scene): - if len(_data.geometry) == 1: - _mesh = list(_data.geometry.values())[0] - else: - print(f"[load_trimesh] concatenating {len(_data.geometry)} meshes.") - _concat = [] - # loop the scene graph and apply transform to each mesh - scene_graph = _data.graph.to_flattened() # dict {name: {transform: 4x4 mat, geometry: str}} - for k, v in scene_graph.items(): - name = v['geometry'] - if name in _data.geometry and isinstance(_data.geometry[name], trimesh.Trimesh): - transform = v['transform'] - _concat.append(_data.geometry[name].apply_transform(transform)) - _mesh = trimesh.util.concatenate(_concat) - else: - _mesh = _data - else: - print(f"[load_trimesh] failed to load mesh, either path or given_mesh must be given") - return None - - should_create_empty_albedo = False - if _mesh.visual.kind == 'vertex': - vertex_colors = _mesh.visual.vertex_colors - vertex_colors = np.array(vertex_colors[..., :3]).astype(np.float32) / 255 - mesh.vc = torch.tensor(vertex_colors, dtype=torch.float32, device=device) - print(f"[load_trimesh] use vertex color: {mesh.vc.shape}") - elif _mesh.visual.kind == 'texture': - _material = _mesh.visual.material - if isinstance(_material, trimesh.visual.material.PBRMaterial): - texture = np.array(_material.baseColorTexture).astype(np.float32) / 255 - # load metallicRoughness if present - if _material.metallicRoughnessTexture is not None: - metallicRoughness = np.array(_material.metallicRoughnessTexture).astype(np.float32) / 255 - mesh.metallicRoughness = torch.tensor(metallicRoughness, dtype=torch.float32, device=device).contiguous() - elif isinstance(_material, trimesh.visual.material.SimpleMaterial): - texture = np.array(_material.to_pbr().baseColorTexture).astype(np.float32) / 255 - else: - raise NotImplementedError(f"material type {type(_material)} not supported!") - - if texture.ndim ==3: - mesh.albedo = torch.tensor(texture[..., :3], dtype=torch.float32, device=device).contiguous() - print(f"[load_trimesh] loaded albedo texture: {texture.shape}") - else: - should_create_empty_albedo = True - else: - should_create_empty_albedo = True - - if should_create_empty_albedo: - mesh.set_new_albedo(1024, 1024) - print(f"[load_trimesh] failed to load albedo texture, create empty abbedo texture with shape: {mesh.albedo.shape}") - - vertices = _mesh.vertices - - try: - texcoords = _mesh.visual.uv - texcoords[:, 1] = 1 - texcoords[:, 1] - except Exception as e: - texcoords = None - - try: - normals = _mesh.vertex_normals - except Exception as e: - normals = None - - # trimesh only support vertex uv... - faces = tfaces = nfaces = _mesh.faces - - mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device) - mesh.vt = ( - torch.tensor(texcoords, dtype=torch.float32, device=device) - if texcoords is not None - else None - ) - mesh.vn = ( - torch.tensor(normals, dtype=torch.float32, device=device) - if normals is not None - else None - ) - - mesh.f = torch.tensor(faces, dtype=torch.int32, device=device) - mesh.ft = ( - torch.tensor(tfaces, dtype=torch.int32, device=device) - if texcoords is not None - else None - ) - mesh.fn = ( - torch.tensor(nfaces, dtype=torch.int32, device=device) - if normals is not None - else None - ) - - return mesh - - def set_new_albedo(self, res_H, res_W): - if self.albedo is None: - texture = np.ones((res_H, res_W, 3), dtype=np.float32) * np.array([0.5, 0.5, 0.5]) - self.albedo = torch.tensor(texture, dtype=torch.float32, device=self.device) - else: - self.albedo = prepare_torch_img(self.albedo.unsqueeze(0), res_H, res_W, self.device).squeeze(0).permute(1, 2, 0).contiguous() # (1, 3, H, W) -> (H, W, 3) - - # aabb - def aabb(self): - """get the axis-aligned bounding box of the mesh. - - Returns: - Tuple[torch.Tensor]: the min xyz and max xyz of the mesh. - """ - return torch.min(self.v, dim=0).values, torch.max(self.v, dim=0).values - - # unit size - @torch.no_grad() - def auto_size(self, bound=0.9): - """auto resize the mesh. - - Args: - bound (float, optional): resizing into ``[-bound, bound]^3``. Defaults to 0.9. - """ - vmin, vmax = self.aabb() - self.ori_center = (vmax + vmin) / 2 - self.ori_scale = 2 * bound / torch.max(vmax - vmin).item() - self.v = (self.v - self.ori_center) * self.ori_scale - - def auto_normal(self): - """auto calculate the vertex normals. - """ - i0, i1, i2 = self.f[:, 0].long(), self.f[:, 1].long(), self.f[:, 2].long() - v0, v1, v2 = self.v[i0, :], self.v[i1, :], self.v[i2, :] - - face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1) - - # Splat face normals to vertices - vn = torch.zeros_like(self.v) - vn.scatter_add_(0, i0[:, None].repeat(1, 3), face_normals) - vn.scatter_add_(0, i1[:, None].repeat(1, 3), face_normals) - vn.scatter_add_(0, i2[:, None].repeat(1, 3), face_normals) - - # Normalize, replace zero (degenerated) normals with some default value - vn = torch.where( - dot(vn, vn) > 1e-20, - vn, - torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device=vn.device), - ) - vn = safe_normalize(vn) - - self.vn = vn - self.fn = self.f - - def auto_uv(self, cache_path=None, vmap=True): - """auto calculate the uv coordinates. - - Args: - cache_path (str, optional): path to save/load the uv cache as a npz file, this can avoid calculating uv every time when loading the same mesh, which is time-consuming. Defaults to None. - vmap (bool, optional): remap vertices based on uv coordinates, so each v correspond to a unique vt (necessary for formats like gltf). - Usually this will duplicate the vertices on the edge of uv atlas. Defaults to True. - """ - # try to load cache - if cache_path is not None: - cache_path = os.path.splitext(cache_path)[0] + "_uv.npz" - if cache_path is not None and os.path.exists(cache_path): - data = np.load(cache_path) - vt_np, ft_np, vmapping = data["vt"], data["ft"], data["vmapping"] - else: - import xatlas - - v_np = self.v.detach().cpu().numpy() - f_np = self.f.detach().int().cpu().numpy() - atlas = xatlas.Atlas() - atlas.add_mesh(v_np, f_np) - chart_options = xatlas.ChartOptions() - # chart_options.max_iterations = 4 - atlas.generate(chart_options=chart_options) - vmapping, ft_np, vt_np = atlas[0] # [N], [M, 3], [N, 2] - - # save to cache - if cache_path is not None: - np.savez(cache_path, vt=vt_np, ft=ft_np, vmapping=vmapping) - - vt = torch.from_numpy(vt_np.astype(np.float32)).to(self.device) - ft = torch.from_numpy(ft_np.astype(np.int32)).to(self.device) - self.vt = vt - self.ft = ft - - if vmap: - # remap v/f to vt/ft, so each v correspond to a unique vt. (necessary for gltf) - vmapping = torch.from_numpy(vmapping.astype(np.int64)).long().to(self.device) - self.align_v_to_vt(vmapping) - - def align_v_to_vt(self, vmapping=None): - """ remap v/f and vn/fn to vt/ft. - - Args: - vmapping (np.ndarray, optional): the mapping relationship from f to ft. Defaults to None. - """ - if vmapping is None: - vt2v_mapping = self.get_default_vt_to_v_mapping() - else: - vt2v_mapping = vmapping - - self.v = self.v[vt2v_mapping] - if self.vc is not None: - self.vc = self.vc[vt2v_mapping] - self.f = self.ft - # assume fn == f - if self.vn is not None: - if vmapping is None: - vt2vn_mapping = self.get_default_vt_to_vn_mapping() - else: - vt2vn_mapping = vmapping - self.vn = self.vn[vt2vn_mapping] - self.fn = self.ft - - def get_default_vt_to_v_mapping(self): - """map from ft to f - - Returns: - vmapping: tensor array - """ - ft = self.ft.view(-1).long() - f = self.f.view(-1).long() - vmapping = torch.zeros(self.vt.shape[0], dtype=torch.long, device=self.device) - vmapping[ft] = f # scatter, randomly choose one if index is not unique - return vmapping - - def get_default_vt_to_vn_mapping(self): - """map from ft to f - - Returns: - vmapping: tensor array - """ - ft = self.ft.view(-1).long() - fn = self.fn.view(-1).long() - vmapping = torch.zeros(self.vt.shape[0], dtype=torch.long, device=self.device) - vmapping[ft] = fn # scatter, randomly choose one if index is not unique - return vmapping - - def to(self, device): - """move all tensor attributes to device. - - Args: - device (torch.device): target device. - - Returns: - Mesh: self. - """ - self.device = device - for name in ["v", "f", "vn", "fn", "vt", "ft", "albedo", "vc", "metallicRoughness"]: - tensor = getattr(self, name) - if tensor is not None: - setattr(self, name, tensor.to(device)) - return self - - def write(self, path): - """write the mesh to a path. - - Args: - path (str): path to write, supports ply, obj and glb. - """ - if path.endswith(".ply"): - self.write_ply(path) - elif path.endswith(".obj"): - self.write_obj(path) - elif path.endswith(".glb") or path.endswith(".gltf"): - self.write_glb(path) - else: - raise NotImplementedError(f"format {path} not supported!") - - def write_ply(self, path): - """write the mesh in ply format. Only for geometry! - - Args: - path (str): path to write. - """ - - if self.albedo is not None: - print(f'[WARN] ply format does not support exporting texture, will ignore!') - - v_np = self.v.detach().cpu().numpy() - f_np = self.f.detach().cpu().numpy() - - _mesh = trimesh.Trimesh(vertices=v_np, faces=f_np) - _mesh.export(path) - - - def write_glb(self, path): - """write the mesh in glb/gltf format. - This will create a scene with a single mesh. - - Args: - path (str): path to write. - """ - - # assert self.v.shape[0] == self.vn.shape[0] and self.v.shape[0] == self.vt.shape[0] - if self.vt is not None and self.v.shape[0] != self.vt.shape[0]: - self.align_v_to_vt() - - import pygltflib - - f_np = self.f.detach().cpu().numpy().astype(np.uint32) - f_np_blob = f_np.flatten().tobytes() - - v_np = self.v.detach().cpu().numpy().astype(np.float32) - v_np_blob = v_np.tobytes() - - blob = f_np_blob + v_np_blob - byteOffset = len(blob) - - # base mesh - gltf = pygltflib.GLTF2( - scene=0, - scenes=[pygltflib.Scene(nodes=[0])], - nodes=[pygltflib.Node(mesh=0)], - meshes=[pygltflib.Mesh(primitives=[pygltflib.Primitive( - # indices to accessors (0 is triangles) - attributes=pygltflib.Attributes( - POSITION=1, - ), - indices=0, - )])], - buffers=[ - pygltflib.Buffer(byteLength=len(f_np_blob) + len(v_np_blob)) - ], - # buffer view (based on dtype) - bufferViews=[ - # triangles; as flatten (element) array - pygltflib.BufferView( - buffer=0, - byteLength=len(f_np_blob), - target=pygltflib.ELEMENT_ARRAY_BUFFER, # GL_ELEMENT_ARRAY_BUFFER (34963) - ), - # positions; as vec3 array - pygltflib.BufferView( - buffer=0, - byteOffset=len(f_np_blob), - byteLength=len(v_np_blob), - byteStride=12, # vec3 - target=pygltflib.ARRAY_BUFFER, # GL_ARRAY_BUFFER (34962) - ), - ], - accessors=[ - # 0 = triangles - pygltflib.Accessor( - bufferView=0, - componentType=pygltflib.UNSIGNED_INT, # GL_UNSIGNED_INT (5125) - count=f_np.size, - type=pygltflib.SCALAR, - max=[int(f_np.max())], - min=[int(f_np.min())], - ), - # 1 = positions - pygltflib.Accessor( - bufferView=1, - componentType=pygltflib.FLOAT, # GL_FLOAT (5126) - count=len(v_np), - type=pygltflib.VEC3, - max=v_np.max(axis=0).tolist(), - min=v_np.min(axis=0).tolist(), - ), - ], - ) - - # append texture info - if self.vt is not None and self.albedo is not None: - - vt_np = self.vt.detach().cpu().numpy().astype(np.float32) - vt_np_blob = vt_np.tobytes() - - albedo = self.albedo.detach().cpu().numpy() - albedo = (albedo * 255).astype(np.uint8) - albedo = cv2.cvtColor(albedo, cv2.COLOR_RGB2BGR) - albedo_blob = cv2.imencode('.png', albedo)[1].tobytes() - - # update primitive - gltf.meshes[0].primitives[0].attributes.TEXCOORD_0 = 2 - gltf.meshes[0].primitives[0].material = 0 - - # update materials - gltf.materials.append(pygltflib.Material( - pbrMetallicRoughness=pygltflib.PbrMetallicRoughness( - baseColorTexture=pygltflib.TextureInfo(index=0, texCoord=0), - metallicFactor=0.0, - roughnessFactor=1.0, - ), - alphaMode=pygltflib.OPAQUE, - alphaCutoff=None, - doubleSided=True, - )) - - gltf.textures.append(pygltflib.Texture(sampler=0, source=0)) - gltf.samplers.append(pygltflib.Sampler(magFilter=pygltflib.LINEAR, minFilter=pygltflib.LINEAR_MIPMAP_LINEAR, wrapS=pygltflib.REPEAT, wrapT=pygltflib.REPEAT)) - gltf.images.append(pygltflib.Image(bufferView=3, mimeType="image/png")) - - # update buffers - gltf.bufferViews.append( - # index = 2, texcoords; as vec2 array - pygltflib.BufferView( - buffer=0, - byteOffset=byteOffset, - byteLength=len(vt_np_blob), - byteStride=8, # vec2 - target=pygltflib.ARRAY_BUFFER, - ) - ) - - gltf.accessors.append( - # 2 = texcoords - pygltflib.Accessor( - bufferView=2, - componentType=pygltflib.FLOAT, - count=len(vt_np), - type=pygltflib.VEC2, - max=vt_np.max(axis=0).tolist(), - min=vt_np.min(axis=0).tolist(), - ) - ) - - blob += vt_np_blob - byteOffset += len(vt_np_blob) - - gltf.bufferViews.append( - # index = 3, albedo texture; as none target - pygltflib.BufferView( - buffer=0, - byteOffset=byteOffset, - byteLength=len(albedo_blob), - ) - ) - - blob += albedo_blob - byteOffset += len(albedo_blob) - - gltf.buffers[0].byteLength = byteOffset - - # append metllic roughness - if self.metallicRoughness is not None: - metallicRoughness = self.metallicRoughness.detach().cpu().numpy() - metallicRoughness = (metallicRoughness * 255).astype(np.uint8) - metallicRoughness = cv2.cvtColor(metallicRoughness, cv2.COLOR_RGB2BGR) - metallicRoughness_blob = cv2.imencode('.png', metallicRoughness)[1].tobytes() - - # update texture definition - gltf.materials[0].pbrMetallicRoughness.metallicFactor = 1.0 - gltf.materials[0].pbrMetallicRoughness.roughnessFactor = 1.0 - gltf.materials[0].pbrMetallicRoughness.metallicRoughnessTexture = pygltflib.TextureInfo(index=1, texCoord=0) - - gltf.textures.append(pygltflib.Texture(sampler=1, source=1)) - gltf.samplers.append(pygltflib.Sampler(magFilter=pygltflib.LINEAR, minFilter=pygltflib.LINEAR_MIPMAP_LINEAR, wrapS=pygltflib.REPEAT, wrapT=pygltflib.REPEAT)) - gltf.images.append(pygltflib.Image(bufferView=4, mimeType="image/png")) - - # update buffers - gltf.bufferViews.append( - # index = 4, metallicRoughness texture; as none target - pygltflib.BufferView( - buffer=0, - byteOffset=byteOffset, - byteLength=len(metallicRoughness_blob), - ) - ) - - blob += metallicRoughness_blob - byteOffset += len(metallicRoughness_blob) - - gltf.buffers[0].byteLength = byteOffset - - - # set actual data - gltf.set_binary_blob(blob) - - # glb = b"".join(gltf.save_to_bytes()) - gltf.save(path) - - - def write_obj(self, path): - """write the mesh in obj format. Will also write the texture and mtl files. - - Args: - path (str): path to write. - """ - - mtl_path = path.replace(".obj", ".mtl") - albedo_path = path.replace(".obj", "_albedo.png") - metallic_path = path.replace(".obj", "_metallic.png") - roughness_path = path.replace(".obj", "_roughness.png") - - v_np = self.v.detach().cpu().numpy() - vt_np = self.vt.detach().cpu().numpy() if self.vt is not None else None - vn_np = self.vn.detach().cpu().numpy() if self.vn is not None else None - vc_np = self.vc.detach().cpu().numpy() if self.vc is not None else None - f_np = self.f.detach().cpu().numpy() - ft_np = self.ft.detach().cpu().numpy() if self.ft is not None else None - fn_np = self.fn.detach().cpu().numpy() if self.fn is not None else None - - with open(path, "w") as fp: - fp.write(f"mtllib {os.path.basename(mtl_path)} \n") - - if vc_np is None: - for v in v_np: - fp.write(f"v {v[0]} {v[1]} {v[2]} \n") - else: - for v_i in range(v_np.shape[0]): - v = v_np[v_i] - vc = vc_np[v_i] - fp.write(f"v {v[0]} {v[1]} {v[2]} {vc[0]} {vc[1]} {vc[2]}\n") - - if vt_np is not None: - for v in vt_np: - fp.write(f"vt {v[0]} {1 - v[1]} \n") - - if vn_np is not None: - for v in vn_np: - fp.write(f"vn {v[0]} {v[1]} {v[2]} \n") - - fp.write(f"usemtl defaultMat \n") - for i in range(len(f_np)): - fp.write( - f'f {f_np[i, 0] + 1}/{ft_np[i, 0] + 1 if ft_np is not None else ""}/{fn_np[i, 0] + 1 if fn_np is not None else ""} \ - {f_np[i, 1] + 1}/{ft_np[i, 1] + 1 if ft_np is not None else ""}/{fn_np[i, 1] + 1 if fn_np is not None else ""} \ - {f_np[i, 2] + 1}/{ft_np[i, 2] + 1 if ft_np is not None else ""}/{fn_np[i, 2] + 1 if fn_np is not None else ""} \n' - ) - - with open(mtl_path, "w") as fp: - fp.write(f"newmtl defaultMat \n") - fp.write(f"Ka 1 1 1 \n") - fp.write(f"Kd 1 1 1 \n") - fp.write(f"Ks 1 1 1 \n") - #fp.write(f"Tr 1 \n") # will cause the mesh materials in three.js completely transparent - fp.write(f"illum 1 \n") - fp.write(f"Ns 10 \n") - if self.albedo is not None: - fp.write(f"map_Kd {os.path.basename(albedo_path)} \n") - if self.metallicRoughness is not None: - # ref: https://en.wikipedia.org/wiki/Wavefront_.obj_file#Physically-based_Rendering - fp.write(f"map_Pm {os.path.basename(metallic_path)} \n") - fp.write(f"map_Pr {os.path.basename(roughness_path)} \n") - - if self.albedo is not None: - albedo = self.albedo.detach().cpu().numpy() - albedo = (albedo * 255).astype(np.uint8) - cv2.imwrite(albedo_path, cv2.cvtColor(albedo, cv2.COLOR_RGB2BGR)) - - if self.metallicRoughness is not None: - metallicRoughness = self.metallicRoughness.detach().cpu().numpy() - metallicRoughness = (metallicRoughness * 255).astype(np.uint8) - cv2.imwrite(metallic_path, metallicRoughness[..., 2]) - cv2.imwrite(roughness_path, metallicRoughness[..., 1]) - - def convert_to_pointcloud(self): - xyz = self.v.detach().cpu().numpy() - num_pts = self.v.shape[0] - shs = np.random.random((num_pts, 3)) / 255.0 - normals = np.zeros((num_pts, 3)) - pcd = PointCloud(points=xyz, colors=SH2RGB(shs), normals=normals) - return pcd - - -class PointCloud(NamedTuple): - points: np.array - colors: np.array - normals: np.array \ No newline at end of file diff --git a/mesh_processer/mesh_utils.py b/mesh_processer/mesh_utils.py deleted file mode 100644 index c3799df..0000000 --- a/mesh_processer/mesh_utils.py +++ /dev/null @@ -1,681 +0,0 @@ -import torch -import numpy as np -from kornia.geometry.conversions import ( - quaternion_to_axis_angle, - axis_angle_to_quaternion, -) -import pymeshlab as pml -from plyfile import PlyData, PlyElement - -from typing import Optional - -pytorch3d_capable = True -try: - import pytorch3d - from pytorch3d.ops import knn_points -except ImportError: - pytorch3d_capable = False - -from .mesh import PointCloud -from shared_utils.sh_utils import SH2RGB, RGB2SH - -def _base_face_areas(face_vertices_0, face_vertices_1, face_vertices_2): - """Base function to compute the face areas.""" - x1, x2, x3 = torch.split(face_vertices_0 - face_vertices_1, 1, dim=-1) - y1, y2, y3 = torch.split(face_vertices_1 - face_vertices_2, 1, dim=-1) - - a = (x2 * y3 - x3 * y2) ** 2 - b = (x3 * y1 - x1 * y3) ** 2 - c = (x1 * y2 - x2 * y1) ** 2 - areas = torch.sqrt(a + b + c) * 0.5 - - return areas - - -def _base_sample_points_selected_faces(face_vertices, face_features=None): - """Base function to sample points over selected faces, sample one point per face. - The coordinates of the face vertices are interpolated to generate new samples. - - Args: - face_vertices (tuple of torch.Tensor): - Coordinates of vertices, corresponding to selected faces to sample from. - A tuple of 3 entries corresponding to each of the face vertices. - Each entry is a torch.Tensor of shape :math:`(\\text{batch_size}, \\text{num_samples}, 3)`. - face_features (tuple of torch.Tensor, Optional): - Features of face vertices, corresponding to selected faces to sample from. - A tuple of 3 entries corresponding to each of the face vertices. - Each entry is a torch.Tensor of shape - :math:`(\\text{batch_size}, \\text{num_samples}, \\text{feature_dim})`. - - Returns: - (torch.Tensor, torch.Tensor): - Sampled point coordinates of shape :math:`(\\text{batch_size}, \\text{num_samples}, 3)`. - Sampled points interpolated features of shape - :math:`(\\text{batch_size}, \\text{num_samples}, \\text{feature_dim})`. - If `face_vertices_features` arg is not specified, the returned interpolated features are None. - """ - - face_vertices0, face_vertices1, face_vertices2 = face_vertices - - sampling_shape = tuple(int(d) for d in face_vertices0.shape[:-1]) + (1,) - # u is proximity to middle point between v1 and v2 against v0. - # v is proximity to v2 against v1. - # - # The probability density for u should be f_U(u) = 2u. - # However, torch.rand use a uniform (f_X(x) = x) distribution, - # so using torch.sqrt we make a change of variable to have the desired density - # f_Y(y) = f_X(y ^ 2) * |d(y ^ 2) / dy| = 2y - u = torch.sqrt(torch.rand(sampling_shape, - device=face_vertices0.device, - dtype=face_vertices0.dtype)) - - v = torch.rand(sampling_shape, - device=face_vertices0.device, - dtype=face_vertices0.dtype) - w0 = 1 - u - w1 = u * (1 - v) - w2 = u * v - - points = w0 * face_vertices0 + w1 * face_vertices1 + w2 * face_vertices2 - - features = None - if face_features is not None: - face_features0, face_features1, face_features2 = face_features - features = w0 * face_features0 + w1 * face_features1 + \ - w2 * face_features2 - - return points, features - -# Modified from https://github.com/NVIDIAGameWorks/kaolin/blob/master/kaolin/ops/mesh/trianglemesh.py#L159 -def sample_points(vertices, faces, num_samples, areas=None, face_features=None): - r"""Uniformly sample points over the surface of triangle meshes. - - First, face on which the point is sampled is randomly selected, - with the probability of selection being proportional to the area of the face. - then the coordinate on the face is uniformly sampled. - - If ``face_features`` is defined for the mesh faces, - the sampled points will be returned with interpolated features as well, - otherwise, no feature interpolation will occur. - - Args: - vertices (torch.Tensor): - The vertices of the meshes, of shape - :math:`(\text{batch_size}, \text{num_vertices}, 3)`. - faces (torch.LongTensor): - The faces of the mesh, of shape :math:`(\text{num_faces}, 3)`. - num_samples (int): - The number of point sampled per mesh. - Also the number of faces sampled per mesh, and then sample a single point per face. - areas (torch.Tensor, optional): - The areas of each face, of shape :math:`(\text{batch_size}, \text{num_faces})`, - can be preprocessed, for fast on-the-fly sampling, - will be computed if None (default). - face_features (torch.Tensor, optional): - Per-vertex-per-face features, matching ``faces`` order, - of shape :math:`(\text{batch_size}, \text{num_faces}, 3, \text{feature_dim})`. - For example: - - 1. Texture uv coordinates would be of shape - :math:`(\text{batch_size}, \text{num_faces}, 3, 2)`. - 2. RGB color values would be of shape - :math:`(\text{batch_size}, \text{num_faces}, 3, 3)`. - - When specified, it is used to interpolate the features for new sampled points. - - See also: - :func:`~kaolin.ops.mesh.index_vertices_by_faces` for conversion of features defined per vertex - and need to be converted to per-vertex-per-face shape of :math:`(\text{num_faces}, 3)`. - - Returns: - (torch.Tensor, torch.LongTensor, (optional) torch.Tensor): - the pointclouds of shape :math:`(\text{batch_size}, \text{num_samples}, 3)`, - and the indexes of the faces selected, - of shape :math:`(\text{batch_size}, \text{num_samples})`. - - If ``face_features`` arg is specified, then the interpolated features of sampled points of shape - :math:`(\text{batch_size}, \text{num_samples}, \text{feature_dim})` are also returned. - """ - if faces.shape[-1] != 3: - raise NotImplementedError("sample_points is only implemented for triangle meshes") - faces_0, faces_1, faces_2 = torch.split(faces, 1, dim=1) # (num_faces, 3) -> tuple of (num_faces,) - face_v_0 = torch.index_select(vertices, 1, faces_0.reshape(-1)) # (batch_size, num_faces, 3) - face_v_1 = torch.index_select(vertices, 1, faces_1.reshape(-1)) # (batch_size, num_faces, 3) - face_v_2 = torch.index_select(vertices, 1, faces_2.reshape(-1)) # (batch_size, num_faces, 3) - - if areas is None: - areas = _base_face_areas(face_v_0, face_v_1, face_v_2).squeeze(-1) - face_dist = torch.distributions.Categorical(areas) - face_choices = face_dist.sample([num_samples]).transpose(0, 1) - _face_choices = face_choices.unsqueeze(-1).repeat(1, 1, 3) - v0 = torch.gather(face_v_0, 1, _face_choices) # (batch_size, num_samples, 3) - v1 = torch.gather(face_v_1, 1, _face_choices) # (batch_size, num_samples, 3) - v2 = torch.gather(face_v_2, 1, _face_choices) # (batch_size, num_samples, 3) - face_vertices_choices = (v0, v1, v2) - - # UV coordinates are available, make sure to calculate them for sampled points as well - face_features_choices = None - if face_features is not None: - feat_dim = face_features.shape[-1] - # (num_faces, 3) -> tuple of (num_faces,) - _face_choices = face_choices[..., None, None].repeat(1, 1, 3, feat_dim) - face_features_choices = torch.gather(face_features, 1, _face_choices) - face_features_choices = tuple( - tmp_feat.squeeze(2) for tmp_feat in torch.split(face_features_choices, 1, dim=2)) - - points, point_features = _base_sample_points_selected_faces( - face_vertices_choices, face_features_choices) - - if point_features is not None: - return points, face_choices, point_features - else: - return points, face_choices - -def poisson_mesh_reconstruction(points, normals=None): - # points/normals: [N, 3] np.ndarray - - import open3d as o3d - - pcd = o3d.geometry.PointCloud() - pcd.points = o3d.utility.Vector3dVector(points) - - # outlier removal - pcd, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=10) - - # normals - if normals is None: - pcd.estimate_normals() - else: - pcd.normals = o3d.utility.Vector3dVector(normals[ind]) - - # visualize - o3d.visualization.draw_geometries([pcd], point_show_normal=False) - - mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson( - pcd, depth=9 - ) - vertices_to_remove = densities < np.quantile(densities, 0.1) - mesh.remove_vertices_by_mask(vertices_to_remove) - - # visualize - o3d.visualization.draw_geometries([mesh]) - - vertices = np.asarray(mesh.vertices) - triangles = np.asarray(mesh.triangles) - - print( - f"[INFO] poisson mesh reconstruction: {points.shape} --> {vertices.shape} / {triangles.shape}" - ) - - return vertices, triangles - - -def decimate_mesh( - verts, faces, target=5e4, backend="pymeshlab", remesh=False, optimalplacement=True, verbose=True -): - """ perform mesh decimation. - - Args: - verts (np.ndarray): mesh vertices, float [N, 3] - faces (np.ndarray): mesh faces, int [M, 3] - target (int): targeted number of faces - backend (str, optional): algorithm backend, can be "pymeshlab" or "pyfqmr". Defaults to "pymeshlab". - remesh (bool, optional): whether to remesh after decimation. Defaults to False. - optimalplacement (bool, optional): For flat mesh, use False to prevent spikes. Defaults to True. - verbose (bool, optional): whether to print the decimation process. Defaults to True. - - Returns: - Tuple[np.ndarray]: vertices and faces after decimation. - """ - - _ori_vert_shape = verts.shape - _ori_face_shape = faces.shape - - if backend == "pyfqmr": - import pyfqmr - - solver = pyfqmr.Simplify() - solver.setMesh(verts, faces) - solver.simplify_mesh(target_count=target, preserve_border=False, verbose=False) - verts, faces, normals = solver.getMesh() - else: - m = pml.Mesh(verts, faces) - ms = pml.MeshSet() - ms.add_mesh(m, "mesh") # will copy! - - ms.meshing_merge_close_vertices() - - # filters - # ms.meshing_decimation_clustering(threshold=pml.PercentageValue(1)) - ms.meshing_decimation_quadric_edge_collapse( - targetfacenum=int(target), optimalplacement=optimalplacement - ) - - if remesh: - # ms.apply_coord_taubin_smoothing() - ms.meshing_isotropic_explicit_remeshing( - iterations=3, targetlen=pml.PercentageValue(1) - ) - - # extract mesh - m = ms.current_mesh() - m.compact() - verts = m.vertex_matrix() - faces = m.face_matrix() - - if verbose: - print(f"[INFO] mesh decimation: {_ori_vert_shape} --> {verts.shape}, {_ori_face_shape} --> {faces.shape}") - - return verts, faces - - -def clean_mesh( - verts, - faces, - v_pct=1, - min_f=64, - min_d=20, - repair=True, - remesh=True, - remesh_size=0.01, -): - # verts: [N, 3] - # faces: [N, 3] - - _ori_vert_shape = verts.shape - _ori_face_shape = faces.shape - - m = pml.Mesh(verts, faces) - ms = pml.MeshSet() - ms.add_mesh(m, "mesh") # will copy! - - # filters - ms.meshing_remove_unreferenced_vertices() # verts not refed by any faces - - if v_pct > 0: - ms.meshing_merge_close_vertices( - threshold=pml.Percentage(v_pct) - ) # 1/10000 of bounding box diagonal - - ms.meshing_remove_duplicate_faces() # faces defined by the same verts - ms.meshing_remove_null_faces() # faces with area == 0 - - if min_d > 0: - ms.meshing_remove_connected_component_by_diameter( - mincomponentdiag=pml.Percentage(min_d) - ) - - if min_f > 0: - ms.meshing_remove_connected_component_by_face_number(mincomponentsize=min_f) - - if repair: - # ms.meshing_remove_t_vertices(method=0, threshold=40, repeat=True) - ms.meshing_repair_non_manifold_edges(method=0) - ms.meshing_repair_non_manifold_vertices(vertdispratio=0) - - if remesh: - # ms.apply_coord_taubin_smoothing() - ms.meshing_isotropic_explicit_remeshing( - iterations=3, targetlen=pml.AbsoluteValue(remesh_size) - ) - - # extract mesh - m = ms.current_mesh() - verts = m.vertex_matrix() - faces = m.face_matrix() - - print( - f"[INFO] mesh cleaning: {_ori_vert_shape} --> {verts.shape}, {_ori_face_shape} --> {faces.shape}" - ) - - return verts, faces - -def construct_list_of_gs_attributes(features_dc, features_rest, scaling, rotation): - l = ['x', 'y', 'z', 'nx', 'ny', 'nz'] - for i in range(features_dc.shape[1]*features_dc.shape[2]): - l.append('f_dc_{}'.format(i)) - for i in range(features_rest.shape[1]*features_rest.shape[2]): - l.append('f_rest_{}'.format(i)) - l.append('opacity') - for i in range(scaling.shape[1]): - l.append('scale_{}'.format(i)) - for i in range(rotation.shape[1]): - l.append('rot_{}'.format(i)) - return l - -def calculate_max_sh_degree_from_gs_ply(plydata): - extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")] - #assert len(extra_f_names)!=3*(max_sh_degree + 1) ** 2 - 3: - max_sh_degree = int(((len(extra_f_names) + 3) / 3) ** 0.5 - 1) - return max_sh_degree, extra_f_names - -def write_gs_ply(xyz, normals, f_dc, f_rest, opacities, scale, rotation, list_of_attributes): - dtype_full = [(attribute, 'f4') for attribute in list_of_attributes] - - elements = np.empty(xyz.shape[0], dtype=dtype_full) - attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1) - elements[:] = list(map(tuple, attributes)) - el = PlyElement.describe(elements, 'vertex') - return PlyData([el]) - -def read_gs_ply(plydata): - xyz = np.stack((np.asarray(plydata.elements[0]["x"]), - np.asarray(plydata.elements[0]["y"]), - np.asarray(plydata.elements[0]["z"])), axis=1) - opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis] - - features_dc = np.zeros((xyz.shape[0], 3, 1)) - features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"]) - features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"]) - features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"]) - - max_sh_degree, extra_f_names = calculate_max_sh_degree_from_gs_ply(plydata) - - features_extra = np.zeros((xyz.shape[0], len(extra_f_names))) - for idx, attr_name in enumerate(extra_f_names): - features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name]) - # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC) - features_extra = features_extra.reshape((features_extra.shape[0], 3, (max_sh_degree + 1) ** 2 - 1)) - - scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")] - scales = np.zeros((xyz.shape[0], len(scale_names))) - for idx, attr_name in enumerate(scale_names): - scales[:, idx] = np.asarray(plydata.elements[0][attr_name]) - - rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")] - rots = np.zeros((xyz.shape[0], len(rot_names))) - for idx, attr_name in enumerate(rot_names): - rots[:, idx] = np.asarray(plydata.elements[0][attr_name]) - - return xyz, features_dc, features_extra, opacities, scales, rots - -def ply_to_points_cloud(plydata): - xyz, features_dc, features_extra, opacities, scales, rots = read_gs_ply(plydata) - - features_dc = np.transpose(features_dc, (0, 2, 1)) # equivalent of torch.transpose(features_dc, 1, 2) - features_extra = np.transpose(features_extra, (0, 2, 1)) - shs = np.concatenate((features_dc, features_extra), axis=1) - normals = np.zeros_like(xyz) - pcd = PointCloud(points=xyz, colors=SH2RGB(shs), normals=normals) - return pcd - - -def get_target_axis_and_scale(axis_string, scale_value=1.0): - """ - Coordinate system inverts when: - 1. Any of the axis inverts - 2. Two of the axises switch - - If coordinate system inverts twice in a row then it will not be inverted - """ - axis_names = ["x", "y", "z"] - - target_axis, target_scale, coordinate_invert_count = [], [], 0 - axis_switch_count = 0 - for i in range(len(axis_names)): - s = axis_string[i] - if s[0] == "-": - target_scale.append(-scale_value) - coordinate_invert_count += 1 - else: - target_scale.append(scale_value) - - new_axis_i = axis_names.index(s[1]) - if new_axis_i != i: - axis_switch_count += 1 - target_axis.append(new_axis_i) - - if axis_switch_count == 2: - coordinate_invert_count += 1 - - return target_axis, target_scale, coordinate_invert_count - -def switch_vector_axis(vector3s, target_axis): - """ - Example: - vector3s = torch.tensor([[1, 2, 3], [3, 2, 1], [2, 3, 1]]) # shape (N, 3) - - target_axis = (2, 0, 1) # or [2, 0, 1] - vector3s[:, [0, 1, 2]] = vector3s[:, target_axis] - - # Result: tensor([[3, 1, 2], [1, 3, 2], [1, 2, 3]]) - """ - vector3s[:, [0, 1, 2]] = vector3s[:, target_axis] - return vector3s - -def switch_ply_axis_and_scale(plydata, target_axis, target_scale, coordinate_invert_count): - """ - Args: - target_axis (array): shape (3) - target_scale (array): shape (3) - """ - xyz, features_dc, features_extra, opacities, scales, rots = read_gs_ply(plydata) - normals = np.zeros_like(xyz) - features_dc_2d = features_dc.reshape(features_dc.shape[0], features_dc.shape[1]*features_dc.shape[2]) - features_extra_2d = features_extra.reshape(features_extra.shape[0], features_extra.shape[1]*features_extra.shape[2]) - - target_scale = torch.tensor(target_scale).float().cuda() - xyz = switch_vector_axis(torch.tensor(xyz).float().cuda() * target_scale, target_axis).detach().cpu().numpy() - scales = switch_vector_axis(torch.tensor(scales).float().cuda(), target_axis).detach().cpu().numpy() - - # change rotation representation from quaternion (w, x, y, z) to axis angle vector (x, y, z) to make swich axis easier - rots_axis_angle = quaternion_to_axis_angle(torch.tensor(rots).float().cuda()) - rots_axis_angle = switch_vector_axis(rots_axis_angle * target_scale, target_axis) - """ - Since axis–angle vector is composed of axis (unit vector/direction) and clockwise radians angle (vector magnitude), - so in order to invert the sign of angle when coordinate system inverts, we also need to invert the direction of axis–angle vector - """ - if coordinate_invert_count % 2 != 0: - rots_axis_angle = -rots_axis_angle - rots = axis_angle_to_quaternion(rots_axis_angle).detach().cpu().numpy() - - return write_gs_ply(xyz, normals, features_dc_2d, features_extra_2d, opacities, scales, rots, construct_list_of_gs_attributes(features_dc, features_extra, scales, rots)) - -def switch_mesh_axis_and_scale(mesh, target_axis, target_scale, flip_normal=False): - """ - Args: - target_axis (array): shape (3) - target_scale (array): shape (3) - """ - target_scale = torch.tensor(target_scale).float().cuda() - mesh.v = switch_vector_axis(mesh.v * target_scale, target_axis) - mesh.vn = switch_vector_axis(mesh.vn * target_scale, target_axis) - if flip_normal: - mesh.vn *= -1 - return mesh - - -def marching_cubes_density_to_mesh(get_density_func, grid_size=256, S=128, density_thresh=10, decimate_target=5e4): - from mcubes import marching_cubes - from kiui.mesh_utils import clean_mesh, decimate_mesh - - - sigmas = np.zeros([grid_size, grid_size, grid_size], dtype=np.float32) - - X = torch.linspace(-1, 1, grid_size).split(S) - Y = torch.linspace(-1, 1, grid_size).split(S) - Z = torch.linspace(-1, 1, grid_size).split(S) - - for xi, xs in enumerate(X): - for yi, ys in enumerate(Y): - for zi, zs in enumerate(Z): - xx, yy, zz = torch.meshgrid(xs, ys, zs, indexing='ij') - pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1) # [S, 3] - val = get_density_func(pts) - sigmas[xi * S: xi * S + len(xs), yi * S: yi * S + len(ys), zi * S: zi * S + len(zs)] = val.reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy() # [S, 1] --> [x, y, z] - - print(f'[INFO] marching cubes thresh: {density_thresh} ({sigmas.min()} ~ {sigmas.max()})') - - vertices, triangles = marching_cubes(sigmas, density_thresh) - vertices = vertices / (grid_size - 1.0) * 2 - 1 - - # clean - vertices = vertices.astype(np.float32) - triangles = triangles.astype(np.int32) - vertices, triangles = clean_mesh(vertices, triangles, remesh=True, remesh_size=0.01) - if triangles.shape[0] > decimate_target: - vertices, triangles = decimate_mesh(vertices, triangles, decimate_target, optimalplacement=False) - - return vertices, triangles - -def color_func_to_albedo(mesh, get_rgb_func, texture_resolution=1024, padding=2, batch_size=640000, device="cuda", force_cuda_rast=False, glctx=None): - import nvdiffrast.torch as dr - from kiui.op import uv_padding - - if glctx is None: - if force_cuda_rast: - glctx = dr.RasterizeCudaContext() - else: - glctx = dr.RasterizeGLContext() - - # render uv maps - h = w = texture_resolution - if mesh.vt is None: - mesh.auto_uv() - - uv = mesh.vt * 2.0 - 1.0 # uvs to range [-1, 1] - uv = torch.cat((uv, torch.zeros_like(uv[..., :1]), torch.ones_like(uv[..., :1])), dim=-1) # [N, 4] - - rast, _ = dr.rasterize(glctx, uv.unsqueeze(0), mesh.ft, (h, w)) # [1, h, w, 4] - xyzs, _ = dr.interpolate(mesh.v.unsqueeze(0), rast, mesh.f) # [1, h, w, 3] - mask, _ = dr.interpolate(torch.ones_like(mesh.v[:, :1]).unsqueeze(0), rast, mesh.f) # [1, h, w, 1] - - # masked query - xyzs = xyzs.view(-1, 3) - mask = (mask > 0).view(-1) - - albedo = torch.zeros(h * w, 3, device=device, dtype=torch.float32) - - if mask.any(): - print(f"[INFO] querying texture...") - - xyzs = xyzs[mask] # [M, 3] - - # batched inference to avoid OOM - batch = [] - head = 0 - while head < xyzs.shape[0]: - tail = min(head + batch_size, xyzs.shape[0]) - batch.append(get_rgb_func(xyzs[head:tail]).float()) - head += batch_size - - albedo[mask] = torch.cat(batch, dim=0) - - albedo = albedo.view(h, w, -1) - mask = mask.view(h, w) - albedo = uv_padding(albedo, mask, padding) - - return albedo - -@torch.no_grad() -def K_nearest_neighbors_func( - points: torch.Tensor, - K: int, - query: Optional[torch.Tensor] = None, - return_dist=False, - skip_first_index=True, -): - if not pytorch3d_capable: - raise ImportError("pytorch3d is not installed, which is required for KNN") - - # query/points: Tensor of shape (N, P1/P2, D) giving a batch of N point clouds, each containing up to P1/P2 points of dimension D - if query is None: - query = points - dist, idx, nn = knn_points(query[None, ...], points[None, ...], K=K, return_nn=True) - - # idx: Tensor of shape (N, P1, K) - # nn: Tensor of shape (N, P1, K, D) - - # take the index 1 since index 0 is the point itself - if skip_first_index: - if not return_dist: - return nn[0, :, 1:, :], idx[0, :, 1:] - else: - return nn[0, :, 1:, :], idx[0, :, 1:], dist[0, :, 1:] - else: - if not return_dist: - return nn[0, :, :, :], idx[0, :, :] - else: - return nn[0, :, :, :], idx[0, :, :], dist[0, :, :] - -def interpolate_texture_map_attr(mesh, texture_size: int = 256, batch_size: int = 64, interpolate_color=True, interpolate_position=False): - # Get UV coordinates and faces - if mesh.vt is None: - mesh.auto_uv() - - # Get Faces on UV - texture_size_minus_one = texture_size - 1 - faces = mesh.ft - verts_uvs = mesh.vt * texture_size_minus_one - verts_uvs_idx = verts_uvs.to(torch.long) - - vmapping = mesh.get_default_vt_to_v_mapping() - # Get vertex colors - interpolate_color = interpolate_color and mesh.vc is not None - texture_map = mesh.albedo - if interpolate_color: - verts_colors = mesh.vc[vmapping] - # Create a blank texture map - texture_map = torch.zeros((texture_size, texture_size, 3), device=verts_colors.device) - - # Get vertex positions - position_map = None - if interpolate_position: - verts_positions = mesh.v[vmapping] - position_map = torch.zeros((texture_size, texture_size, 3), device=verts_colors.device) - - # Create a grid of UV coordinates - grid_x, grid_y = torch.meshgrid(torch.linspace(0, texture_size_minus_one, texture_size, dtype=verts_uvs.dtype, device=verts_colors.device), torch.linspace(0, texture_size_minus_one, texture_size, dtype=verts_uvs.dtype, device=verts_colors.device)) - grid = torch.stack([grid_x, grid_y], dim=-1) - - batch_idx_range = torch.arange(0, batch_size**2, 1, device=verts_colors.device, dtype=torch.long) - - all_face_verts = verts_uvs_idx[faces] - for uv_i in range(0, texture_size_minus_one, batch_size): - for uv_j in range(0, texture_size_minus_one, batch_size): - # Divide faces into squre batchs - faces_mask = (uv_i <= all_face_verts[:, :, 0]) & (all_face_verts[:, :, 0] <= uv_i + batch_size) & (uv_j <= all_face_verts[:, :, 1]) & (all_face_verts[:, :, 1] <= uv_j + batch_size) - faces_mask, _ = torch.max(faces_mask, dim=1) - faces_mask = faces_mask.nonzero().squeeze(-1) - batch_faces = faces[faces_mask] - - uvs = verts_uvs[batch_faces] - if interpolate_color: - colors = verts_colors[batch_faces] - if interpolate_position: - positions = verts_positions[batch_faces] - - v0, v1, v2 = uvs[:, 0], uvs[:, 1], uvs[:, 2] - v0_0, v0_1 = v0[:, 0], v0[:, 1] - v1_0, v1_1 = v1[:, 0], v1[:, 1] - v2_0, v2_1 = v2[:, 0], v2[:, 1] - - sub_grid = grid[uv_i:uv_i + batch_size, uv_j:uv_j + batch_size].reshape(-1, 2).unsqueeze(1) - - # Compute barycentric coordinates for the batch of faces - denom = (v1_1 - v2_1) * (v0_0 - v2_0) + (v2_0 - v1_0) * (v0_1 - v2_1) - a = ((v1_1 - v2_1) * (sub_grid[:, :, 0] - v2_0) + (v2_0 - v1_0) * (sub_grid[:, :, 1] - v2_1)) / denom - b = ((v2_1 - v0_1) * (sub_grid[:, :, 0] - v2_0) + (v0_0 - v2_0) * (sub_grid[:, :, 1] - v2_1)) / denom - c = 1 - a - b - - # Mask for points inside the triangle - mask = (a >= 0) & (b >= 0) & (c >= 0) - mask, mask_idx = torch.max(mask, dim=1) - - # Get texture coordinates - sub_grid = sub_grid.squeeze(1) - uv_coords = sub_grid[mask].to(torch.long) - - # Interpolate colors - a, b, c = a[batch_idx_range, mask_idx][mask], b[batch_idx_range, mask_idx][mask], c[batch_idx_range, mask_idx][mask] - - if interpolate_color: - colors = colors[mask_idx][mask] - interpolated_colors = (a[:, None] * colors[:, 0] + (b[:, None] * colors[:, 1]) + (c[:, None] * colors[:, 2])) - texture_map[uv_coords[:, 1], uv_coords[:, 0]] = interpolated_colors - if interpolate_position: - positions = positions[mask_idx][mask] - interpolated_positions = (a[:, None] * positions[:, 0] + (b[:, None] * positions[:, 1]) + (c[:, None] * positions[:, 2])) - position_map[uv_coords[:, 1], uv_coords[:, 0]] = interpolated_positions - - return texture_map, position_map \ No newline at end of file