import os import sys sys.path.append(os.path.join(os.path.dirname(__file__), '..')) import json import shutil import argparse import torch import numpy as np import pandas as pd from tqdm import tqdm from easydict import EasyDict as edict from concurrent.futures import ThreadPoolExecutor from queue import Queue from utils import parse_view_indices import pixal3d.models as models torch.set_grad_enabled(False) def clear_cuda_error(): torch.cuda.synchronize() torch.cuda.empty_cache() if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--root', type=str, required=True, help='Directory to save the metadata') parser.add_argument('--shape_latent_root', type=str, default=None, help='Directory containing the shape latent files') parser.add_argument('--ss_latent_root', type=str, default=None, help='Directory to save the ss latent files') parser.add_argument('--filter_low_aesthetic_score', type=float, default=None, help='Filter objects with aesthetic score lower than this value') parser.add_argument('--resolution', type=int, default=32, help='SS latent resolution') parser.add_argument('--shape_latent_name', type=str, required=True, help='Name of the shape latent files (e.g., shape_enc_next_dc_f16c32_fp16_512)') parser.add_argument('--enc_pretrained', type=str, default='microsoft/TRELLIS-image-large/ckpts/ss_enc_conv3d_16l8_fp16', help='Pretrained encoder model') parser.add_argument('--model_root', type=str, help='Root directory of models') parser.add_argument('--enc_model', type=str, help='Encoder model. if specified, use this model instead of pretrained model') parser.add_argument('--ckpt', type=str, help='Checkpoint to load') parser.add_argument('--instances', type=str, default=None, help='Instances to process') parser.add_argument('--view_indices', type=str, default=None, help='View indices to process, e.g., "0,1,2" or "0-5". None for all views') parser.add_argument('--num_views', type=int, default=24, help='Total number of views (used when view_indices is None)') parser.add_argument('--rank', type=int, default=0) parser.add_argument('--world_size', type=int, default=1) opt = parser.parse_args() opt = edict(vars(opt)) opt.shape_latent_root = opt.shape_latent_root or opt.root opt.ss_latent_root = opt.ss_latent_root or opt.root # Parse view_indices view_indices = parse_view_indices(opt.view_indices) if view_indices is None: view_indices = list(range(opt.num_views)) print(f'View indices to process: {view_indices}') if opt.enc_model is None: latent_name = f'{opt.enc_pretrained.split("/")[-1]}_{opt.resolution}' encoder = models.from_pretrained(opt.enc_pretrained).eval().cuda() else: latent_name = f'{opt.enc_model.split("/")[-1]}_{opt.ckpt}_{opt.resolution}' cfg = edict(json.load(open(os.path.join(opt.model_root, opt.enc_model, 'config.json'), 'r'))) encoder = getattr(models, cfg.models.encoder.name)(**cfg.models.encoder.args).cuda() ckpt_path = os.path.join(opt.model_root, opt.enc_model, 'ckpts', f'encoder_{opt.ckpt}.pt') encoder.load_state_dict(torch.load(ckpt_path), strict=False) encoder.eval() print(f'Loaded model from {ckpt_path}') # Multi-view shape_latent and ss_latent directory names shape_latent_view_name = f'{opt.shape_latent_name}_view' ss_latent_view_name = f'{latent_name}_view' os.makedirs(os.path.join(opt.ss_latent_root, 'ss_latents', ss_latent_view_name, 'new_records'), exist_ok=True) # Get file list if not os.path.exists(os.path.join(opt.root, 'metadata.csv')): raise ValueError('metadata.csv not found') metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256') if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')): aesthetic_metadata = pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256') metadata = metadata.join(aesthetic_metadata, how='left', rsuffix='_aesthetic') # Check shape_latent_view metadata shape_latent_view_metadata_path = os.path.join(opt.shape_latent_root, 'shape_latents', shape_latent_view_name, 'metadata.csv') if os.path.exists(shape_latent_view_metadata_path): shape_latent_view_metadata = pd.read_csv(shape_latent_view_metadata_path).set_index('sha256') metadata = metadata.join(shape_latent_view_metadata, how='left', rsuffix='_shape_latent_view') print(f'Loaded shape_latent_view metadata with {len(shape_latent_view_metadata)} records') else: print(f'Warning: shape_latent_view metadata not found at {shape_latent_view_metadata_path}') # Check ss_latent_view metadata (used to skip already completed tasks) ss_latent_view_metadata_path = os.path.join(opt.ss_latent_root, 'ss_latents', ss_latent_view_name, 'metadata.csv') if os.path.exists(ss_latent_view_metadata_path): ss_latent_view_metadata = pd.read_csv(ss_latent_view_metadata_path).set_index('sha256') metadata = metadata.join(ss_latent_view_metadata, how='left', rsuffix='_ss_latent_view') print(f'Loaded ss_latent_view metadata with {len(ss_latent_view_metadata)} records') else: print(f'Warning: ss_latent_view metadata not found at {ss_latent_view_metadata_path}') metadata = metadata.reset_index() if opt.instances is None: if opt.filter_low_aesthetic_score is not None: metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score] # Filter to objects that have shape_latent_view data # Use first view as indicator first_view_col = f'shape_latent_view{view_indices[0]:02d}_encoded' if first_view_col in metadata.columns: metadata = metadata[metadata[first_view_col] == True] else: print(f'Warning: Column {first_view_col} not found in metadata, will check files directly') else: if os.path.exists(opt.instances): with open(opt.instances, 'r') as f: instances = f.read().splitlines() else: instances = opt.instances.split(',') metadata = metadata[metadata['sha256'].isin(instances)] start = len(metadata) * opt.rank // opt.world_size end = len(metadata) * (opt.rank + 1) // opt.world_size metadata = metadata[start:end] records = [] # Build task list: (sha256, view_idx), filter already completed tasks via metadata tasks = [] skipped_count = 0 # Pre-fetch completion status columns for each view encoded_cols = {view_idx: f'ss_latent_view{view_idx:02d}_encoded' for view_idx in view_indices} for _, row in metadata.iterrows(): sha256 = row['sha256'] for view_idx in view_indices: encoded_col = encoded_cols[view_idx] # Check if already marked as completed in metadata if encoded_col in metadata.columns and row.get(encoded_col, False) == True: skipped_count += 1 continue tasks.append((sha256, view_idx)) print(f'Total tasks: {len(tasks) + skipped_count}, Already done (from metadata): {skipped_count}, To process: {len(tasks)}') load_queue = Queue(maxsize=32) with ThreadPoolExecutor(max_workers=32) as loader_executor, \ ThreadPoolExecutor(max_workers=32) as saver_executor: def loader(task): sha256, view_idx = task try: # Check if output file already exists, skip if so output_path = os.path.join( opt.ss_latent_root, 'ss_latents', ss_latent_view_name, sha256, f'view{view_idx:02d}.npz' ) if os.path.exists(output_path): load_queue.put((sha256, view_idx, None)) return # shape_latent_view path: shape_latents/{shape_latent_view_name}/{sha256}/view{idx:02d}.npz npz_path = os.path.join( opt.shape_latent_root, 'shape_latents', shape_latent_view_name, sha256, f'view{view_idx:02d}.npz' ) if not os.path.exists(npz_path): print(f"[Loader Skip] {sha256}/view{view_idx:02d}: npz file not found at {npz_path}") load_queue.put((sha256, view_idx, None)) return data = np.load(npz_path) coords = data['coords'] # Validate coords are within resolution range assert np.all(coords < opt.resolution), f"{sha256}/view{view_idx:02d}: Invalid coords (max={coords.max()}, resolution={opt.resolution})" coords = torch.from_numpy(coords).long() ss = torch.zeros(1, opt.resolution, opt.resolution, opt.resolution, dtype=torch.long) ss[:, coords[:, 0], coords[:, 1], coords[:, 2]] = 1 load_queue.put((sha256, view_idx, ss)) except Exception as e: print(f"[Loader Error] {sha256}/view{view_idx:02d}: {e}") load_queue.put((sha256, view_idx, None)) loader_executor.map(loader, tasks) def saver(sha256, view_idx, pack): sha256_dir = os.path.join(opt.ss_latent_root, 'ss_latents', ss_latent_view_name, sha256) os.makedirs(sha256_dir, exist_ok=True) save_path = os.path.join(sha256_dir, f'view{view_idx:02d}.npz') np.savez_compressed(save_path, **pack) # Copy scale.json from shape_latent_view directory src_scale_path = os.path.join( opt.shape_latent_root, 'shape_latents', shape_latent_view_name, sha256, f'view{view_idx:02d}_scale.json' ) dst_scale_path = os.path.join(sha256_dir, f'view{view_idx:02d}_scale.json') if os.path.exists(src_scale_path) and not os.path.exists(dst_scale_path): shutil.copy2(src_scale_path, dst_scale_path) records.append({ 'sha256': sha256, f'ss_latent_view{view_idx:02d}_encoded': True, }) for _ in tqdm(range(len(tasks)), desc="Extracting SS view latents"): try: sha256, view_idx, ss = load_queue.get() if ss is None: continue ss = ss.cuda()[None].float() z = encoder(ss, sample_posterior=False) torch.cuda.synchronize() if not torch.isfinite(z).all(): print(f"[Skip] {sha256}/view{view_idx:02d}: Non-finite latent") clear_cuda_error() continue pack = { 'z': z[0].cpu().numpy(), } saver_executor.submit(saver, sha256, view_idx, pack) except Exception as e: print(f"[Error] {sha256}/view{view_idx:02d}: {e}") clear_cuda_error() continue saver_executor.shutdown(wait=True) records = pd.DataFrame.from_records(records) records.to_csv(os.path.join(opt.ss_latent_root, 'ss_latents', ss_latent_view_name, 'new_records', f'part_{opt.rank}.csv'), index=False)