import os import cv2 from typing import NamedTuple, Optional import torch from torch import Tensor import trimesh import numpy as np from PIL import Image 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, load_pbr=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, load_pbr=load_pbr, **{**kwargs, 'device': device}) elif path.endswith((".glb", ".gltf")): mesh = cls._load_gltf_fast(path, lazy_texture=lazy_texture, load_pbr=load_pbr, **{**kwargs, 'device': device}) 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() if resize: mesh.auto_size(bound=bound) if renormal or mesh.vn is None: mesh.auto_normal() if retex or (mesh.albedo is not None and mesh.vt is None): mesh.auto_uv(cache_path=path) if front_dir != "+z": mesh._apply_front_dir_rotation(front_dir) print(f"[FastMesh] Loading completed successfully!") # Show final texture status has_albedo = mesh.albedo is not None has_pbr = mesh.metallicRoughness is not None return mesh @classmethod def _load_gltf_fast(cls, path, device=None, lazy_texture=True, load_pbr=True, **kwargs): """Fast loading of GLB/GLTF through mesh_processor""" try: # Load GLTF document doc = load_gltf_or_glb(path) # Extract geometry vertices, faces, mesh_groups = get_all_meshes_triangles(doc, transform_to_global=True) if len(vertices) == 0 or len(faces) == 0: return None mesh = cls(device=device) # Convert to torch tensors in parallel mesh.v = torch.from_numpy(vertices.astype(np.float32)).to(device) mesh.f = torch.from_numpy(faces.astype(np.int32)).to(device) # Extract attributes and textures in parallel uvs, normals, colors = cls._extract_all_mesh_attributes(doc, mesh_groups, device) if load_pbr: cls._extract_gltf_textures_fast(mesh, doc, device, lazy_empty=lazy_texture) else: mesh._create_empty_albedo() # Set attributes (avoid redundant assignments) if uvs is not None: mesh.vt, mesh.ft = uvs, mesh.f if normals is not None: mesh.vn, mesh.fn = normals, mesh.f if colors is not None: mesh.vc = colors 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: # Try both possible key names group_mesh_idx = group.get('mesh_idx', group.get('mesh_index')) group_primitive_idx = group.get('primitive_idx', group.get('primitive_index')) if group_mesh_idx == mesh_idx and group_primitive_idx == primitive_idx: mesh_group = group break if mesh_group is None: continue # Extract vertex info from mesh_group (handle different key names) vertices_range = mesh_group.get('vertices_range', (0, 0)) vertex_offset = vertices_range[0] if isinstance(vertices_range, tuple) else mesh_group.get('vertex_offset', 0) vertex_count = vertices_range[1] - vertices_range[0] if isinstance(vertices_range, tuple) else mesh_group.get('vertex_count', 0) # 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) all_uvs.append(uv_data) except Exception: 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) except Exception: 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) except Exception: all_colors.append(None) else: all_colors.append(None) # Concatenate all data final_uvs = None final_normals = None final_colors = None # Optimized concatenation - only process non-None arrays if any(uv is not None for uv in all_uvs): valid_uvs = [uv for uv in all_uvs if uv is not None] if valid_uvs: if len(valid_uvs) == 1: final_uvs = torch.from_numpy(valid_uvs[0]).to(device) else: final_uvs = torch.from_numpy(np.concatenate(valid_uvs, axis=0)).to(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: if len(valid_normals) == 1: final_normals = torch.from_numpy(valid_normals[0]).to(device) else: final_normals = torch.from_numpy(np.concatenate(valid_normals, axis=0)).to(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: if len(valid_colors) == 1: final_colors = torch.from_numpy(valid_colors[0]).to(device) else: final_colors = torch.from_numpy(np.concatenate(valid_colors, axis=0)).to(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 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) except Exception as e: print(f"[FastMesh] Error colors: {e}") @classmethod def _extract_gltf_textures_fast(cls, mesh, doc, device, lazy_empty=True): """Fast texture extraction""" if not (doc.materials() and doc.textures() and doc.images()): if not lazy_empty: mesh._create_empty_albedo() return try: material = doc.materials()[0] pbr = material.get("pbrMetallicRoughness", {}) albedo_texture = None mr_texture = None # 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)) # Metallic-Roughness texture mr_texture_info = pbr.get("metallicRoughnessTexture") if mr_texture_info: mr_texture = cls._extract_texture_by_index(doc, mr_texture_info.get("index", 0)) # Convert to tensors if albedo_texture is not None: mesh.albedo = torch.from_numpy(albedo_texture.astype(np.float32) / 255.0).to(device).contiguous() if mr_texture is not None: mesh.metallicRoughness = torch.from_numpy(mr_texture.astype(np.float32) / 255.0).to(device).contiguous() except Exception as e: print(f"[FastMesh] Error extracting textures: {e}") if not lazy_empty: mesh._create_empty_albedo() @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 # 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() 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, load_pbr=True, **kwargs): """Fast loading of OBJ using direct parsing""" try: return cls._load_obj_direct(path, device=device, load_pbr=load_pbr, **kwargs) except Exception as e: standard_mesh = Mesh.load_obj(path, device=device, **kwargs) return cls._from_standard_mesh(standard_mesh) @classmethod def _load_obj_direct(cls, path, device=None, load_pbr=True, **kwargs): """Direct OBJ parsing without trimesh dependency""" if device is None: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Parse OBJ file vertices, normals, uvs, faces = [], [], [], [] mtl_path = None with open(path, 'r', encoding='utf-8') as f: for line in f: line = line.strip() if not line or line.startswith('#'): continue parts = line.split() if not parts: continue prefix = parts[0] if prefix == 'v': vertices.append([float(parts[i]) for i in (1, 2, 3)]) elif prefix == 'vn': normals.append([float(parts[i]) for i in (1, 2, 3)]) elif prefix == 'vt': u, v = float(parts[1]), float(parts[2]) uvs.append([u, 1.0 - v]) # Flip V for consistency elif prefix == 'f': # Parse face indices face_verts = [] for vert_str in parts[1:]: indices = vert_str.split('/') vi = int(indices[0]) - 1 if indices[0] else 0 ti = int(indices[1]) - 1 if len(indices) > 1 and indices[1] else 0 ni = int(indices[2]) - 1 if len(indices) > 2 and indices[2] else 0 face_verts.append((vi, ti, ni)) # Triangulate face (fan triangulation) for i in range(1, len(face_verts) - 1): faces.extend([face_verts[0], face_verts[i], face_verts[i + 1]]) elif prefix == 'mtllib': mtl_path = os.path.join(os.path.dirname(path), parts[1]) if not vertices or not faces: raise ValueError("No valid geometry found in OBJ file") # Deduplicate vertices and build final arrays pos, normal, uv, indices = cls._deduplicate_obj_vertices( vertices, normals, uvs, faces ) # Create FastMesh mesh = cls(device=device) mesh.v = torch.tensor(pos, dtype=torch.float32, device=device) mesh.f = torch.tensor(indices, dtype=torch.int32, device=device).reshape(-1, 3) if normal: mesh.vn = torch.tensor(normal, dtype=torch.float32, device=device) mesh.fn = mesh.f if uv: mesh.vt = torch.tensor(uv, dtype=torch.float32, device=device) mesh.ft = mesh.f # Load textures from MTL if mtl_path and os.path.exists(mtl_path): cls._load_mtl_textures(mesh, mtl_path, device, load_pbr=load_pbr) else: # Create empty texture mesh._create_empty_albedo() return mesh @classmethod def _deduplicate_obj_vertices(cls, vertices, normals, uvs, faces): """Deduplicate vertices based on position/normal/uv combination""" pos, normal, uv = [], [], [] vertex_map = {} indices = [] for vi, ti, ni in faces: # Create unique key for this vertex combination key = (vi, ti if ti < len(uvs) else -1, ni if ni < len(normals) else -1) if key not in vertex_map: vertex_map[key] = len(pos) # Add vertex position pos.append(vertices[vi]) # Add UV coordinates if ti < len(uvs): uv.append(uvs[ti]) else: uv.append([0.0, 0.0]) # Add normal if ni < len(normals): normal.append(normals[ni]) else: normal.append([0.0, 0.0, 1.0]) indices.append(vertex_map[key]) return pos, normal, uv, indices @classmethod def _load_mtl_textures(cls, mesh, mtl_path, device, load_pbr=True): """Load textures from MTL file""" textures = {} base_color = [1.0, 1.0, 1.0] try: with open(mtl_path, 'r', encoding='utf-8') as f: for line in f: line = line.strip() if not line or line.startswith('#'): continue parts = line.split() if not parts: continue if parts[0] == 'Kd': base_color = [float(parts[i]) for i in (1, 2, 3)] elif parts[0] == 'map_Kd': textures['albedo'] = os.path.join(os.path.dirname(mtl_path), parts[1]) elif parts[0] == 'map_Pm': textures['metallic'] = os.path.join(os.path.dirname(mtl_path), parts[1]) elif parts[0] == 'map_Pr': textures['roughness'] = os.path.join(os.path.dirname(mtl_path), parts[1]) # Load albedo texture if 'albedo' in textures and os.path.exists(textures['albedo']): albedo_img = cv2.imread(textures['albedo'], cv2.IMREAD_COLOR) if albedo_img is not None: albedo_img = cv2.cvtColor(albedo_img, cv2.COLOR_BGR2RGB) albedo_tensor = torch.from_numpy(albedo_img.astype(np.float32) / 255.0).to(device) mesh.albedo = albedo_tensor.contiguous() else: # Create texture with base color texture = np.ones((1024, 1024, 3), dtype=np.float32) * np.array(base_color) mesh.albedo = torch.tensor(texture, dtype=torch.float32, device=device) # Load and combine metallic/roughness textures (only if PBR requested) if load_pbr and 'metallic' in textures and 'roughness' in textures: if os.path.exists(textures['metallic']) and os.path.exists(textures['roughness']): mr_texture = cls._combine_metallic_roughness( textures['metallic'], textures['roughness'] ) mesh.metallicRoughness = torch.from_numpy(mr_texture.astype(np.float32) / 255.0).to(device).contiguous() elif not load_pbr and ('metallic' in textures or 'roughness' in textures): print(f"[FastMesh] PBR disabled - skipping metallic/roughness textures") except Exception as e: print(f"[FastMesh] Error loading MTL textures: {e}") mesh._create_empty_albedo() @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 without PBR materials""" 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) self.metallicRoughness = None # Clear PBR materials 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: # Check if we need to align vertices to UV coordinates if self.vt is not None and self.v.shape[0] != self.vt.shape[0]: self._align_vertices_to_uv() # 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 # Add UV coordinates if available if self.vt is not None: uv_data = self.vt.detach().cpu().numpy().astype(np.float32) uv_accessor, _ = append_accessor_and_bufferview( doc, uv_data, component_type=5126, # FLOAT element_type="VEC2", target=34962 # ARRAY_BUFFER ) primitive["attributes"]["TEXCOORD_0"] = uv_accessor # Add texture if available if self.albedo is not None and self.vt is not None: self._add_texture_to_gltf(doc, primitive) # 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 _add_texture_to_gltf(self, doc, primitive): """Add texture data to GLTF document using fast approach""" try: # Convert albedo texture to PNG bytes albedo_np = self.albedo.detach().cpu().numpy() # Ensure values are in [0, 1] range albedo_np = np.clip(albedo_np, 0.0, 1.0) albedo_uint8 = (albedo_np * 255).astype(np.uint8) # Encode as PNG import cv2 albedo_bgr = cv2.cvtColor(albedo_uint8, cv2.COLOR_RGB2BGR) success, png_bytes = cv2.imencode('.png', albedo_bgr) if not success: print(f"[FastMesh] Failed to encode texture as PNG") return png_data = png_bytes.tobytes() # Add image data to buffer using helper function image_accessor, _ = append_accessor_and_bufferview( doc, np.frombuffer(png_data, dtype=np.uint8), component_type=5121, # UNSIGNED_BYTE element_type="SCALAR", target=None # No specific target for image data ) # Create GLTF structures gltf_json = doc.json() # Initialize arrays if they don't exist if "materials" not in gltf_json: gltf_json["materials"] = [] if "textures" not in gltf_json: gltf_json["textures"] = [] if "images" not in gltf_json: gltf_json["images"] = [] if "samplers" not in gltf_json: gltf_json["samplers"] = [] # Add image image_idx = len(gltf_json["images"]) gltf_json["images"].append({ "bufferView": image_accessor, "mimeType": "image/png" }) # Add sampler sampler_idx = len(gltf_json["samplers"]) gltf_json["samplers"].append({ "magFilter": 9729, # LINEAR "minFilter": 9987, # LINEAR_MIPMAP_LINEAR "wrapS": 10497, # REPEAT "wrapT": 10497 # REPEAT }) # Add texture texture_idx = len(gltf_json["textures"]) gltf_json["textures"].append({ "sampler": sampler_idx, "source": image_idx }) # Add material material_idx = len(gltf_json["materials"]) material = { "pbrMetallicRoughness": { "baseColorTexture": { "index": texture_idx, "texCoord": 0 }, "metallicFactor": 0.0, "roughnessFactor": 1.0 }, "alphaMode": "OPAQUE", "doubleSided": True } # Add metallic-roughness texture if available if self.metallicRoughness is not None: mr_np = self.metallicRoughness.detach().cpu().numpy() mr_uint8 = (mr_np * 255).astype(np.uint8) mr_bgr = cv2.cvtColor(mr_uint8, cv2.COLOR_RGB2BGR) mr_success, mr_png_bytes = cv2.imencode('.png', mr_bgr) if mr_success: mr_png_data = mr_png_bytes.tobytes() mr_image_accessor, _ = append_accessor_and_bufferview( doc, np.frombuffer(mr_png_data, dtype=np.uint8), component_type=5121, # UNSIGNED_BYTE element_type="SCALAR", target=None ) # Add MR image and texture mr_image_idx = len(gltf_json["images"]) gltf_json["images"].append({ "bufferView": mr_image_accessor, "mimeType": "image/png" }) mr_texture_idx = len(gltf_json["textures"]) gltf_json["textures"].append({ "sampler": sampler_idx, # Reuse same sampler "source": mr_image_idx }) # Update material material["pbrMetallicRoughness"]["metallicRoughnessTexture"] = { "index": mr_texture_idx, "texCoord": 0 } material["pbrMetallicRoughness"]["metallicFactor"] = 1.0 material["pbrMetallicRoughness"]["roughnessFactor"] = 1.0 gltf_json["materials"].append(material) # Link material to primitive primitive["material"] = material_idx except Exception as e: print(f"[FastMesh] Error adding texture to GLTF: {e}") def _align_vertices_to_uv(self): """Align vertices to UV coordinates (like align_v_to_vt in standard Mesh)""" try: # Use the standard Mesh method for alignment 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 ) # Apply alignment temp_mesh.align_v_to_vt() # Copy back aligned data 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 except Exception as e: print(f"[FastMesh] Error during vertex-UV alignment: {e}") 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 with MTL support""" vertices = self.v.detach().cpu().numpy() faces = self.f.detach().cpu().numpy() # Prepare file paths 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") # Get UV and normal data if available uvs = self.vt.detach().cpu().numpy() if self.vt is not None else None normals = self.vn.detach().cpu().numpy() if self.vn is not None else None uv_faces = self.ft.detach().cpu().numpy() if self.ft is not None else None normal_faces = self.fn.detach().cpu().numpy() if self.fn is not None else None # Write OBJ file with open(path, 'w') as f: if self.albedo is not None: f.write(f"mtllib {os.path.basename(mtl_path)}\n") # Write vertices for v in vertices: f.write(f"v {v[0]} {v[1]} {v[2]}\n") # Write UV coordinates if uvs is not None: for uv in uvs: f.write(f"vt {uv[0]} {1.0 - uv[1]}\n") # Flip V back for OBJ # Write normals if normals is not None: for n in normals: f.write(f"vn {n[0]} {n[1]} {n[2]}\n") # Write material usage if self.albedo is not None: f.write("usemtl defaultMat\n") # Write faces for i, face in enumerate(faces): face_str = "f" for j in range(3): v_idx = face[j] + 1 uv_idx = uv_faces[i][j] + 1 if uv_faces is not None else "" n_idx = normal_faces[i][j] + 1 if normal_faces is not None else "" if uvs is not None and normals is not None: face_str += f" {v_idx}/{uv_idx}/{n_idx}" elif uvs is not None: face_str += f" {v_idx}/{uv_idx}" elif normals is not None: face_str += f" {v_idx}//{n_idx}" else: face_str += f" {v_idx}" f.write(face_str + "\n") # Write MTL file if we have textures if self.albedo is not None: self._write_mtl_file(mtl_path, albedo_path, metallic_path, roughness_path) # Save texture files self._save_texture_files(albedo_path, metallic_path, roughness_path) def _write_mtl_file(self, mtl_path, albedo_path, metallic_path, roughness_path): """Write MTL file for OBJ""" with open(mtl_path, 'w') as f: f.write("newmtl defaultMat\n") f.write("Ka 1 1 1\n") f.write("Kd 1 1 1\n") f.write("Ks 1 1 1\n") f.write("illum 1\n") f.write("Ns 10\n") if self.albedo is not None: f.write(f"map_Kd {os.path.basename(albedo_path)}\n") if self.metallicRoughness is not None: f.write(f"map_Pm {os.path.basename(metallic_path)}\n") f.write(f"map_Pr {os.path.basename(roughness_path)}\n") def _save_texture_files(self, albedo_path, metallic_path, roughness_path): """Save texture files for OBJ""" try: # Save albedo texture if self.albedo is not None: albedo_np = self.albedo.detach().cpu().numpy() albedo_uint8 = (np.clip(albedo_np, 0, 1) * 255).astype(np.uint8) albedo_bgr = cv2.cvtColor(albedo_uint8, cv2.COLOR_RGB2BGR) cv2.imwrite(albedo_path, albedo_bgr) # Save metallic/roughness textures if self.metallicRoughness is not None: mr_np = self.metallicRoughness.detach().cpu().numpy() mr_uint8 = (np.clip(mr_np, 0, 1) * 255).astype(np.uint8) # Extract metallic (B channel) and roughness (G channel) metallic_uint8 = mr_uint8[:, :, 2] # B channel roughness_uint8 = mr_uint8[:, :, 1] # G channel cv2.imwrite(metallic_path, metallic_uint8) cv2.imwrite(roughness_path, roughness_uint8) except Exception as e: print(f"[FastMesh] Error saving texture files: {e}") 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) @classmethod def create_glb_with_pbr_materials(cls, obj_path, textures_dict, output_path, device=None): """Create GLB with PBR materials from separate texture files""" if device is None: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") try: mesh = cls.load(obj_path, **{'device': device}) if mesh is None: raise ValueError(f"Failed to load mesh from {obj_path}") # Combine metallic and roughness if both provided if "metallic" in textures_dict and "roughness" in textures_dict: mr_texture = cls._combine_metallic_roughness( textures_dict["metallic"], textures_dict["roughness"] ) mesh.metallicRoughness = torch.from_numpy(mr_texture.astype(np.float32) / 255.0).to(device).contiguous() # Load albedo texture if "albedo" in textures_dict and os.path.exists(textures_dict["albedo"]): albedo_img = np.array(Image.open(textures_dict["albedo"]).convert("RGB")) mesh.albedo = torch.from_numpy(albedo_img.astype(np.float32) / 255.0).to(device).contiguous() # Generate UV if not present if mesh.vt is None: mesh.auto_uv() # Write GLB mesh.write(output_path) print(f"PBR GLB file saved: {output_path}") except Exception as e: print(f"Error creating GLB with PBR materials: {e}") raise e @staticmethod def _combine_metallic_roughness(metallic_path, roughness_path): """Combine metallic and roughness maps into single texture""" metallic_img = Image.open(metallic_path).convert("L") roughness_img = Image.open(roughness_path).convert("L") if metallic_img.size != roughness_img.size: roughness_img = roughness_img.resize(metallic_img.size) width, height = metallic_img.size metallic_array = np.array(metallic_img) roughness_array = np.array(roughness_img) # Create combined array (R, G, B) = (AO, Roughness, Metallic) combined_array = np.zeros((height, width, 3), dtype=np.uint8) combined_array[:, :, 0] = 255 # R: AO (white if no AO map) combined_array[:, :, 1] = roughness_array # G: Roughness combined_array[:, :, 2] = metallic_array # B: Metallic return combined_array @classmethod def obj_to_glb_converter(cls, obj_path, output_path, device=None): """Convert OBJ to GLB using FastMesh direct parsing""" if device is None: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") try: # Load OBJ using direct parsing mesh = cls._load_obj_direct(obj_path, device=device) if mesh is None: raise ValueError(f"Failed to load OBJ file: {obj_path}") # Generate UV if not present if mesh.vt is None and mesh.albedo is not None: mesh.auto_uv() # Write as GLB mesh.write(output_path) return True except Exception as e: print(f"[FastMesh] OBJ to GLB conversion failed: {e}") return False @classmethod def read_obj_structure(cls, path): """Read OBJ file structure like your original code""" vertices, normals, uvs, faces = [], [], [], [] mtl_path = None with open(path, 'r', encoding='utf-8') as f: for line in f: line = line.strip() if not line or line.startswith('#'): continue parts = line.split() if not parts: continue prefix = parts[0] if prefix == 'v': vertices.append([float(parts[i]) for i in (1, 2, 3)]) elif prefix == 'vn': normals.append([float(parts[i]) for i in (1, 2, 3)]) elif prefix == 'vt': u, v = float(parts[1]), float(parts[2]) uvs.append([u, 1.0 - v]) # Flip V for glTF elif prefix == 'f': face_indices = [] for vert_str in parts[1:]: indices = vert_str.split('/') vi = int(indices[0]) - 1 if indices[0] else 0 ti = int(indices[1]) - 1 if len(indices) > 1 and indices[1] else 0 ni = int(indices[2]) - 1 if len(indices) > 2 and indices[2] else 0 face_indices.append((vi, ti, ni)) # Triangulate (fan triangulation) for i in range(1, len(face_indices) - 1): faces.extend([face_indices[0], face_indices[i], face_indices[i + 1]]) elif prefix == 'mtllib': mtl_path = os.path.join(os.path.dirname(path), parts[1]) # Parse MTL file textures = {} base_color = [1.0, 1.0, 1.0] if mtl_path and os.path.exists(mtl_path): try: with open(mtl_path, 'r', encoding='utf-8') as f: for line in f: line = line.strip() if not line or line.startswith('#'): continue parts = line.split() if not parts: continue if parts[0] == 'Kd': base_color = [float(parts[i]) for i in (1, 2, 3)] elif parts[0] == 'map_Kd': textures['baseColor'] = os.path.join(os.path.dirname(mtl_path), parts[1]) elif parts[0] == 'map_Pm': textures['metallic'] = os.path.join(os.path.dirname(mtl_path), parts[1]) elif parts[0] == 'map_Pr': textures['roughness'] = os.path.join(os.path.dirname(mtl_path), parts[1]) elif parts[0] == 'map_Bump': if '-bm' in parts: idx = parts.index('-bm') textures['normal'] = os.path.join(os.path.dirname(mtl_path), parts[idx+2]) else: textures['normal'] = os.path.join(os.path.dirname(mtl_path), parts[1]) except Exception as e: print(f"[FastMesh] Error parsing MTL: {e}") # Deduplicate vertices pos, normal, uv, indices, original_vi = cls._deduplicate_obj_vertices_with_mapping( vertices, normals, uvs, faces ) return pos, normal, uv, indices, original_vi, textures, base_color @classmethod def _deduplicate_obj_vertices_with_mapping(cls, vertices, normals, uvs, faces): """Deduplicate vertices and keep original vertex mapping""" pos, normal, uv = [], [], [] vertex_map = {} indices = [] original_vi = [] # Maps new vertex index to original vertex index for vi, ti, ni in faces: # Create unique key for this vertex combination key = (vi, ti if ti < len(uvs) else -1, ni if ni < len(normals) else -1) if key not in vertex_map: vertex_map[key] = len(pos) # Add vertex position pos.append(vertices[vi]) original_vi.append(vi) # Track original vertex index # Add UV coordinates if ti < len(uvs): uv.append(uvs[ti]) else: uv.append([0.0, 0.0]) # Add normal if ni < len(normals): normal.append(normals[ni]) else: normal.append([0.0, 0.0, 1.0]) indices.append(vertex_map[key]) return pos, normal, uv, indices, original_vi