Spaces:
Running on Zero
Running on Zero
File size: 6,748 Bytes
5c331a4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | """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]
|