import os import json from typing import * import numpy as np import torch import utils3d from .. import models from .components import ImageConditionedMixin, ViewImageConditionedMixin from ..modules.sparse import SparseTensor from .structured_latent import SLatVisMixin, SLat from ..utils.render_utils import get_renderer, yaw_pitch_r_fov_to_extrinsics_intrinsics from ..utils.data_utils import load_balanced_group_indices class SLatShapeVisMixin(SLatVisMixin): def _loading_slat_dec(self): if self.slat_dec is not None: return if self.slat_dec_path is not None: cfg = json.load(open(os.path.join(self.slat_dec_path, 'config.json'), 'r')) decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args']) ckpt_path = os.path.join(self.slat_dec_path, 'ckpts', f'decoder_{self.slat_dec_ckpt}.pt') decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True)) else: decoder = models.from_pretrained(self.pretrained_slat_dec) decoder.set_resolution(self.resolution) self.slat_dec = decoder.cuda().eval() @torch.no_grad() def visualize_sample( self, x_0: Union[SparseTensor, dict], camera_angle_x: Optional[torch.Tensor] = None, camera_distance: Optional[torch.Tensor] = None, mesh_scale: Optional[torch.Tensor] = None, ): """ Visualize shape samples. Args: x_0: SparseTensor or dict containing 'x_0' camera_angle_x: Optional [B] camera FOV angle in radians camera_distance: Optional [B] camera distance for GT view rendering mesh_scale: Optional [B] mesh scale factor for coordinate alignment Returns: dict with: 'multiview': [B, 3, 1024, 1024] - 4 fixed views rendered in 2x2 grid (normal) 'gt_view': [B, 3, 512, 512] - GT camera view (if camera params provided) """ x_0 = x_0 if isinstance(x_0, SparseTensor) else x_0['x_0'] reps = self.decode_latent(x_0.cuda()) # build fixed camera views (4 views: 0°, 90°, 180°, 270°) yaw = [0, np.pi/2, np.pi, 3*np.pi/2] yaw_offset = -16 / 180 * np.pi yaw = [y + yaw_offset for y in yaw] pitch = [20 / 180 * np.pi for _ in range(4)] fixed_exts, fixed_ints = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaw, pitch, 2, 30) # Check if we have GT camera parameters for GT view rendering has_gt_camera = ( camera_angle_x is not None and camera_distance is not None and mesh_scale is not None ) # render renderer = get_renderer(reps[0]) multiview_images = [] gt_view_images = [] for i, representation in enumerate(reps): # Render 4 fixed views (2x2 grid) image = torch.zeros(3, 1024, 1024).cuda() tile = [2, 2] # Validate mesh data before rasterization verts = representation.vertices faces = representation.faces if verts.shape[0] == 0 or faces.shape[0] == 0: print(f"[visualize_sample] Warning: sample {i} has empty mesh, skipping") multiview_images.append(image) continue if faces.max() >= verts.shape[0]: print(f"[visualize_sample] Warning: sample {i} has out-of-bound face indices " f"(max face idx={faces.max().item()}, num verts={verts.shape[0]}), skipping") multiview_images.append(image) continue if torch.isnan(verts).any() or torch.isinf(verts).any(): print(f"[visualize_sample] Warning: sample {i} has NaN/Inf vertices, skipping") multiview_images.append(image) continue try: for j, (ext, intr) in enumerate(zip(fixed_exts, fixed_ints)): res = renderer.render(representation, ext, intr) image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['normal'] except RuntimeError as e: print(f"[visualize_sample] Warning: render failed for sample {i}: {e}") image = torch.zeros(3, 1024, 1024).cuda() multiview_images.append(image) # Render GT camera view using the fixed front view (same as sparse_structure_latent.py) if has_gt_camera: # The GT view should match exactly how ProjGrid projects 3D points to 2D. # # In image_conditioned_proj.py (ProjGrid.forward): # 1. grid_points are in [-1, 1]^3 (from torch.linspace(-1, 1, res)) # 2. grid_points are rotated by rotation_matrix (Y-Z swap): x'=x, y'=-z, z'=y # 3. grid_points are scaled: grid_points / mesh_scale / 2 # 4. Points are projected using front_view_transform_matrix with distance # # Mesh vertices are in [-0.5, 0.5]^3. To match ProjGrid's coordinate space, # we need to scale them: vertices / mesh_scale -> [-0.5/s, 0.5/s]^3 # This is equivalent to ProjGrid's: [-1,1]^3 / scale / 2 -> [-0.5/s, 0.5/s]^3 # # Camera position: ProjGrid camera is at (0, -distance, 0) in Blender coords (Z-up). # After inverse rotation to mesh space, camera is at (0, 0, distance). scale = mesh_scale[i].item() distance = camera_distance[i].item() fov = camera_angle_x[i].item() device = representation.vertices.device # Scale mesh vertices to match ProjGrid's projection space from ..representations import Mesh scaled_rep = Mesh( vertices=representation.vertices / scale, faces=representation.faces, ) cam_pos = torch.tensor([0.0, 0.0, distance], device=device) look_at = torch.tensor([0.0, 0.0, 0.0], device=device) cam_up = torch.tensor([0.0, 1.0, 0.0], device=device) gt_ext = utils3d.torch.extrinsics_look_at(cam_pos, look_at, cam_up) gt_int = utils3d.torch.intrinsics_from_fov_xy( torch.tensor(fov, device=device), torch.tensor(fov, device=device) ) gt_ext = gt_ext.to(device) gt_int = gt_int.to(device) # Use scaled mesh renderer with appropriate near/far for smaller mesh mesh_half_size = 0.5 / scale renderer.rendering_options.near = max(0.01, distance - mesh_half_size - 0.5) renderer.rendering_options.far = distance + mesh_half_size + 0.5 try: gt_res = renderer.render(scaled_rep, gt_ext, gt_int) gt_view_images.append(gt_res['normal']) except RuntimeError as e: print(f"[visualize_sample] Warning: GT view render failed for sample {i}: {e}") gt_view_images.append(torch.full((3, 512, 512), 0.5, device=device)) result = { 'multiview': torch.stack(multiview_images), } if has_gt_camera and len(gt_view_images) > 0: result['gt_view'] = torch.stack(gt_view_images) return result class SLatShape(SLatShapeVisMixin, SLat): """ structured latent for shape generation Args: roots (str): path to the dataset resolution (int): resolution of the shape min_aesthetic_score (float): minimum aesthetic score max_tokens (int): maximum number of tokens latent_key (str): key of the latent to be used normalization (dict): normalization stats pretrained_slat_dec (str): name of the pretrained slat decoder slat_dec_path (str): path to the slat decoder, if given, will override the pretrained_slat_dec slat_dec_ckpt (str): name of the slat decoder checkpoint skip_list (str, optional): path to a file containing sha256 hashes to skip skip_aesthetic_score_datasets (list, optional): list of dataset names to skip aesthetic score check """ def __init__(self, roots: str, *, resolution: int, min_aesthetic_score: float = 5.0, max_tokens: int = 32768, normalization: Optional[dict] = None, pretrained_slat_dec: str = 'microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16', slat_dec_path: Optional[str] = None, slat_dec_ckpt: Optional[str] = None, skip_list: Optional[str] = None, skip_aesthetic_score_datasets: Optional[list] = None, ): super().__init__( roots, min_aesthetic_score=min_aesthetic_score, max_tokens=max_tokens, latent_key='shape_latent', normalization=normalization, pretrained_slat_dec=pretrained_slat_dec, slat_dec_path=slat_dec_path, slat_dec_ckpt=slat_dec_ckpt, skip_list=skip_list, skip_aesthetic_score_datasets=skip_aesthetic_score_datasets, ) self.resolution = resolution class ImageConditionedSLatShape(ImageConditionedMixin, SLatShape): """ Image conditioned structured latent for shape generation """ pass class SLatShapeView(SLatShapeVisMixin, SLat): """ View-based structured latent for shape generation. Data format: {sha256}/view{XX}.npz where each npz contains 'coords' and 'feats' keys. Args: roots (str): path to the dataset resolution (int): resolution of the shape min_aesthetic_score (float): minimum aesthetic score max_tokens (int): maximum number of tokens num_views (int): Number of views to use (0 to num_views-1). Default is 2. normalization (dict): normalization stats pretrained_slat_dec (str): name of the pretrained slat decoder slat_dec_path (str): path to the slat decoder, if given, will override the pretrained_slat_dec slat_dec_ckpt (str): name of the slat decoder checkpoint skip_list (str, optional): path to a file containing sha256 hashes to skip skip_aesthetic_score_datasets (list, optional): list of dataset names to skip aesthetic score check """ def __init__(self, roots: str, *, resolution: int, min_aesthetic_score: float = 5.0, max_tokens: int = 32768, num_views: int = 2, normalization: Optional[dict] = None, pretrained_slat_dec: str = 'microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16', slat_dec_path: Optional[str] = None, slat_dec_ckpt: Optional[str] = None, skip_list: Optional[str] = None, skip_aesthetic_score_datasets: Optional[list] = None, ): self.normalization = normalization self.min_aesthetic_score = min_aesthetic_score self.max_tokens = max_tokens self.num_views = num_views self.latent_key = 'shape_latent' self.value_range = (0, 1) # Initialize parent with SLatVisMixin parameters from .components import StandardDatasetBase SLatVisMixin.__init__( self, roots, pretrained_slat_dec=pretrained_slat_dec, slat_dec_path=slat_dec_path, slat_dec_ckpt=slat_dec_ckpt, ) StandardDatasetBase.__init__(self, roots, skip_list=skip_list, skip_aesthetic_score_datasets=skip_aesthetic_score_datasets) self.resolution = resolution # Calculate loads for load balancing self.loads = [] for _, sha256, _ in self.instances: if f'{self.latent_key}_tokens' in self.metadata.columns: try: self.loads.append(self.metadata.loc[sha256, f'{self.latent_key}_tokens']) except: self.loads.append(self.max_tokens) else: self.loads.append(self.max_tokens) if self.normalization is not None: self.mean = torch.tensor(self.normalization['mean']).reshape(1, -1) self.std = torch.tensor(self.normalization['std']).reshape(1, -1) def filter_metadata(self, metadata, dataset_name=None): stats = {} # View-based shape_latent uses columns like shape_latent_view00_encoded, shape_latent_view01_encoded, etc. required_view_cols = [f'shape_latent_view{i:02d}_encoded' for i in range(self.num_views)] existing_view_cols = [col for col in required_view_cols if col in metadata.columns] if existing_view_cols: # Filter rows where all required views are encoded # Note: NaN should be treated as False, so use == True for explicit comparison has_all_views = (metadata[existing_view_cols] == True).all(axis=1) metadata = metadata[has_all_views] stats[f'With {self.num_views} view latents'] = len(metadata) else: # Fallback: check shape_latent_encoded column if f'{self.latent_key}_encoded' in metadata.columns: metadata = metadata[metadata[f'{self.latent_key}_encoded'] == True] stats['With latent'] = len(metadata) # Skip aesthetic score check for specified datasets (e.g., texverse) or if column doesn't exist skip_aesthetic = ( (dataset_name and dataset_name.lower() in [d.lower() for d in self.skip_aesthetic_score_datasets]) or ('aesthetic_score' not in metadata.columns) ) if skip_aesthetic: stats[f'Aesthetic score check skipped'] = len(metadata) else: metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score] stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata) # Filter by max_tokens if column exists tokens_col = f'{self.latent_key}_tokens' if tokens_col in metadata.columns: metadata = metadata[metadata[tokens_col] <= self.max_tokens] stats[f'Num tokens <= {self.max_tokens}'] = len(metadata) return metadata, stats def get_instance(self, root, instance): # View-based format: directory with view{XX}.npz files latent_dir = os.path.join(root[self.latent_key], instance) # Randomly select a view from the configured range view_idx = np.random.randint(0, self.num_views) view_file = f'view{view_idx:02d}.npz' # Store view info for ViewImageConditionedMixin self._current_view_idx = view_idx self._current_latent_dir = latent_dir data = np.load(os.path.join(latent_dir, view_file)) coords = torch.tensor(data['coords']).int() feats = torch.tensor(data['feats']).float() if self.normalization is not None: feats = (feats - self.mean) / self.std return { 'coords': coords, 'feats': feats, 'view_idx': view_idx, } @staticmethod def collate_fn(batch, split_size=None): if split_size is None: group_idx = [list(range(len(batch)))] else: group_idx = load_balanced_group_indices([b['coords'].shape[0] for b in batch], split_size) packs = [] for group in group_idx: sub_batch = [batch[i] for i in group] pack = {} coords = [] feats = [] layout = [] start = 0 for i, b in enumerate(sub_batch): coords.append(torch.cat([torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32), b['coords']], dim=-1)) feats.append(b['feats']) layout.append(slice(start, start + b['coords'].shape[0])) start += b['coords'].shape[0] coords = torch.cat(coords) feats = torch.cat(feats) pack['x_0'] = SparseTensor( coords=coords, feats=feats, ) pack['x_0']._shape = torch.Size([len(group), *sub_batch[0]['feats'].shape[1:]]) pack['x_0'].register_spatial_cache('layout', layout) # collate other data keys = [k for k in sub_batch[0].keys() if k not in ['coords', 'feats']] for k in keys: if isinstance(sub_batch[0][k], torch.Tensor): pack[k] = torch.stack([b[k] for b in sub_batch]) elif isinstance(sub_batch[0][k], list): pack[k] = sum([b[k] for b in sub_batch], []) else: pack[k] = [b[k] for b in sub_batch] packs.append(pack) if split_size is None: return packs[0] return packs class ViewImageConditionedSLatShapeView(ViewImageConditionedMixin, SLatShapeView): """ Image-conditioned view-based structured latent for shape generation. Loads shape_latent from {sha256}/view{XX}.npz format and pairs with corresponding view from render_cond. Uses ViewImageConditionedMixin which reads mesh_scale from view{XX}_scale.json. """ pass