"""Preserve original faces, UVs, material channels and scene transforms when splitting.""" import copy from pathlib import Path import numpy as np from PIL import Image import trimesh CHANNELS = ("baseColorTexture", "metallicRoughnessTexture", "normalTexture", "occlusionTexture", "emissiveTexture") def load_mesh(path: Path) -> trimesh.Trimesh: scene = trimesh.load(path, force="scene", process=False) meshes = [] for node in scene.graph.nodes_geometry: transform, geometry = scene.graph[node] mesh = scene.geometry[geometry].copy(include_cache=True) mesh.apply_transform(transform) meshes.append(mesh) if not meshes: raise ValueError("The generated asset contains no geometry.") # TRELLIS exports one atlas. Refuse ambiguous material merging rather than losing channels. if len(meshes) != 1: raise ValueError("Expected one TRELLIS atlas mesh before segmentation.") mesh = meshes[0] if not len(mesh.faces) or not np.isfinite(mesh.vertices).all(): raise ValueError("The generated geometry is empty or non-finite.") if mesh.visual.kind != "texture" or mesh.visual.uv is None: raise ValueError("The generated mesh is missing UVs or textures.") if getattr(mesh.visual.material, "baseColorTexture", None) is None: raise ValueError("The generated mesh is missing its base-color texture.") return mesh def split(mesh, labels, folder): labels = np.asarray(labels) if labels.shape != (len(mesh.faces),) or not np.issubdtype(labels.dtype, np.integer): raise ValueError("Segmentation must return one integer label per original face.") folder.mkdir(parents=True, exist_ok=True) parts = [] for index, label in enumerate(np.unique(labels)): face_ids = np.flatnonzero(labels == label) vertex_ids, faces = np.unique(mesh.faces[face_ids], return_inverse=True) material = copy.deepcopy(mesh.visual.material) uv = mesh.visual.uv[vertex_ids].copy() # Crop each part's atlas bounds and remap UVs without rebaking geometry. # Out-of-range tiled UVs retain the atlas to preserve sampler meaning. if np.all((uv >= 0) & (uv <= 1)): image = material.baseColorTexture width, height = image.size x0 = max(0, int(np.floor(uv[:, 0].min() * width)) - 4) x1 = min(width, int(np.ceil(uv[:, 0].max() * width)) + 4) y0 = max(0, int(np.floor((1 - uv[:, 1].max()) * height)) - 4) y1 = min(height, int(np.ceil((1 - uv[:, 1].min()) * height)) + 4) if x1 > x0 and y1 > y0: for channel in CHANNELS: texture = getattr(material, channel, None) if texture is not None: w, h = texture.size # Common normalized crop requires matching atlas dimensions. texture = texture.resize((width, height), Image.Resampling.BILINEAR) if (w, h) != (width, height) else texture setattr(material, channel, texture.crop((x0, y0, x1, y1))) uv[:, 0] = (uv[:, 0] * width - x0) / (x1 - x0) uv[:, 1] = (uv[:, 1] * height - (height - y1)) / (y1 - y0) part = trimesh.Trimesh(vertices=mesh.vertices[vertex_ids], faces=faces.reshape(-1, 3), vertex_normals=mesh.vertex_normals[vertex_ids], visual=trimesh.visual.TextureVisuals(uv=uv, material=material), process=False) part_id = f"part-{index:03d}" part.metadata["name"] = part_id directory = folder / part_id directory.mkdir() np.save(directory / "source-faces.npy", face_ids) part.export(directory / "source.glb") for channel in CHANNELS: texture = getattr(material, channel, None) if texture is not None: texture.save(directory / f"{channel}.png") parts.append((part_id, part)) if sum(len(part.faces) for _, part in parts) != len(mesh.faces): raise ValueError("Segmentation lost source faces.") return parts def export(parts, path): scene = trimesh.Scene() for part_id, part in parts: label = part.metadata.get("name", part_id) scene.add_geometry(part, node_name=f"{part_id} {label}" if label != part_id else part_id, geom_name=part_id) scene.export(path) check = trimesh.load(path, force="scene", process=False) if len(check.geometry) != len(parts): raise ValueError("GLB round trip lost components.") for mesh in check.geometry.values(): if mesh.visual.kind != "texture" or getattr(mesh.visual.material, "baseColorTexture", None) is None: raise ValueError("GLB round trip lost textures.") return {"components": len(parts), "faces": sum(len(p.faces) for _, p in parts), "vertices": sum(len(p.vertices) for _, p in parts), "export_reopened": True} def apply_refinement(parts, plan): by_id = dict(parts) edits = [edit.part_id for edit in plan.parts] if len(set(edits)) != len(edits) or any(key not in by_id for key in edits): raise ValueError("Refinement contains duplicate or unknown component IDs.") result = [(key, part.copy(include_cache=True)) for key, part in parts] for key, part in result: normals = part.vertex_normals.copy() if plan.remove_degenerate_faces: part.update_faces(part.nondegenerate_faces()) if plan.remove_duplicate_faces: part.update_faces(part.unique_faces()) if not len(part.faces): raise ValueError(f"Refinement would erase {key}.") referenced = np.unique(part.faces) part.remove_unreferenced_vertices() part.vertex_normals = normals[referenced] part.visual.material = copy.deepcopy(part.visual.material) for edit in plan.parts: if edit.part_id != key: continue part.metadata["name"] = edit.name material = part.visual.material if edit.roughness is not None: material.roughnessFactor = edit.roughness if edit.metallic is not None: material.metallicFactor = edit.metallic rgba = np.array(material.baseColorTexture.convert("RGBA")) rgba[:, :, :3] = np.clip(rgba[:, :, :3].astype(float) * edit.color_gain, 0, 255).astype(np.uint8) material.baseColorTexture = Image.fromarray(rgba) return result def describe(parts): return [{"part_id": key, "faces": len(part.faces), "bounds": part.bounds.tolist(), "texture_size": list(part.visual.material.baseColorTexture.size)} for key, part in parts]