import os import copy import sys import importlib import argparse import pandas as pd import pickle import numpy as np import torch from easydict import EasyDict as edict from functools import partial import o_voxel def _pbr_voxelize(file, metadatum, pbr_dump_root, root): sha256 = metadatum['sha256'] try: pack = {'sha256': sha256} dump = None for res in opt.resolution: need_process = False # check if already processed if os.path.exists(os.path.join(root, f'pbr_voxels_{res}', f'{sha256}.vxz')): try: info = o_voxel.io.read_vxz_info(os.path.join(root, f'pbr_voxels_{res}', f'{sha256}.vxz')) pack[f'pbr_voxelized_{res}'] = True pack[f'num_pbr_voxels_{res}'] = info['num_voxel'] except Exception as e: print(f'Error reading {sha256}.vxz: {e}') need_process = True else: need_process = True # process if necessary if need_process: if dump == None: with open(os.path.join(pbr_dump_root, 'pbr_dumps', f'{sha256}.pickle'), 'rb') as f: dump = pickle.load(f) # Fix dump alpha map for mat in dump['materials']: if mat['alphaTexture'] is not None and mat['alphaMode'] == 'OPAQUE': mat['alphaMode'] = 'BLEND' dump['materials'].append({ "baseColorFactor": [0.8, 0.8, 0.8], "alphaFactor": 1.0, "metallicFactor": 0.0, "roughnessFactor": 0.5, "alphaMode": "OPAQUE", "alphaCutoff": 0.5, "baseColorTexture": None, "alphaTexture": None, "metallicTexture": None, "roughnessTexture": None, }) # append default material dump['objects'] = [ obj for obj in dump['objects'] if obj['vertices'].size != 0 and obj['faces'].size != 0 ] vertices = torch.from_numpy(np.concatenate([obj['vertices'] for obj in dump['objects']], axis=0)).float() vertices_min = vertices.min(dim=0)[0] vertices_max = vertices.max(dim=0)[0] center = (vertices_min + vertices_max) / 2 scale = 0.99999 / (vertices_max - vertices_min).max() for obj in dump['objects']: obj['vertices'] = (torch.from_numpy(obj['vertices']).float() - center) * scale obj['vertices'] = obj['vertices'].numpy() obj['mat_ids'][obj['mat_ids'] == -1] = len(dump['materials']) - 1 assert np.all(obj['mat_ids'] >= 0), 'invalid mat_ids' assert np.all(obj['vertices'] >= -0.5) and np.all(obj['vertices'] <= 0.5), 'vertices out of range' coord, attr = o_voxel.convert.blender_dump_to_volumetric_attr(dump, grid_size=res, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], mip_level_offset=0, verbose=False, timing=False) del attr['normal'] del attr['emissive'] o_voxel.io.write_vxz(os.path.join(root, f'pbr_voxels_{res}', f'{sha256}.vxz'), coord, attr) pack[f'pbr_voxelized_{res}'] = True pack[f'num_pbr_voxels_{res}'] = len(coord) return pack except Exception as e: print(f'Error voxelizing {sha256}: {e}') 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('--pbr_dump_root', type=str, default=None, help='Directory to load mesh dumps') parser.add_argument('--pbr_voxel_root', type=str, default=None, help='Directory to save voxelized pbr attributes') 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') dataset_utils.add_args(parser) parser.add_argument('--resolution', type=str, default=1024) parser.add_argument('--rank', type=int, default=0) 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 = sorted([int(x) for x in opt.resolution.split(',')], reverse=True) opt.pbr_dump_root = opt.pbr_dump_root or opt.root opt.pbr_voxel_root = opt.pbr_voxel_root or opt.root for res in opt.resolution: os.makedirs(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{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.pbr_dump_root, 'pbr_dumps', 'metadata.csv')): metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')).set_index('sha256')) for res in opt.resolution: if os.path.exists(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{res}', 'metadata.csv')): pbr_voxel_metadata = pd.read_csv(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{res}','metadata.csv')).set_index('sha256') pbr_voxel_metadata = pbr_voxel_metadata.rename(columns={'pbr_voxelized': f'pbr_voxelized_{res}', 'num_pbr_voxels': f'num_pbr_voxels_{res}'}) metadata = metadata.combine_first(pbr_voxel_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['pbr_dumped'] == True] mask = np.zeros(len(metadata), dtype=bool) for res in opt.resolution: if f'pbr_voxelized_{res}' in metadata.columns: mask |= metadata[f'pbr_voxelized_{res}'] != True else: mask[:] = True break metadata = metadata[mask] 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...') # process objects func = partial(_pbr_voxelize, pbr_dump_root=opt.pbr_dump_root, root=opt.pbr_voxel_root) pbr_voxelized = dataset_utils.foreach_instance(metadata, None, func, max_workers=opt.max_workers, no_file=True, desc='Voxelizing') if 'error' in pbr_voxelized.columns: errors = pbr_voxelized[pbr_voxelized['error'].notna()] with open('errors.txt', 'w') as f: f.write('\n'.join(errors['sha256'].tolist())) for res in opt.resolution: if f'pbr_voxelized_{res}' in pbr_voxelized.columns: pbr_voxel_metadata = pbr_voxelized[pbr_voxelized[f'pbr_voxelized_{res}'] == True] if len(pbr_voxel_metadata) > 0: pbr_voxel_metadata = pbr_voxel_metadata[['sha256', f'pbr_voxelized_{res}', f'num_pbr_voxels_{res}']] pbr_voxel_metadata = pbr_voxel_metadata.rename(columns={f'pbr_voxelized_{res}': 'pbr_voxelized', f'num_pbr_voxels_{res}': 'num_pbr_voxels'}) pbr_voxel_metadata.to_csv(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{res}', 'new_records', f'part_{opt.rank}.csv'), index=False)