Spaces:
Running on Zero
Running on Zero
| """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] | |