Spaces:
Running on Zero
Running on Zero
| import logging as log | |
| import os.path as osp | |
| import imageio | |
| import matplotlib | |
| matplotlib.use('Agg') | |
| import matplotlib.colors as mcolors | |
| import numpy as np | |
| import torch | |
| from sklearn.neighbors import KDTree | |
| from tqdm import tqdm | |
| from third_party.TRELLIS.trellis.representations.mesh.cube2mesh import MeshExtractResult | |
| from third_party.TRELLIS.trellis.renderers.mesh_renderer import MeshRenderer | |
| from third_party.TRELLIS.trellis.utils.render_utils import yaw_pitch_r_fov_to_extrinsics_intrinsics | |
| def map_voxel_labels_to_vertices(mesh_vertices, voxel_coords, voxel_labels): | |
| """Assign each mesh vertex the label of its nearest voxel via KD-tree lookup. | |
| Args: | |
| mesh_vertices: np.ndarray (V, 3) — mesh vertex positions | |
| voxel_coords: np.ndarray (M, 3) — voxel center positions (same coordinate space) | |
| voxel_labels: np.ndarray (M,) — integer label per voxel | |
| Returns: | |
| np.ndarray (V,) — integer label per vertex | |
| """ | |
| tree = KDTree(voxel_coords) | |
| _, idx = tree.query(mesh_vertices, k=1) | |
| return voxel_labels[idx.flatten()] | |
| def visualize_and_save(structure, partfield_clusters, output_folder, | |
| output_name='cluster_visualization.mp4', | |
| num_frames=300, resolution=512): | |
| """Render a smooth rotating orbit video of the structure colored by cluster labels. | |
| Uses the TRELLIS MeshRenderer (nvdiffrast, ssaa=4) and the same yaw/pitch | |
| camera trajectory as the pipeline's Gaussian scene videos so all sides are visible. | |
| Args: | |
| structure: trimesh.Trimesh — normalized structure mesh | |
| (vertices in [-0.5, 0.5]) | |
| partfield_clusters: np.ndarray (V,) — per-vertex integer cluster labels, | |
| where V matches structure.vertices.shape[0] | |
| output_folder: str — directory in which to save the video | |
| output_name: str — filename for the saved video | |
| (default 'cluster_visualization.mp4') | |
| num_frames: int — number of video frames (default 300) | |
| resolution: int — output resolution in pixels (default 512) | |
| """ | |
| partfield_clusters = np.asarray(partfield_clusters, dtype=int) | |
| num_clusters = int(partfield_clusters.max()) + 1 | |
| # Evenly-spaced HSV hues → maximally distinct colors | |
| hues = np.linspace(0, 1, num_clusters, endpoint=False) | |
| cluster_colors = np.array([mcolors.hsv_to_rgb([h, 0.90, 0.95]) for h in hues], | |
| dtype=np.float32) # (num_clusters, 3) | |
| vertex_colors = cluster_colors[partfield_clusters] # (V, 3) | |
| # Build MeshExtractResult for the TRELLIS MeshRenderer | |
| verts_t = torch.tensor(structure.vertices, dtype=torch.float32, device='cuda') | |
| faces_t = torch.tensor(structure.faces, dtype=torch.long, device='cuda') | |
| colors_t = torch.tensor(vertex_colors, dtype=torch.float32, device='cuda') | |
| mesh_result = MeshExtractResult(vertices=verts_t, faces=faces_t, vertex_attrs=colors_t) | |
| # Camera trajectory: identical to render_utils.render_video | |
| yaws = torch.linspace(0, 2 * torch.pi, num_frames).tolist() | |
| pitch = (0.25 + 0.5 * torch.sin(torch.linspace(0, 2 * torch.pi, num_frames))).tolist() | |
| extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitch, 2, 40) | |
| renderer = MeshRenderer(rendering_options={ | |
| "resolution": resolution, | |
| "near": 1, | |
| "far": 100, | |
| "ssaa": 4, | |
| }) | |
| frames = [] | |
| log.info(f"Rendering {num_frames} frames → {output_name} ...") | |
| with torch.no_grad(): | |
| for extr, intr in tqdm(zip(extrinsics, intrinsics), total=num_frames, | |
| desc=f'Rendering {output_name}'): | |
| res = renderer.render(mesh_result, extr, intr, | |
| return_types=["color", "mask"]) | |
| # composite cluster colors over black background | |
| frame = (res['color'] * res['mask']).permute(1, 2, 0).clamp(0.0, 1.0) | |
| frames.append((frame.cpu().numpy() * 255).astype(np.uint8)) | |
| video_path = osp.join(output_folder, output_name) | |
| imageio.mimsave(video_path, frames, fps=30) | |
| log.info(f"Saved {video_path}") | |