""" dual_grid_view.py - Multi-view transform dual grid processing Extends dual_grid.py with scale and mesh rotation logic Based on test_ovoxel_transform.py implementation """ import os import sys import importlib import argparse import json import pandas as pd import numpy as np import torch import pickle import o_voxel from easydict import EasyDict as edict from functools import partial from utils import get_new_camera_matrix, transform_mesh, sphere_normalize_torch def _dual_grid_mesh_view(file, sha256, mesh_dump_root, transform_root, root, view_indices=None): """ Process multi-view dual grid conversion for a single sha256. Args: file: local_path from metadata sha256: sha256 string mesh_dump_root: directory containing mesh dump files transform_root: directory containing transform json files root: output directory for dual grids view_indices: list of view indices to process, None for all views """ try: pack = {'sha256': sha256} vertices_sphere = None sphere_radius = None faces = None # Load transforms transform_path = os.path.join(transform_root, sha256, 'transforms.json') if not os.path.exists(transform_path): print(f'Transform file not found for {sha256}, skipping') return {'sha256': sha256, 'error': 'Transform file not found'} with open(transform_path, 'r') as f: transforms_json = json.load(f) transform_mats = transforms_json['frames'] # Determine views to process if view_indices is None: view_indices = list(range(len(transform_mats))) else: view_indices = [i for i in view_indices if i < len(transform_mats)] # Track processed and skipped counts processed_count = 0 skipped_count = 0 for view_idx in view_indices: for res in opt.resolution: need_process = False # Check if already processed # Path structure: dual_grid_view_{res}/{sha256}/view{idx:02d}.vxz sha256_dir = os.path.join(root, f'dual_grid_view_{res}', sha256) vxz_path = os.path.join(sha256_dir, f'view{view_idx:02d}.vxz') if os.path.exists(vxz_path): try: info = o_voxel.io.read_vxz_info(vxz_path) pack[f'dual_grid_view{view_idx:02d}_converted_{res}'] = True pack[f'dual_grid_view{view_idx:02d}_size_{res}'] = info['num_voxel'] skipped_count += 1 except Exception as e: print(f'Error reading {sha256}/view{view_idx:02d}.vxz: {e}') need_process = True else: need_process = True # Process mesh if need_process: # Lazy load mesh if vertices_sphere is None: mesh_file = os.path.join(mesh_dump_root, 'mesh_dumps', f'{sha256}.pickle') if not os.path.exists(mesh_file): print(f'Mesh dump not found for {sha256}, skipping') return {'sha256': sha256, 'error': 'Mesh dump not found'} with open(mesh_file, 'rb') as f: dump = pickle.load(f) start = 0 vertices_list = [] faces_list = [] for obj in dump['objects']: if obj['vertices'].size == 0 or obj['faces'].size == 0: continue vertices_list.append(obj['vertices']) faces_list.append(obj['faces'] + start) start += len(obj['vertices']) if len(vertices_list) == 0: print(f'No valid mesh data for {sha256}, skipping') return {'sha256': sha256, 'error': 'No valid mesh data'} vertices = torch.from_numpy(np.concatenate(vertices_list, axis=0)).float().contiguous() faces = torch.from_numpy(np.concatenate(faces_list, axis=0)).long().contiguous() # Sphere normalization (for multi-view transform) - CPU only vertices_sphere, sphere_center, sphere_radius = sphere_normalize_torch(vertices) # Get transform for current view transform = transform_mats[view_idx] # Multi-view transform - CPU only transformed_vertices = transform_mesh(vertices_sphere, transform) # Post-transform normalization: scale by abs max to [-0.5, 0.5]^3 # Only scale, no center shift, to preserve relative model position abs_max = transformed_vertices.abs().max().item() box_scale = 0.49999 / abs_max # Normalize to [-0.5, 0.5] range transformed_normalized = transformed_vertices * box_scale transformed_normalized_cpu = transformed_normalized.contiguous() # Compute total scale (from original mesh to final normalized mesh) total_scale = box_scale / sphere_radius.item() # Validate range assert torch.all(transformed_normalized_cpu >= -0.5) and torch.all(transformed_normalized_cpu <= 0.5), \ f'vertices out of range for {sha256} view {view_idx}' # Ensure vertices and faces are on CPU with correct types and contiguous memory # CPU only, consistent with process_dual_grid in test_ovoxel_transform.py vertices_for_grid = transformed_normalized_cpu.float().contiguous() faces_for_grid = faces.long().contiguous() data_for_grid = {'vertices': vertices_for_grid, 'faces': faces_for_grid} # Dual grid encoding voxel_indices, dual_vertices, intersected = o_voxel.convert.mesh_to_flexible_dual_grid( **data_for_grid, grid_size=res, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], face_weight=1.0, boundary_weight=0.2, regularization_weight=1e-2, timing=False, ) # Convert to intra-voxel offsets and quantize dual_vertices = dual_vertices.float() voxel_indices_float = voxel_indices.float() dual_vertices = dual_vertices * res - voxel_indices_float assert torch.all(dual_vertices >= -1e-3) and torch.all(dual_vertices <= 1+1e-3), \ f'dual_vertices out of range for {sha256} view {view_idx}' dual_vertices = torch.clamp(dual_vertices, 0, 1) dual_vertices = (dual_vertices * 255).type(torch.uint8) intersected = (intersected[:, 0:1] + 2 * intersected[:, 1:2] + 4 * intersected[:, 2:3]).type(torch.uint8) # Save .vxz file os.makedirs(sha256_dir, exist_ok=True) o_voxel.io.write_vxz( vxz_path, voxel_indices, {'vertices': dual_vertices, 'intersected': intersected}, ) # Save scale info scale_path = os.path.join(sha256_dir, f'view{view_idx:02d}_scale.json') scale_info = { 'sha256': sha256, 'view_idx': view_idx, 'total_scale': total_scale, 'sphere_radius': sphere_radius.item(), 'box_scale': box_scale, } with open(scale_path, 'w') as f: json.dump(scale_info, f, indent=2) pack[f'dual_grid_view{view_idx:02d}_converted_{res}'] = True pack[f'dual_grid_view{view_idx:02d}_size_{res}'] = len(voxel_indices) pack[f'dual_grid_view{view_idx:02d}_scale_{res}'] = total_scale processed_count += 1 # Record processing stats pack['_processed_count'] = processed_count pack['_skipped_count'] = skipped_count return pack except Exception as e: print(f'Error processing {sha256}: {e}') import traceback traceback.print_exc() return {'sha256': sha256, 'error': str(e)} if __name__ == '__main__': dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}') parser = argparse.ArgumentParser() parser.add_argument('--root', type=str, required=True, help='Directory to save the metadata') parser.add_argument('--mesh_dump_root', type=str, default=None, help='Directory to load mesh dumps') parser.add_argument('--transform_root', type=str, default=None, help='Directory to load transform json files (renders_cond)') parser.add_argument('--dual_grid_root', type=str, default=None, help='Directory to save dual grids') parser.add_argument('--filter_low_aesthetic_score', type=float, default=None, help='Filter objects with aesthetic score lower than this value') 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') dataset_utils.add_args(parser) parser.add_argument('--rank', type=int, default=0) parser.add_argument('--resolution', type=str, default='256') parser.add_argument('--world_size', type=int, default=1) parser.add_argument('--max_workers', type=int, default=0) opt = parser.parse_args(sys.argv[2:]) opt = edict(vars(opt)) opt.resolution = [int(x) for x in opt.resolution.split(',')] opt.mesh_dump_root = opt.mesh_dump_root or opt.root opt.transform_root = opt.transform_root or os.path.join(opt.root, 'renders_cond') opt.dual_grid_root = opt.dual_grid_root or opt.root # Parse view_indices view_indices = None if opt.view_indices is not None: view_indices = [] for part in opt.view_indices.split(','): if '-' in part: start, end = map(int, part.split('-')) view_indices.extend(range(start, end + 1)) else: view_indices.append(int(part)) view_indices = list(set(view_indices)) # Deduplicate view_indices.sort() for res in opt.resolution: os.makedirs(os.path.join(opt.dual_grid_root, f'dual_grid_view_{res}', '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')): metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256')) if os.path.exists(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')): metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')).set_index('sha256')) # Check already processed dual_grid_view for res in opt.resolution: if os.path.exists(os.path.join(opt.dual_grid_root, f'dual_grid_view_{res}', 'metadata.csv')): dual_grid_metadata = pd.read_csv(os.path.join(opt.dual_grid_root, f'dual_grid_view_{res}', 'metadata.csv')).set_index('sha256') metadata = metadata.combine_first(dual_grid_metadata) 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] metadata = metadata[metadata['mesh_dumped'] == True] # Filter out objects with all views already processed if view_indices is not None: for res in opt.resolution: # Check if each specified view is already processed all_views_done_col = f'_all_views_done_{res}' metadata[all_views_done_col] = True for view_idx in view_indices: col_name = f'dual_grid_view{view_idx:02d}_converted_{res}' if col_name in metadata.columns: metadata[all_views_done_col] = metadata[all_views_done_col] & (metadata[col_name] == True) else: metadata[all_views_done_col] = False break # Keep objects with at least one incomplete resolution any_incomplete = None for res in opt.resolution: all_views_done_col = f'_all_views_done_{res}' if all_views_done_col in metadata.columns: if any_incomplete is None: any_incomplete = ~metadata[all_views_done_col] else: any_incomplete = any_incomplete | ~metadata[all_views_done_col] if any_incomplete is not None: before_filter = len(metadata) metadata = metadata[any_incomplete] print(f'Filtered out {before_filter - len(metadata)} already completed objects') 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] print(f'Processing {len(metadata)} objects...') if view_indices: print(f'View indices to process: {view_indices}') else: print('Processing all available views') # Process objects func = partial(_dual_grid_mesh_view, root=opt.dual_grid_root, mesh_dump_root=opt.mesh_dump_root, transform_root=opt.transform_root, view_indices=view_indices) dual_grids = dataset_utils.foreach_instance(metadata, opt.root, func, max_workers=opt.max_workers, desc='Dual griding views', timeout=300) # Processing summary total_processed = dual_grids['_processed_count'].sum() if '_processed_count' in dual_grids.columns else 0 total_skipped = dual_grids['_skipped_count'].sum() if '_skipped_count' in dual_grids.columns else 0 print(f'\n========== Processing Summary ==========') print(f'Total processed (new): {int(total_processed)}') print(f'Total skipped (existing): {int(total_skipped)}') print(f'Total items: {int(total_processed + total_skipped)}') print(f'=========================================\n') if 'error' in dual_grids.columns: errors = dual_grids[dual_grids['error'].notna()] if len(errors) > 0: with open('errors_view.txt', 'w') as f: f.write('\n'.join(errors['sha256'].tolist())) print(f'Errors written to errors_view.txt ({len(errors)} errors)') # Save metadata for res in opt.resolution: # Collect all view-related columns view_cols = [col for col in dual_grids.columns if f'dual_grid_view' in col and f'_{res}' in col and 'converted' in col] if view_cols: # Save metadata for each view dual_grid_metadata = dual_grids[dual_grids[view_cols].any(axis=1)] if len(dual_grid_metadata) > 0: # Save simplified metadata cols_to_save = ['sha256'] + [col for col in dual_grids.columns if f'_{res}' in col] cols_to_save = [col for col in cols_to_save if col in dual_grids.columns] dual_grid_metadata[cols_to_save].to_csv( os.path.join(opt.dual_grid_root, f'dual_grid_view_{res}', 'new_records', f'part_{opt.rank}.csv'), index=False ) print('Done!')