renamed mesh_processor
This commit is contained in:
DenisKochetov
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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))
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
+3
-1
@@ -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__
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user