test new uv load
This commit is contained in:
DenisKochetov
Binary file not shown.
+269
@@ -0,0 +1,269 @@
|
||||
class Preview_3DMesh:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"mesh_file_path": ("STRING", {"default": '', "multiline": False}),
|
||||
},
|
||||
}
|
||||
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "preview_mesh"
|
||||
CATEGORY = "Comfy3D/Visualize"
|
||||
|
||||
def preview_mesh(self, mesh_file_path):
|
||||
|
||||
mesh_folder_path, filename = os.path.split(mesh_file_path)
|
||||
|
||||
if not os.path.isabs(mesh_file_path):
|
||||
mesh_file_path = os.path.join(comfy_paths.output_directory, mesh_folder_path, filename)
|
||||
|
||||
if not filename.lower().endswith(SUPPORTED_3D_EXTENSIONS):
|
||||
cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3D file extensions: {SUPPORTED_3D_EXTENSIONS}").error.print()
|
||||
mesh_file_path = ""
|
||||
|
||||
print(f"[Preview_3DMesh] Final mesh path: {mesh_file_path}")
|
||||
print(f"[Preview_3DMesh] File exists: {os.path.exists(mesh_file_path) if mesh_file_path else False}")
|
||||
|
||||
previews = [
|
||||
{
|
||||
"filepath": mesh_file_path,
|
||||
}
|
||||
]
|
||||
return {"ui": {"previews": previews}, "result": ()}
|
||||
|
||||
class Load_3D_Mesh:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"mesh_file_path": ("STRING", {"default": '', "multiline": False}),
|
||||
"use_fastmesh": ("BOOLEAN", {"default": True, "tooltip": "Использовать FastMesh для GLB/GLTF файлов"}),
|
||||
"resize": ("BOOLEAN", {"default": False},),
|
||||
"renormal": ("BOOLEAN", {"default": True},),
|
||||
"retex": ("BOOLEAN", {"default": False},),
|
||||
"optimizable": ("BOOLEAN", {"default": False},),
|
||||
"clean": ("BOOLEAN", {"default": False},),
|
||||
"resize_bound": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1000.0, "step": 0.001}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"MESH",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"mesh",
|
||||
)
|
||||
FUNCTION = "load_mesh"
|
||||
CATEGORY = "Comfy3D/Import|Export"
|
||||
|
||||
def load_mesh(self, mesh_file_path, use_fastmesh, resize, renormal, retex, optimizable, clean, resize_bound):
|
||||
import time
|
||||
mesh = None
|
||||
|
||||
if not os.path.isabs(mesh_file_path):
|
||||
mesh_file_path = os.path.join(comfy_paths.input_directory, mesh_file_path)
|
||||
|
||||
if os.path.exists(mesh_file_path):
|
||||
folder, filename = os.path.split(mesh_file_path)
|
||||
if filename.lower().endswith(SUPPORTED_3D_EXTENSIONS):
|
||||
with torch.inference_mode(not optimizable):
|
||||
start_time = time.time()
|
||||
if use_fastmesh and mesh_file_path.endswith((".glb", ".gltf")):
|
||||
mesh = FastMesh.load(mesh_file_path, resize, renormal, retex, clean, resize_bound)
|
||||
print(f"[FastMesh] Loading completed in {time.time() - start_time:.3f}s")
|
||||
else:
|
||||
mesh = Mesh.load(mesh_file_path, resize, renormal, retex, clean, resize_bound)
|
||||
print(f"[Mesh] Loading completed in {time.time() - start_time:.3f}s")
|
||||
else:
|
||||
cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3D file extensions: {SUPPORTED_3D_EXTENSIONS}").error.print()
|
||||
else:
|
||||
cstr(f"[{self.__class__.__name__}] File {mesh_file_path} does not exist").error.print()
|
||||
return (mesh, )
|
||||
|
||||
class Save_3D_Mesh:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"mesh": ("MESH",),
|
||||
"save_path": ("STRING", {"default": 'Mesh_%Y-%m-%d-%M-%S-%f.glb', "multiline": False}),
|
||||
"use_fastmesh": ("BOOLEAN", {"default": True, "tooltip": "Использовать FastMesh для быстрого сохранения GLB/PLY/OBJ"}),
|
||||
},
|
||||
}
|
||||
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = (
|
||||
"STRING",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"save_path",
|
||||
)
|
||||
FUNCTION = "save_mesh"
|
||||
CATEGORY = "Comfy3D/Import|Export"
|
||||
|
||||
def save_mesh(self, mesh, save_path, use_fastmesh):
|
||||
import time
|
||||
save_path = parse_save_filename(save_path, comfy_paths.output_directory, SUPPORTED_3D_EXTENSIONS, self.__class__.__name__)
|
||||
|
||||
if save_path is not None:
|
||||
start_time = time.time()
|
||||
|
||||
# Используем FastMesh методы для быстрого сохранения
|
||||
if use_fastmesh and save_path.endswith((".glb", ".ply", ".obj")):
|
||||
# Если это уже FastMesh - используем напрямую
|
||||
print(f"[Save_3D_Mesh] FastMesh detected")
|
||||
if hasattr(mesh, '_write_glb_fast'):
|
||||
mesh.write(save_path)
|
||||
save_time = time.time() - start_time
|
||||
print(f"[FastMesh] Saving completed in {save_time:.3f}s")
|
||||
else:
|
||||
# Если обычный Mesh - создаем временный FastMesh для быстрых методов
|
||||
print(f"[Save_3D_Mesh] REGULAR Mesh detected")
|
||||
fast_mesh = FastMesh(
|
||||
v=mesh.v, f=mesh.f, vn=mesh.vn, fn=mesh.fn,
|
||||
vt=mesh.vt, ft=mesh.ft, vc=mesh.vc,
|
||||
albedo=mesh.albedo, metallicRoughness=getattr(mesh, 'metallicRoughness', None),
|
||||
device=mesh.device
|
||||
)
|
||||
fast_mesh.ori_center = getattr(mesh, 'ori_center', 0)
|
||||
fast_mesh.ori_scale = getattr(mesh, 'ori_scale', 1)
|
||||
fast_mesh.write(save_path)
|
||||
save_time = time.time() - start_time
|
||||
print(f"[FastMesh] Saving completed in {save_time:.3f}s")
|
||||
else:
|
||||
# Обычное сохранение через стандартный метод
|
||||
mesh.write(save_path)
|
||||
save_time = time.time() - start_time
|
||||
print(f"[Mesh] Saving completed in {save_time:.3f}s")
|
||||
|
||||
return (save_path, )
|
||||
|
||||
|
||||
class Preview_3D_FastMesh:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"mesh_file_path": ("STRING", {"default": '', "multiline": False}),
|
||||
},
|
||||
}
|
||||
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "preview_mesh"
|
||||
CATEGORY = "Comfy3D/Visualize"
|
||||
|
||||
def preview_mesh(self, mesh_file_path):
|
||||
|
||||
mesh_folder_path, filename = os.path.split(mesh_file_path)
|
||||
|
||||
if not os.path.isabs(mesh_file_path):
|
||||
mesh_file_path = os.path.join(comfy_paths.output_directory, mesh_folder_path, filename)
|
||||
|
||||
if not filename.lower().endswith(SUPPORTED_3D_EXTENSIONS):
|
||||
cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3D file extensions: {SUPPORTED_3D_EXTENSIONS}").error.print()
|
||||
mesh_file_path = ""
|
||||
|
||||
print(f"[Preview_3D_FastMesh] Final mesh path: {mesh_file_path}")
|
||||
print(f"[Preview_3D_FastMesh] File exists: {os.path.exists(mesh_file_path) if mesh_file_path else False}")
|
||||
|
||||
previews = [
|
||||
{
|
||||
"filepath": mesh_file_path,
|
||||
}
|
||||
]
|
||||
return {"ui": {"previews": previews}, "result": ()}
|
||||
|
||||
|
||||
class Load_3D_FastMesh:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"mesh_file_path": ("STRING", {"default": '', "multiline": False}),
|
||||
"use_fastmesh": ("BOOLEAN", {"default": True, "tooltip": "Использовать FastMesh загрузчик"}),
|
||||
"resize": ("BOOLEAN", {"default": False},),
|
||||
"renormal": ("BOOLEAN", {"default": True},),
|
||||
"retex": ("BOOLEAN", {"default": False},),
|
||||
"optimizable": ("BOOLEAN", {"default": False},),
|
||||
"clean": ("BOOLEAN", {"default": False},),
|
||||
"resize_bound": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1000.0, "step": 0.001}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"MESH",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"mesh",
|
||||
)
|
||||
FUNCTION = "load_mesh"
|
||||
CATEGORY = "Comfy3D/Import|Export"
|
||||
|
||||
def load_mesh(self, mesh_file_path, use_fastmesh, resize, renormal, retex, optimizable, clean, resize_bound):
|
||||
import time
|
||||
mesh = None
|
||||
|
||||
if not os.path.isabs(mesh_file_path):
|
||||
mesh_file_path = os.path.join(comfy_paths.input_directory, mesh_file_path)
|
||||
|
||||
if os.path.exists(mesh_file_path):
|
||||
folder, filename = os.path.split(mesh_file_path)
|
||||
if filename.lower().endswith(SUPPORTED_3D_EXTENSIONS):
|
||||
with torch.inference_mode(not optimizable):
|
||||
start_time = time.time()
|
||||
if use_fastmesh and mesh_file_path.endswith((".glb", ".gltf")):
|
||||
from mesh_processor.mesh import FastMesh
|
||||
mesh = FastMesh.load(mesh_file_path, resize, renormal, retex, clean, resize_bound)
|
||||
print(f"[FastMesh] Loading completed in {time.time() - start_time:.3f}s")
|
||||
else:
|
||||
mesh = Mesh.load(mesh_file_path, resize, renormal, retex, clean, resize_bound)
|
||||
mesh_type = "FastMesh" if use_fastmesh else "Mesh"
|
||||
print(f"[{mesh_type}] Loading completed in {time.time() - start_time:.3f}s")
|
||||
else:
|
||||
cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3D file extensions: {SUPPORTED_3D_EXTENSIONS}").error.print()
|
||||
else:
|
||||
cstr(f"[{self.__class__.__name__}] File {mesh_file_path} does not exist").error.print()
|
||||
return (mesh, )
|
||||
|
||||
|
||||
class Save_3D_FastMesh:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"mesh": ("MESH",),
|
||||
"save_path": ("STRING", {"default": 'FastMesh_%Y-%m-%d-%M-%S-%f.glb', "multiline": False}),
|
||||
},
|
||||
}
|
||||
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = (
|
||||
"STRING",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"save_path",
|
||||
)
|
||||
FUNCTION = "save_mesh"
|
||||
CATEGORY = "Comfy3D/Import|Export"
|
||||
|
||||
def save_mesh(self, mesh, save_path):
|
||||
import time
|
||||
save_path = parse_save_filename(save_path, comfy_paths.output_directory, SUPPORTED_3D_EXTENSIONS, self.__class__.__name__)
|
||||
|
||||
if save_path is not None:
|
||||
start_time = time.time()
|
||||
mesh.write(save_path)
|
||||
save_time = time.time() - start_time
|
||||
mesh_type = "FastMesh" if hasattr(mesh, '_create_empty_albedo') else "Mesh"
|
||||
print(f"[{mesh_type}] Saving completed in {save_time:.3f}s")
|
||||
|
||||
return (save_path, )
|
||||
@@ -34,9 +34,9 @@ from .rigging_ops import (
|
||||
from .validate import validate_before_save
|
||||
from .mesh import (
|
||||
Mesh,
|
||||
FastMesh,
|
||||
PointCloud,
|
||||
)
|
||||
from .fastmesh import FastMesh
|
||||
|
||||
__all__ = [
|
||||
"GltfDocument",
|
||||
|
||||
@@ -5,7 +5,8 @@ Direct functions WITHOUT trimesh dependency
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
from .mesh import Mesh, FastMesh
|
||||
from .mesh import Mesh
|
||||
from .fastmesh import FastMesh
|
||||
|
||||
|
||||
def export_to_fastmesh(mesh_output, device=None):
|
||||
|
||||
@@ -0,0 +1,629 @@
|
||||
import os
|
||||
import cv2
|
||||
from typing import NamedTuple, Optional
|
||||
import torch
|
||||
from torch import Tensor
|
||||
import trimesh
|
||||
import numpy as np
|
||||
|
||||
from kiui.op import safe_normalize, dot
|
||||
from kiui.typing import *
|
||||
|
||||
from .io_gltf import load_gltf_or_glb, get_binary_data, save_glb, GltfDocument
|
||||
from .mesh_ops import get_all_meshes_triangles
|
||||
from .accessors import access_data, append_accessor_and_bufferview
|
||||
from .mesh import Mesh
|
||||
|
||||
class FastMesh:
|
||||
"""
|
||||
Fast implementation of mesh based on FastGLB without dependency on trimesh.
|
||||
Uses direct loading of GLB/GLTF through mesh_processor.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
v: Optional[Tensor] = None,
|
||||
f: Optional[Tensor] = None,
|
||||
vn: Optional[Tensor] = None,
|
||||
fn: Optional[Tensor] = None,
|
||||
vt: Optional[Tensor] = None,
|
||||
ft: Optional[Tensor] = None,
|
||||
vc: Optional[Tensor] = None,
|
||||
albedo: Optional[Tensor] = None,
|
||||
metallicRoughness: Optional[Tensor] = None,
|
||||
device: Optional[torch.device] = None,
|
||||
):
|
||||
"""Initialize FastMesh with the same parameters as Mesh"""
|
||||
self.device = device if device is not None else torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
self.v = v
|
||||
self.vn = vn
|
||||
self.vt = vt
|
||||
self.f = f
|
||||
self.fn = fn
|
||||
self.ft = ft
|
||||
self.vc = vc
|
||||
self.albedo = albedo
|
||||
self.metallicRoughness = metallicRoughness
|
||||
self.ori_center = 0
|
||||
self.ori_scale = 1
|
||||
self._lazy_texture_doc = None
|
||||
|
||||
@classmethod
|
||||
def load(cls, path, resize=True, renormal=True, retex=False, clean=False, bound=0.5, front_dir='+z', lazy_texture=True, **kwargs):
|
||||
"""Loading mesh through FastGLB approach"""
|
||||
|
||||
device = kwargs.get('device')
|
||||
if device is None:
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
if path.endswith(".obj"):
|
||||
mesh = cls._load_obj_fast(path, device=device, **kwargs)
|
||||
elif path.endswith((".glb", ".gltf")):
|
||||
mesh = cls._load_gltf_fast(path, device=device, lazy_texture=lazy_texture, **kwargs)
|
||||
else:
|
||||
standard_mesh = Mesh.load(path, resize=False, renormal=False, retex=False, clean=False, **kwargs)
|
||||
mesh = cls._from_standard_mesh(standard_mesh)
|
||||
|
||||
if mesh is None:
|
||||
return None
|
||||
|
||||
if clean:
|
||||
mesh._clean_mesh()
|
||||
|
||||
print(f"[FastMesh loading] v: {mesh.v.shape}, f: {mesh.f.shape}")
|
||||
|
||||
if resize:
|
||||
mesh.auto_size(bound=bound)
|
||||
|
||||
if renormal or mesh.vn is None:
|
||||
mesh.auto_normal()
|
||||
print(f"[FastMesh loading] vn: {mesh.vn.shape}, fn: {mesh.fn.shape}")
|
||||
|
||||
print(f"[FastMesh] UV check: retex={retex}, albedo={mesh.albedo is not None}, vt={mesh.vt is not None}")
|
||||
if retex or (mesh.albedo is not None and mesh.vt is None):
|
||||
print(f"[FastMesh] Starting UV generation for {mesh.v.shape[0]} vertices...")
|
||||
mesh.auto_uv(cache_path=path)
|
||||
print(f"[FastMesh loading] vt: {mesh.vt.shape}, ft: {mesh.ft.shape}")
|
||||
|
||||
print(f"[FastMesh] Checking front_dir: {front_dir}")
|
||||
if front_dir != "+z":
|
||||
print(f"[FastMesh] Applying front_dir rotation...")
|
||||
mesh._apply_front_dir_rotation(front_dir)
|
||||
|
||||
print(f"[FastMesh] Loading completed successfully!")
|
||||
|
||||
return mesh
|
||||
|
||||
@classmethod
|
||||
def _load_gltf_fast(cls, path, device=None, lazy_texture=True, **kwargs):
|
||||
"""Fast loading of GLB/GLTF through mesh_processor"""
|
||||
|
||||
try:
|
||||
print(f"[FastMesh] Loading GLB/GLTF: {path}")
|
||||
|
||||
# Load GLTF document
|
||||
|
||||
doc = load_gltf_or_glb(path)
|
||||
print(f"[FastMesh] GLTF document: {len(doc.meshes())} meshes")
|
||||
|
||||
# Extract geometry
|
||||
vertices, faces, mesh_groups = get_all_meshes_triangles(doc, transform_to_global=True)
|
||||
|
||||
if len(vertices) == 0 or len(faces) == 0:
|
||||
print(f"[FastMesh] No geometry in {path}")
|
||||
return None
|
||||
|
||||
mesh = cls(device=device)
|
||||
mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device)
|
||||
mesh.f = torch.tensor(faces, dtype=torch.int32, device=device)
|
||||
|
||||
# Extract UV, normals, colors from all meshes
|
||||
uvs, normals, colors = cls._extract_all_mesh_attributes(doc, mesh_groups, device)
|
||||
if uvs is not None:
|
||||
mesh.vt = uvs
|
||||
mesh.ft = mesh.f # Same topology
|
||||
print(f"[FastMesh] UV: {mesh.vt.shape}")
|
||||
if normals is not None:
|
||||
mesh.vn = normals
|
||||
mesh.fn = mesh.f
|
||||
print(f"[FastMesh] Normals: {mesh.vn.shape}")
|
||||
if colors is not None:
|
||||
mesh.vc = colors
|
||||
print(f"[FastMesh] Colors: {mesh.vc.shape}")
|
||||
if not lazy_texture:
|
||||
cls._extract_gltf_textures(mesh, doc, device)
|
||||
else:
|
||||
mesh._lazy_texture_doc = doc
|
||||
|
||||
print(f"[FastMesh] Loaded: {len(vertices)} vertices, {len(faces)} faces")
|
||||
return mesh
|
||||
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Error loading GLB/GLTF {path}: {e}")
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _extract_all_mesh_attributes(cls, doc, mesh_groups, device):
|
||||
"""Extract UV, normals, colors from all meshes in GLTF"""
|
||||
|
||||
all_uvs = []
|
||||
all_normals = []
|
||||
all_colors = []
|
||||
|
||||
for mesh_idx, gltf_mesh in enumerate(doc.meshes()):
|
||||
for primitive_idx, primitive in enumerate(gltf_mesh.get("primitives", [])):
|
||||
attributes = primitive.get("attributes", {})
|
||||
|
||||
# Find corresponding mesh group
|
||||
mesh_group = None
|
||||
for group in mesh_groups:
|
||||
if group.get('mesh_idx') == mesh_idx and group.get('primitive_idx') == primitive_idx:
|
||||
mesh_group = group
|
||||
break
|
||||
|
||||
if mesh_group is None:
|
||||
continue
|
||||
|
||||
vertex_offset = mesh_group['vertex_offset']
|
||||
vertex_count = mesh_group['vertex_count']
|
||||
|
||||
# Extract UV coordinates
|
||||
if "TEXCOORD_0" in attributes:
|
||||
try:
|
||||
uv_accessor_idx = int(attributes["TEXCOORD_0"])
|
||||
uv_data = access_data(doc, uv_accessor_idx).astype(np.float32)
|
||||
uv_data[:, 1] = 1.0 - uv_data[:, 1] # Flip V
|
||||
all_uvs.append(uv_data)
|
||||
print(f"[FastMesh] Mesh {mesh_idx}.{primitive_idx} UV: {uv_data.shape}")
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Error extracting UV from mesh {mesh_idx}.{primitive_idx}: {e}")
|
||||
all_uvs.append(None)
|
||||
else:
|
||||
all_uvs.append(None)
|
||||
|
||||
# Extract normals
|
||||
if "NORMAL" in attributes:
|
||||
try:
|
||||
normal_accessor_idx = int(attributes["NORMAL"])
|
||||
normal_data = access_data(doc, normal_accessor_idx).astype(np.float32)
|
||||
all_normals.append(normal_data)
|
||||
print(f"[FastMesh] Mesh {mesh_idx}.{primitive_idx} Normals: {normal_data.shape}")
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Error extracting normals from mesh {mesh_idx}.{primitive_idx}: {e}")
|
||||
all_normals.append(None)
|
||||
else:
|
||||
all_normals.append(None)
|
||||
|
||||
# Extract vertex colors
|
||||
if "COLOR_0" in attributes:
|
||||
try:
|
||||
color_accessor_idx = int(attributes["COLOR_0"])
|
||||
color_data = access_data(doc, color_accessor_idx).astype(np.float32)
|
||||
if color_data.shape[1] > 3:
|
||||
color_data = color_data[:, :3]
|
||||
all_colors.append(color_data)
|
||||
print(f"[FastMesh] Mesh {mesh_idx}.{primitive_idx} Colors: {color_data.shape}")
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Error extracting colors from mesh {mesh_idx}.{primitive_idx}: {e}")
|
||||
all_colors.append(None)
|
||||
else:
|
||||
all_colors.append(None)
|
||||
|
||||
# Concatenate all data
|
||||
final_uvs = None
|
||||
final_normals = None
|
||||
final_colors = None
|
||||
|
||||
if any(uv is not None for uv in all_uvs):
|
||||
# Pad missing UV data with zeros
|
||||
padded_uvs = []
|
||||
for i, uv in enumerate(all_uvs):
|
||||
if uv is not None:
|
||||
padded_uvs.append(uv)
|
||||
else:
|
||||
# Create dummy UV coordinates
|
||||
vertex_count = mesh_groups[i]['vertex_count'] if i < len(mesh_groups) else 1000
|
||||
dummy_uv = np.zeros((vertex_count, 2), dtype=np.float32)
|
||||
padded_uvs.append(dummy_uv)
|
||||
|
||||
if padded_uvs:
|
||||
final_uvs = torch.tensor(np.concatenate(padded_uvs, axis=0), dtype=torch.float32, device=device)
|
||||
|
||||
if any(normal is not None for normal in all_normals):
|
||||
valid_normals = [n for n in all_normals if n is not None]
|
||||
if valid_normals:
|
||||
final_normals = torch.tensor(np.concatenate(valid_normals, axis=0), dtype=torch.float32, device=device)
|
||||
|
||||
if any(color is not None for color in all_colors):
|
||||
valid_colors = [c for c in all_colors if c is not None]
|
||||
if valid_colors:
|
||||
final_colors = torch.tensor(np.concatenate(valid_colors, axis=0), dtype=torch.float32, device=device)
|
||||
|
||||
return final_uvs, final_normals, final_colors
|
||||
|
||||
@classmethod
|
||||
def _extract_gltf_attributes(cls, mesh, doc, device):
|
||||
"""Extract attributes from GLTF (UV, normals, colors)"""
|
||||
|
||||
if not doc.meshes():
|
||||
return
|
||||
|
||||
first_mesh = doc.meshes()[0]
|
||||
for primitive in first_mesh.get("primitives", []):
|
||||
attributes = primitive.get("attributes", {})
|
||||
|
||||
# UV coordinates
|
||||
if "TEXCOORD_0" in attributes and mesh.vt is None:
|
||||
try:
|
||||
uv_accessor_idx = int(attributes["TEXCOORD_0"])
|
||||
uv_data = access_data(doc, uv_accessor_idx).astype(np.float32)
|
||||
uv_data[:, 1] = 1.0 - uv_data[:, 1] # Flip V
|
||||
mesh.vt = torch.tensor(uv_data, dtype=torch.float32, device=device)
|
||||
mesh.ft = mesh.f
|
||||
print(f"[FastMesh] UV: {mesh.vt.shape}")
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Error UV: {e}")
|
||||
|
||||
# Normals
|
||||
if "NORMAL" in attributes and mesh.vn is None:
|
||||
try:
|
||||
normal_accessor_idx = int(attributes["NORMAL"])
|
||||
normal_data = access_data(doc, normal_accessor_idx).astype(np.float32)
|
||||
mesh.vn = torch.tensor(normal_data, dtype=torch.float32, device=device)
|
||||
mesh.fn = mesh.f
|
||||
print(f"[FastMesh] Normals: {mesh.vn.shape}")
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Error normals: {e}")
|
||||
|
||||
# Vertex colors
|
||||
if "COLOR_0" in attributes and mesh.vc is None:
|
||||
try:
|
||||
color_accessor_idx = int(attributes["COLOR_0"])
|
||||
color_data = access_data(doc, color_accessor_idx).astype(np.float32)
|
||||
if color_data.shape[1] > 3:
|
||||
color_data = color_data[:, :3]
|
||||
mesh.vc = torch.tensor(color_data, dtype=torch.float32, device=device)
|
||||
print(f"[FastMesh] Colors: {mesh.vc.shape}")
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Error colors: {e}")
|
||||
|
||||
@classmethod
|
||||
def _extract_gltf_textures(cls, mesh, doc, device):
|
||||
"""Extract textures from GLTF"""
|
||||
|
||||
should_create_empty = True
|
||||
|
||||
try:
|
||||
if not (doc.materials() and doc.textures() and doc.images()):
|
||||
mesh._create_empty_albedo()
|
||||
return
|
||||
|
||||
material = doc.materials()[0]
|
||||
pbr = material.get("pbrMetallicRoughness", {})
|
||||
|
||||
# Albedo texture
|
||||
base_color_texture = pbr.get("baseColorTexture")
|
||||
if base_color_texture:
|
||||
albedo_texture = cls._extract_texture_by_index(doc, base_color_texture.get("index", 0))
|
||||
if albedo_texture is not None:
|
||||
albedo_float = albedo_texture.astype(np.float32) / 255.0
|
||||
mesh.albedo = torch.tensor(albedo_float, dtype=torch.float32, device=device).contiguous()
|
||||
should_create_empty = False
|
||||
print(f"[FastMesh] Albedo: {mesh.albedo.shape}")
|
||||
|
||||
# Metallic-Roughness texture
|
||||
mr_texture_info = pbr.get("metallicRoughnessTexture")
|
||||
if mr_texture_info and not should_create_empty:
|
||||
mr_texture = cls._extract_texture_by_index(doc, mr_texture_info.get("index", 0))
|
||||
if mr_texture is not None:
|
||||
mr_float = mr_texture.astype(np.float32) / 255.0
|
||||
mesh.metallicRoughness = torch.tensor(mr_float, dtype=torch.float32, device=device).contiguous()
|
||||
print(f"[FastMesh] MetallicRoughness: {mesh.metallicRoughness.shape}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Error extracting textures: {e}")
|
||||
|
||||
if should_create_empty:
|
||||
mesh._create_empty_albedo()
|
||||
|
||||
@classmethod
|
||||
def _extract_texture_by_index(cls, doc, texture_idx):
|
||||
"""Extract texture by index from GLTF"""
|
||||
|
||||
try:
|
||||
if texture_idx >= len(doc.textures()):
|
||||
return None
|
||||
|
||||
texture = doc.textures()[texture_idx]
|
||||
image_idx = texture.get("source", 0)
|
||||
|
||||
if image_idx >= len(doc.images()):
|
||||
return None
|
||||
|
||||
image_info = doc.images()[image_idx]
|
||||
buffer_view_idx = image_info.get("bufferView")
|
||||
|
||||
if buffer_view_idx is None:
|
||||
return None
|
||||
|
||||
buffer_view = doc.bufferViews()[buffer_view_idx]
|
||||
buffer_idx = buffer_view.get("buffer", 0)
|
||||
byte_offset = buffer_view.get("byteOffset", 0)
|
||||
byte_length = buffer_view.get("byteLength", 0)
|
||||
|
||||
binary_data = get_binary_data(doc, buffer_idx)
|
||||
image_bytes = binary_data[byte_offset:byte_offset + byte_length]
|
||||
|
||||
# Decode image
|
||||
image_array = np.frombuffer(image_bytes, dtype=np.uint8)
|
||||
texture_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
|
||||
|
||||
if texture_image is not None:
|
||||
texture_image = cv2.cvtColor(texture_image, cv2.COLOR_BGR2RGB)
|
||||
return texture_image
|
||||
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Error extracting texture {texture_idx}: {e}")
|
||||
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _load_obj_fast(cls, path, device=None, **kwargs):
|
||||
"""Fast loading of OBJ (fallback to standard method)"""
|
||||
standard_mesh = Mesh.load_obj(path, device=device, **kwargs)
|
||||
return cls._from_standard_mesh(standard_mesh)
|
||||
|
||||
@classmethod
|
||||
def _from_standard_mesh(cls, standard_mesh):
|
||||
"""Convert standard Mesh to FastMesh"""
|
||||
if standard_mesh is None:
|
||||
return None
|
||||
|
||||
return cls(
|
||||
v=standard_mesh.v,
|
||||
f=standard_mesh.f,
|
||||
vn=standard_mesh.vn,
|
||||
fn=standard_mesh.fn,
|
||||
vt=standard_mesh.vt,
|
||||
ft=standard_mesh.ft,
|
||||
vc=standard_mesh.vc,
|
||||
albedo=standard_mesh.albedo,
|
||||
metallicRoughness=standard_mesh.metallicRoughness,
|
||||
device=standard_mesh.device
|
||||
)
|
||||
|
||||
def _create_empty_albedo(self):
|
||||
"""Create empty albedo texture"""
|
||||
texture = np.ones((1024, 1024, 3), dtype=np.float32) * np.array([0.5, 0.5, 0.5])
|
||||
self.albedo = torch.tensor(texture, dtype=torch.float32, device=self.device)
|
||||
print(f"[FastMesh] Empty texture: {self.albedo.shape}")
|
||||
|
||||
def _clean_mesh(self):
|
||||
"""Clean mesh (simplified version)"""
|
||||
try:
|
||||
from kiui.mesh_utils import clean_mesh
|
||||
vertices = self.v.detach().cpu().numpy()
|
||||
triangles = self.f.detach().cpu().numpy()
|
||||
vertices, triangles = clean_mesh(vertices, triangles, remesh=False)
|
||||
self.v = torch.from_numpy(vertices).contiguous().float().to(self.device)
|
||||
self.f = torch.from_numpy(triangles).contiguous().int().to(self.device)
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Error cleaning: {e}")
|
||||
|
||||
def _apply_front_dir_rotation(self, front_dir):
|
||||
"""Apply rotation for aligning front_dir with +z"""
|
||||
# Use the same logic as in standard Mesh
|
||||
if front_dir == "+z":
|
||||
return
|
||||
|
||||
# Define transformation matrix
|
||||
if "-z" in front_dir:
|
||||
T = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, -1]], device=self.device, dtype=torch.float32)
|
||||
elif "+x" in front_dir:
|
||||
T = torch.tensor([[0, 0, 1], [0, 1, 0], [1, 0, 0]], device=self.device, dtype=torch.float32)
|
||||
elif "-x" in front_dir:
|
||||
T = torch.tensor([[0, 0, -1], [0, 1, 0], [1, 0, 0]], device=self.device, dtype=torch.float32)
|
||||
elif "+y" in front_dir:
|
||||
T = torch.tensor([[1, 0, 0], [0, 0, 1], [0, 1, 0]], device=self.device, dtype=torch.float32)
|
||||
elif "-y" in front_dir:
|
||||
T = torch.tensor([[1, 0, 0], [0, 0, -1], [0, 1, 0]], device=self.device, dtype=torch.float32)
|
||||
else:
|
||||
T = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], device=self.device, dtype=torch.float32)
|
||||
|
||||
# Additional rotations
|
||||
if '1' in front_dir:
|
||||
T @= torch.tensor([[0, -1, 0], [1, 0, 0], [0, 0, 1]], device=self.device, dtype=torch.float32)
|
||||
elif '2' in front_dir:
|
||||
T @= torch.tensor([[1, 0, 0], [0, -1, 0], [0, 0, 1]], device=self.device, dtype=torch.float32)
|
||||
elif '3' in front_dir:
|
||||
T @= torch.tensor([[0, 1, 0], [-1, 0, 0], [0, 0, 1]], device=self.device, dtype=torch.float32)
|
||||
|
||||
# Apply transformation
|
||||
if self.v is not None:
|
||||
self.v @= T
|
||||
if self.vn is not None:
|
||||
self.vn @= T
|
||||
|
||||
# Use methods from standard Mesh
|
||||
def aabb(self):
|
||||
"""AABB mesh"""
|
||||
return torch.min(self.v, dim=0).values, torch.max(self.v, dim=0).values
|
||||
|
||||
@torch.no_grad()
|
||||
def auto_size(self, bound=0.9):
|
||||
"""Automatic size change"""
|
||||
vmin, vmax = self.aabb()
|
||||
self.ori_center = (vmax + vmin) / 2
|
||||
self.ori_scale = 2 * bound / torch.max(vmax - vmin).item()
|
||||
self.v = (self.v - self.ori_center) * self.ori_scale
|
||||
|
||||
def auto_normal(self):
|
||||
"""Automatic calculation of normals"""
|
||||
i0, i1, i2 = self.f[:, 0].long(), self.f[:, 1].long(), self.f[:, 2].long()
|
||||
v0, v1, v2 = self.v[i0, :], self.v[i1, :], self.v[i2, :]
|
||||
|
||||
face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1)
|
||||
|
||||
vn = torch.zeros_like(self.v)
|
||||
vn.scatter_add_(0, i0[:, None].repeat(1, 3), face_normals)
|
||||
vn.scatter_add_(0, i1[:, None].repeat(1, 3), face_normals)
|
||||
vn.scatter_add_(0, i2[:, None].repeat(1, 3), face_normals)
|
||||
|
||||
vn = torch.where(
|
||||
dot(vn, vn) > 1e-20,
|
||||
vn,
|
||||
torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device=vn.device),
|
||||
)
|
||||
vn = safe_normalize(vn)
|
||||
|
||||
self.vn = vn
|
||||
self.fn = self.f
|
||||
|
||||
def auto_uv(self, cache_path=None, vmap=True):
|
||||
"""Automatic calculation of UV (use method from standard Mesh)"""
|
||||
# Temporarily delegate to standard Mesh
|
||||
temp_mesh = Mesh(
|
||||
v=self.v, f=self.f, vn=self.vn, fn=self.fn,
|
||||
vt=self.vt, ft=self.ft, vc=self.vc,
|
||||
albedo=self.albedo, metallicRoughness=self.metallicRoughness,
|
||||
device=self.device
|
||||
)
|
||||
temp_mesh.auto_uv(cache_path=cache_path, vmap=vmap)
|
||||
|
||||
# Copy result back
|
||||
self.vt = temp_mesh.vt
|
||||
self.ft = temp_mesh.ft
|
||||
if vmap:
|
||||
self.v = temp_mesh.v
|
||||
self.f = temp_mesh.f
|
||||
self.vn = temp_mesh.vn
|
||||
self.fn = temp_mesh.fn
|
||||
if temp_mesh.vc is not None:
|
||||
self.vc = temp_mesh.vc
|
||||
|
||||
def to(self, device):
|
||||
"""Move to another device"""
|
||||
self.device = device
|
||||
for name in ["v", "f", "vn", "fn", "vt", "ft", "albedo", "vc", "metallicRoughness"]:
|
||||
tensor = getattr(self, name)
|
||||
if tensor is not None:
|
||||
setattr(self, name, tensor.to(device))
|
||||
return self
|
||||
|
||||
def write(self, path):
|
||||
"""Fast saving using mesh_processor for GLB/GLTF files"""
|
||||
if path.endswith(".glb"):
|
||||
self._write_glb_fast(path)
|
||||
elif path.endswith(".gltf"):
|
||||
# For GLTF, fallback to standard method (less common)
|
||||
self._write_fallback(path)
|
||||
elif path.endswith(".ply"):
|
||||
self._write_ply_fast(path)
|
||||
elif path.endswith(".obj"):
|
||||
self._write_obj_fast(path)
|
||||
else:
|
||||
# Fallback for unknown formats
|
||||
self._write_fallback(path)
|
||||
|
||||
def _write_glb_fast(self, path):
|
||||
"""Fast GLB writing using mesh_processor"""
|
||||
try:
|
||||
# Prepare data
|
||||
vertices = self.v.detach().cpu().numpy().astype(np.float32)
|
||||
faces = self.f.detach().cpu().numpy().astype(np.uint32)
|
||||
|
||||
# Create minimal GLTF structure
|
||||
gltf_json = {
|
||||
"asset": {"version": "2.0", "generator": "FastMesh"},
|
||||
"scenes": [{"nodes": [0]}],
|
||||
"nodes": [{"mesh": 0}],
|
||||
"meshes": [{
|
||||
"primitives": [{
|
||||
"mode": 4 # TRIANGLES
|
||||
}]
|
||||
}]
|
||||
}
|
||||
|
||||
# Create document with empty buffer
|
||||
doc = GltfDocument(json_data=gltf_json, binary_buffers={0: np.array([], dtype=np.uint8)})
|
||||
|
||||
# Add vertex data using helper function
|
||||
pos_accessor, _ = append_accessor_and_bufferview(
|
||||
doc, vertices,
|
||||
component_type=5126, # FLOAT
|
||||
element_type="VEC3",
|
||||
target=34962 # ARRAY_BUFFER
|
||||
)
|
||||
|
||||
# Add face data using helper function
|
||||
faces_flat = faces.flatten()
|
||||
face_component_type = 5123 if vertices.shape[0] <= 65535 else 5125 # UNSIGNED_SHORT or UNSIGNED_INT
|
||||
face_dtype = np.uint16 if face_component_type == 5123 else np.uint32
|
||||
|
||||
idx_accessor, _ = append_accessor_and_bufferview(
|
||||
doc, faces_flat.astype(face_dtype),
|
||||
component_type=face_component_type,
|
||||
element_type="SCALAR",
|
||||
target=34963 # ELEMENT_ARRAY_BUFFER
|
||||
)
|
||||
|
||||
# Update primitive with accessor indices
|
||||
primitive = doc.json()["meshes"][0]["primitives"][0]
|
||||
primitive["attributes"] = {"POSITION": pos_accessor}
|
||||
primitive["indices"] = idx_accessor
|
||||
|
||||
# Save to file
|
||||
with open(path, "wb") as f:
|
||||
save_glb(doc, f)
|
||||
|
||||
print(f"[FastMesh] Fast GLB saved: {path}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Fast GLB save failed: {e}, using fallback")
|
||||
self._write_fallback(path)
|
||||
|
||||
def _write_ply_fast(self, path):
|
||||
"""Fast PLY writing"""
|
||||
vertices = self.v.detach().cpu().numpy()
|
||||
faces = self.f.detach().cpu().numpy()
|
||||
|
||||
with open(path, 'w') as f:
|
||||
f.write("ply\n")
|
||||
f.write("format ascii 1.0\n")
|
||||
f.write(f"element vertex {len(vertices)}\n")
|
||||
f.write("property float x\n")
|
||||
f.write("property float y\n")
|
||||
f.write("property float z\n")
|
||||
f.write(f"element face {len(faces)}\n")
|
||||
f.write("property list uchar int vertex_indices\n")
|
||||
f.write("end_header\n")
|
||||
|
||||
for v in vertices:
|
||||
f.write(f"{v[0]} {v[1]} {v[2]}\n")
|
||||
|
||||
for face in faces:
|
||||
f.write(f"3 {face[0]} {face[1]} {face[2]}\n")
|
||||
|
||||
def _write_obj_fast(self, path):
|
||||
"""Fast OBJ writing"""
|
||||
vertices = self.v.detach().cpu().numpy()
|
||||
faces = self.f.detach().cpu().numpy()
|
||||
|
||||
with open(path, 'w') as f:
|
||||
for v in vertices:
|
||||
f.write(f"v {v[0]} {v[1]} {v[2]}\n")
|
||||
|
||||
for face in faces:
|
||||
f.write(f"f {face[0]+1} {face[1]+1} {face[2]+1}\n")
|
||||
|
||||
def _write_fallback(self, path):
|
||||
"""Fallback to standard Mesh write method"""
|
||||
temp_mesh = Mesh(
|
||||
v=self.v, f=self.f, vn=self.vn, fn=self.fn,
|
||||
vt=self.vt, ft=self.ft, vc=self.vc,
|
||||
albedo=self.albedo, metallicRoughness=self.metallicRoughness,
|
||||
device=self.device
|
||||
)
|
||||
temp_mesh.ori_center = self.ori_center
|
||||
temp_mesh.ori_scale = self.ori_scale
|
||||
return temp_mesh.write(path)
|
||||
@@ -1130,404 +1130,6 @@ class Mesh:
|
||||
return pcd
|
||||
|
||||
|
||||
class FastMesh:
|
||||
"""
|
||||
Fast implementation of mesh based on FastGLB without dependency on trimesh.
|
||||
Uses direct loading of GLB/GLTF through mesh_processor.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
v: Optional[Tensor] = None,
|
||||
f: Optional[Tensor] = None,
|
||||
vn: Optional[Tensor] = None,
|
||||
fn: Optional[Tensor] = None,
|
||||
vt: Optional[Tensor] = None,
|
||||
ft: Optional[Tensor] = None,
|
||||
vc: Optional[Tensor] = None,
|
||||
albedo: Optional[Tensor] = None,
|
||||
metallicRoughness: Optional[Tensor] = None,
|
||||
device: Optional[torch.device] = None,
|
||||
):
|
||||
"""Initialize FastMesh with the same parameters as Mesh"""
|
||||
self.device = device if device is not None else torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
self.v = v
|
||||
self.vn = vn
|
||||
self.vt = vt
|
||||
self.f = f
|
||||
self.fn = fn
|
||||
self.ft = ft
|
||||
self.vc = vc
|
||||
self.albedo = albedo
|
||||
self.metallicRoughness = metallicRoughness
|
||||
self.ori_center = 0
|
||||
self.ori_scale = 1
|
||||
|
||||
@classmethod
|
||||
def load(cls, path, resize=True, renormal=True, retex=False, clean=False, bound=0.5, front_dir='+z', **kwargs):
|
||||
"""Loading mesh through FastGLB approach"""
|
||||
|
||||
device = kwargs.get('device')
|
||||
if device is None:
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
if path.endswith(".obj"):
|
||||
mesh = cls._load_obj_fast(path, device=device, **kwargs)
|
||||
elif path.endswith((".glb", ".gltf")):
|
||||
mesh = cls._load_gltf_fast(path, device=device, **kwargs)
|
||||
else:
|
||||
standard_mesh = Mesh.load(path, resize=False, renormal=False, retex=False, clean=False, **kwargs)
|
||||
mesh = cls._from_standard_mesh(standard_mesh)
|
||||
|
||||
if mesh is None:
|
||||
return None
|
||||
|
||||
if clean:
|
||||
mesh._clean_mesh()
|
||||
|
||||
print(f"[FastMesh loading] v: {mesh.v.shape}, f: {mesh.f.shape}")
|
||||
|
||||
if resize:
|
||||
mesh.auto_size(bound=bound)
|
||||
|
||||
if renormal or mesh.vn is None:
|
||||
mesh.auto_normal()
|
||||
print(f"[FastMesh loading] vn: {mesh.vn.shape}, fn: {mesh.fn.shape}")
|
||||
|
||||
if retex or (mesh.albedo is not None and mesh.vt is None):
|
||||
mesh.auto_uv(cache_path=path)
|
||||
print(f"[FastMesh loading] vt: {mesh.vt.shape}, ft: {mesh.ft.shape}")
|
||||
|
||||
if front_dir != "+z":
|
||||
mesh._apply_front_dir_rotation(front_dir)
|
||||
|
||||
return mesh
|
||||
|
||||
@classmethod
|
||||
def _load_gltf_fast(cls, path, device=None, **kwargs):
|
||||
"""Fast loading of GLB/GLTF through mesh_processor"""
|
||||
|
||||
try:
|
||||
print(f"[FastMesh] Loading GLB/GLTF: {path}")
|
||||
|
||||
# Load GLTF document
|
||||
from .io_gltf import load_gltf_or_glb, get_binary_data
|
||||
from .mesh_ops import get_all_meshes_triangles
|
||||
from .accessors import access_data
|
||||
|
||||
doc = load_gltf_or_glb(path)
|
||||
print(f"[FastMesh] GLTF document: {len(doc.meshes())} meshes")
|
||||
|
||||
# Extract geometry
|
||||
vertices, faces, mesh_groups = get_all_meshes_triangles(doc, transform_to_global=True)
|
||||
|
||||
if len(vertices) == 0 or len(faces) == 0:
|
||||
print(f"[FastMesh] No geometry in {path}")
|
||||
return None
|
||||
|
||||
mesh = cls(device=device)
|
||||
mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device)
|
||||
mesh.f = torch.tensor(faces, dtype=torch.int32, device=device)
|
||||
|
||||
# Extract additional data
|
||||
cls._extract_gltf_attributes(mesh, doc, device)
|
||||
cls._extract_gltf_textures(mesh, doc, device)
|
||||
|
||||
print(f"[FastMesh] Loaded: {len(vertices)} vertices, {len(faces)} faces")
|
||||
return mesh
|
||||
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Error loading GLB/GLTF {path}: {e}")
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _extract_gltf_attributes(cls, mesh, doc, device):
|
||||
"""Extract attributes from GLTF (UV, normals, colors)"""
|
||||
|
||||
if not doc.meshes():
|
||||
return
|
||||
|
||||
first_mesh = doc.meshes()[0]
|
||||
for primitive in first_mesh.get("primitives", []):
|
||||
attributes = primitive.get("attributes", {})
|
||||
|
||||
# UV coordinates
|
||||
if "TEXCOORD_0" in attributes and mesh.vt is None:
|
||||
try:
|
||||
from .accessors import access_data
|
||||
uv_accessor_idx = int(attributes["TEXCOORD_0"])
|
||||
uv_data = access_data(doc, uv_accessor_idx).astype(np.float32)
|
||||
uv_data[:, 1] = 1.0 - uv_data[:, 1] # Flip V
|
||||
mesh.vt = torch.tensor(uv_data, dtype=torch.float32, device=device)
|
||||
mesh.ft = mesh.f
|
||||
print(f"[FastMesh] UV: {mesh.vt.shape}")
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Error UV: {e}")
|
||||
|
||||
# Normals
|
||||
if "NORMAL" in attributes and mesh.vn is None:
|
||||
try:
|
||||
from .accessors import access_data
|
||||
normal_accessor_idx = int(attributes["NORMAL"])
|
||||
normal_data = access_data(doc, normal_accessor_idx).astype(np.float32)
|
||||
mesh.vn = torch.tensor(normal_data, dtype=torch.float32, device=device)
|
||||
mesh.fn = mesh.f
|
||||
print(f"[FastMesh] Normals: {mesh.vn.shape}")
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Error normals: {e}")
|
||||
|
||||
# Vertex colors
|
||||
if "COLOR_0" in attributes and mesh.vc is None:
|
||||
try:
|
||||
from .accessors import access_data
|
||||
color_accessor_idx = int(attributes["COLOR_0"])
|
||||
color_data = access_data(doc, color_accessor_idx).astype(np.float32)
|
||||
if color_data.shape[1] > 3:
|
||||
color_data = color_data[:, :3]
|
||||
mesh.vc = torch.tensor(color_data, dtype=torch.float32, device=device)
|
||||
print(f"[FastMesh] Colors: {mesh.vc.shape}")
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Error colors: {e}")
|
||||
|
||||
@classmethod
|
||||
def _extract_gltf_textures(cls, mesh, doc, device):
|
||||
"""Extract textures from GLTF"""
|
||||
|
||||
should_create_empty = True
|
||||
|
||||
try:
|
||||
if not (doc.materials() and doc.textures() and doc.images()):
|
||||
mesh._create_empty_albedo()
|
||||
return
|
||||
|
||||
material = doc.materials()[0]
|
||||
pbr = material.get("pbrMetallicRoughness", {})
|
||||
|
||||
# Albedo texture
|
||||
base_color_texture = pbr.get("baseColorTexture")
|
||||
if base_color_texture:
|
||||
albedo_texture = cls._extract_texture_by_index(doc, base_color_texture.get("index", 0))
|
||||
if albedo_texture is not None:
|
||||
albedo_float = albedo_texture.astype(np.float32) / 255.0
|
||||
mesh.albedo = torch.tensor(albedo_float, dtype=torch.float32, device=device).contiguous()
|
||||
should_create_empty = False
|
||||
print(f"[FastMesh] Albedo: {mesh.albedo.shape}")
|
||||
|
||||
# Metallic-Roughness texture
|
||||
mr_texture_info = pbr.get("metallicRoughnessTexture")
|
||||
if mr_texture_info and not should_create_empty:
|
||||
mr_texture = cls._extract_texture_by_index(doc, mr_texture_info.get("index", 0))
|
||||
if mr_texture is not None:
|
||||
mr_float = mr_texture.astype(np.float32) / 255.0
|
||||
mesh.metallicRoughness = torch.tensor(mr_float, dtype=torch.float32, device=device).contiguous()
|
||||
print(f"[FastMesh] MetallicRoughness: {mesh.metallicRoughness.shape}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Error extracting textures: {e}")
|
||||
|
||||
if should_create_empty:
|
||||
mesh._create_empty_albedo()
|
||||
|
||||
@classmethod
|
||||
def _extract_texture_by_index(cls, doc, texture_idx):
|
||||
"""Extract texture by index from GLTF"""
|
||||
|
||||
try:
|
||||
from .io_gltf import get_binary_data
|
||||
|
||||
if texture_idx >= len(doc.textures()):
|
||||
return None
|
||||
|
||||
texture = doc.textures()[texture_idx]
|
||||
image_idx = texture.get("source", 0)
|
||||
|
||||
if image_idx >= len(doc.images()):
|
||||
return None
|
||||
|
||||
image_info = doc.images()[image_idx]
|
||||
buffer_view_idx = image_info.get("bufferView")
|
||||
|
||||
if buffer_view_idx is None:
|
||||
return None
|
||||
|
||||
buffer_view = doc.bufferViews()[buffer_view_idx]
|
||||
buffer_idx = buffer_view.get("buffer", 0)
|
||||
byte_offset = buffer_view.get("byteOffset", 0)
|
||||
byte_length = buffer_view.get("byteLength", 0)
|
||||
|
||||
binary_data = get_binary_data(doc, buffer_idx)
|
||||
image_bytes = binary_data[byte_offset:byte_offset + byte_length]
|
||||
|
||||
# Decode image
|
||||
image_array = np.frombuffer(image_bytes, dtype=np.uint8)
|
||||
texture_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
|
||||
|
||||
if texture_image is not None:
|
||||
texture_image = cv2.cvtColor(texture_image, cv2.COLOR_BGR2RGB)
|
||||
return texture_image
|
||||
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Error extracting texture {texture_idx}: {e}")
|
||||
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _load_obj_fast(cls, path, device=None, **kwargs):
|
||||
"""Fast loading of OBJ (fallback to standard method)"""
|
||||
standard_mesh = Mesh.load_obj(path, device=device, **kwargs)
|
||||
return cls._from_standard_mesh(standard_mesh)
|
||||
|
||||
@classmethod
|
||||
def _from_standard_mesh(cls, standard_mesh):
|
||||
"""Convert standard Mesh to FastMesh"""
|
||||
if standard_mesh is None:
|
||||
return None
|
||||
|
||||
return cls(
|
||||
v=standard_mesh.v,
|
||||
f=standard_mesh.f,
|
||||
vn=standard_mesh.vn,
|
||||
fn=standard_mesh.fn,
|
||||
vt=standard_mesh.vt,
|
||||
ft=standard_mesh.ft,
|
||||
vc=standard_mesh.vc,
|
||||
albedo=standard_mesh.albedo,
|
||||
metallicRoughness=standard_mesh.metallicRoughness,
|
||||
device=standard_mesh.device
|
||||
)
|
||||
|
||||
def _create_empty_albedo(self):
|
||||
"""Create empty albedo texture"""
|
||||
texture = np.ones((1024, 1024, 3), dtype=np.float32) * np.array([0.5, 0.5, 0.5])
|
||||
self.albedo = torch.tensor(texture, dtype=torch.float32, device=self.device)
|
||||
print(f"[FastMesh] Empty texture: {self.albedo.shape}")
|
||||
|
||||
def _clean_mesh(self):
|
||||
"""Clean mesh (simplified version)"""
|
||||
try:
|
||||
from kiui.mesh_utils import clean_mesh
|
||||
vertices = self.v.detach().cpu().numpy()
|
||||
triangles = self.f.detach().cpu().numpy()
|
||||
vertices, triangles = clean_mesh(vertices, triangles, remesh=False)
|
||||
self.v = torch.from_numpy(vertices).contiguous().float().to(self.device)
|
||||
self.f = torch.from_numpy(triangles).contiguous().int().to(self.device)
|
||||
except Exception as e:
|
||||
print(f"[FastMesh] Error cleaning: {e}")
|
||||
|
||||
def _apply_front_dir_rotation(self, front_dir):
|
||||
"""Apply rotation for aligning front_dir with +z"""
|
||||
# Use the same logic as in standard Mesh
|
||||
if front_dir == "+z":
|
||||
return
|
||||
|
||||
# Define transformation matrix
|
||||
if "-z" in front_dir:
|
||||
T = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, -1]], device=self.device, dtype=torch.float32)
|
||||
elif "+x" in front_dir:
|
||||
T = torch.tensor([[0, 0, 1], [0, 1, 0], [1, 0, 0]], device=self.device, dtype=torch.float32)
|
||||
elif "-x" in front_dir:
|
||||
T = torch.tensor([[0, 0, -1], [0, 1, 0], [1, 0, 0]], device=self.device, dtype=torch.float32)
|
||||
elif "+y" in front_dir:
|
||||
T = torch.tensor([[1, 0, 0], [0, 0, 1], [0, 1, 0]], device=self.device, dtype=torch.float32)
|
||||
elif "-y" in front_dir:
|
||||
T = torch.tensor([[1, 0, 0], [0, 0, -1], [0, 1, 0]], device=self.device, dtype=torch.float32)
|
||||
else:
|
||||
T = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], device=self.device, dtype=torch.float32)
|
||||
|
||||
# Additional rotations
|
||||
if '1' in front_dir:
|
||||
T @= torch.tensor([[0, -1, 0], [1, 0, 0], [0, 0, 1]], device=self.device, dtype=torch.float32)
|
||||
elif '2' in front_dir:
|
||||
T @= torch.tensor([[1, 0, 0], [0, -1, 0], [0, 0, 1]], device=self.device, dtype=torch.float32)
|
||||
elif '3' in front_dir:
|
||||
T @= torch.tensor([[0, 1, 0], [-1, 0, 0], [0, 0, 1]], device=self.device, dtype=torch.float32)
|
||||
|
||||
# Apply transformation
|
||||
if self.v is not None:
|
||||
self.v @= T
|
||||
if self.vn is not None:
|
||||
self.vn @= T
|
||||
|
||||
# Use methods from standard Mesh
|
||||
def aabb(self):
|
||||
"""AABB mesh"""
|
||||
return torch.min(self.v, dim=0).values, torch.max(self.v, dim=0).values
|
||||
|
||||
@torch.no_grad()
|
||||
def auto_size(self, bound=0.9):
|
||||
"""Automatic size change"""
|
||||
vmin, vmax = self.aabb()
|
||||
self.ori_center = (vmax + vmin) / 2
|
||||
self.ori_scale = 2 * bound / torch.max(vmax - vmin).item()
|
||||
self.v = (self.v - self.ori_center) * self.ori_scale
|
||||
|
||||
def auto_normal(self):
|
||||
"""Automatic calculation of normals"""
|
||||
i0, i1, i2 = self.f[:, 0].long(), self.f[:, 1].long(), self.f[:, 2].long()
|
||||
v0, v1, v2 = self.v[i0, :], self.v[i1, :], self.v[i2, :]
|
||||
|
||||
face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1)
|
||||
|
||||
vn = torch.zeros_like(self.v)
|
||||
vn.scatter_add_(0, i0[:, None].repeat(1, 3), face_normals)
|
||||
vn.scatter_add_(0, i1[:, None].repeat(1, 3), face_normals)
|
||||
vn.scatter_add_(0, i2[:, None].repeat(1, 3), face_normals)
|
||||
|
||||
vn = torch.where(
|
||||
dot(vn, vn) > 1e-20,
|
||||
vn,
|
||||
torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device=vn.device),
|
||||
)
|
||||
vn = safe_normalize(vn)
|
||||
|
||||
self.vn = vn
|
||||
self.fn = self.f
|
||||
|
||||
def auto_uv(self, cache_path=None, vmap=True):
|
||||
"""Automatic calculation of UV (use method from standard Mesh)"""
|
||||
# Temporarily delegate to standard Mesh
|
||||
temp_mesh = Mesh(
|
||||
v=self.v, f=self.f, vn=self.vn, fn=self.fn,
|
||||
vt=self.vt, ft=self.ft, vc=self.vc,
|
||||
albedo=self.albedo, metallicRoughness=self.metallicRoughness,
|
||||
device=self.device
|
||||
)
|
||||
temp_mesh.auto_uv(cache_path=cache_path, vmap=vmap)
|
||||
|
||||
# Copy result back
|
||||
self.vt = temp_mesh.vt
|
||||
self.ft = temp_mesh.ft
|
||||
if vmap:
|
||||
self.v = temp_mesh.v
|
||||
self.f = temp_mesh.f
|
||||
self.vn = temp_mesh.vn
|
||||
self.fn = temp_mesh.fn
|
||||
if temp_mesh.vc is not None:
|
||||
self.vc = temp_mesh.vc
|
||||
|
||||
def to(self, device):
|
||||
"""Move to another device"""
|
||||
self.device = device
|
||||
for name in ["v", "f", "vn", "fn", "vt", "ft", "albedo", "vc", "metallicRoughness"]:
|
||||
tensor = getattr(self, name)
|
||||
if tensor is not None:
|
||||
setattr(self, name, tensor.to(device))
|
||||
return self
|
||||
|
||||
def write(self, path):
|
||||
"""Saving to file (use methods from standard Mesh)"""
|
||||
# Create temporary standard Mesh for saving
|
||||
temp_mesh = Mesh(
|
||||
v=self.v, f=self.f, vn=self.vn, fn=self.fn,
|
||||
vt=self.vt, ft=self.ft, vc=self.vc,
|
||||
albedo=self.albedo, metallicRoughness=self.metallicRoughness,
|
||||
device=self.device
|
||||
)
|
||||
temp_mesh.ori_center = self.ori_center
|
||||
temp_mesh.ori_scale = self.ori_scale
|
||||
return temp_mesh.write(path)
|
||||
|
||||
|
||||
class PointCloud(NamedTuple):
|
||||
points: np.array
|
||||
colors: np.array
|
||||
|
||||
@@ -38,7 +38,8 @@ import trimesh
|
||||
from PIL import Image
|
||||
|
||||
from mesh_processor.export_utils import export_to_fastmesh, export_to_mesh
|
||||
from mesh_processor.mesh import FastMesh, Mesh
|
||||
from mesh_processor.fastmesh import FastMesh
|
||||
from mesh_processor.mesh import Mesh
|
||||
from mesh_processor.mesh_utils import (
|
||||
ply_to_points_cloud,
|
||||
get_target_axis_and_scale,
|
||||
@@ -288,6 +289,7 @@ class Load_3D_Mesh:
|
||||
return {
|
||||
"required": {
|
||||
"mesh_file_path": ("STRING", {"default": '', "multiline": False}),
|
||||
"use_fastmesh": ("BOOLEAN", {"default": True, "tooltip": "Use FastMesh for GLB/GLTF files"}),
|
||||
"resize": ("BOOLEAN", {"default": False},),
|
||||
"renormal": ("BOOLEAN", {"default": True},),
|
||||
"retex": ("BOOLEAN", {"default": False},),
|
||||
@@ -306,7 +308,7 @@ class Load_3D_Mesh:
|
||||
FUNCTION = "load_mesh"
|
||||
CATEGORY = "Comfy3D/Import|Export"
|
||||
|
||||
def load_mesh(self, mesh_file_path, resize, renormal, retex, optimizable, clean, resize_bound):
|
||||
def load_mesh(self, mesh_file_path, use_fastmesh, resize, renormal, retex, optimizable, clean, resize_bound):
|
||||
mesh = None
|
||||
|
||||
if not os.path.isabs(mesh_file_path):
|
||||
@@ -316,7 +318,10 @@ class Load_3D_Mesh:
|
||||
folder, filename = os.path.split(mesh_file_path)
|
||||
if filename.lower().endswith(SUPPORTED_3D_EXTENSIONS):
|
||||
with torch.inference_mode(not optimizable):
|
||||
mesh = Mesh.load(mesh_file_path, resize, renormal, retex, clean, resize_bound)
|
||||
if use_fastmesh and mesh_file_path.endswith((".glb", ".gltf")):
|
||||
mesh = FastMesh.load(mesh_file_path, resize, renormal, retex, clean, resize_bound)
|
||||
else:
|
||||
mesh = Mesh.load(mesh_file_path, resize, renormal, retex, clean, resize_bound)
|
||||
else:
|
||||
cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3D file extensions: {SUPPORTED_3D_EXTENSIONS}").error.print()
|
||||
else:
|
||||
@@ -366,6 +371,7 @@ class Save_3D_Mesh:
|
||||
"required": {
|
||||
"mesh": ("MESH",),
|
||||
"save_path": ("STRING", {"default": 'Mesh_%Y-%m-%d-%M-%S-%f.glb', "multiline": False}),
|
||||
"use_fastmesh": ("BOOLEAN", {"default": True, "tooltip": "Use FastMesh for fast saving GLB"}),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -379,11 +385,26 @@ class Save_3D_Mesh:
|
||||
FUNCTION = "save_mesh"
|
||||
CATEGORY = "Comfy3D/Import|Export"
|
||||
|
||||
def save_mesh(self, mesh, save_path):
|
||||
def save_mesh(self, mesh, save_path, use_fastmesh):
|
||||
save_path = parse_save_filename(save_path, comfy_paths.output_directory, SUPPORTED_3D_EXTENSIONS, self.__class__.__name__)
|
||||
|
||||
if save_path is not None:
|
||||
mesh.write(save_path)
|
||||
|
||||
if use_fastmesh and save_path.endswith((".glb", ".ply", ".obj")):
|
||||
if hasattr(mesh, '_write_glb_fast'):
|
||||
mesh.write(save_path)
|
||||
else:
|
||||
fast_mesh = FastMesh(
|
||||
v=mesh.v, f=mesh.f, vn=mesh.vn, fn=mesh.fn,
|
||||
vt=mesh.vt, ft=mesh.ft, vc=mesh.vc,
|
||||
albedo=mesh.albedo, metallicRoughness=getattr(mesh, 'metallicRoughness', None),
|
||||
device=mesh.device
|
||||
)
|
||||
fast_mesh.ori_center = getattr(mesh, 'ori_center', 0)
|
||||
fast_mesh.ori_scale = getattr(mesh, 'ori_scale', 1)
|
||||
fast_mesh.write(save_path)
|
||||
else:
|
||||
mesh.write(save_path)
|
||||
|
||||
return (save_path, )
|
||||
|
||||
|
||||
Reference in New Issue
Block a user