Spaces:
Running on Zero
Running on Zero
File size: 4,231 Bytes
43ca9e0 | 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 | 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}")
|