Spaces:
Running
on
Zero
Running
on
Zero
| import os | |
| import shutil | |
| import sys | |
| import time | |
| import importlib | |
| import argparse | |
| import numpy as np | |
| import pandas as pd | |
| from tqdm import tqdm | |
| from easydict import EasyDict as edict | |
| from concurrent.futures import ThreadPoolExecutor | |
| import utils3d | |
| def get_first_directory(path): | |
| with os.scandir(path) as it: | |
| for entry in it: | |
| if entry.is_dir(): | |
| return entry.name | |
| return None | |
| def need_process(key): | |
| return key in opt.field or opt.field == ['all'] | |
| if __name__ == '__main__': | |
| dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}') | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--output_dir', type=str, required=True, | |
| help='Directory to save the metadata') | |
| parser.add_argument('--field', type=str, default='all', | |
| help='Fields to process, separated by commas') | |
| parser.add_argument('--from_file', action='store_true', | |
| help='Build metadata from file instead of from records of processings.' + | |
| 'Useful when some processing fail to generate records but file already exists.') | |
| dataset_utils.add_args(parser) | |
| opt = parser.parse_args(sys.argv[2:]) | |
| opt = edict(vars(opt)) | |
| os.makedirs(opt.output_dir, exist_ok=True) | |
| os.makedirs(os.path.join(opt.output_dir, 'merged_records'), exist_ok=True) | |
| opt.field = opt.field.split(',') | |
| timestamp = str(int(time.time())) | |
| # get file list | |
| if os.path.exists(os.path.join(opt.output_dir, 'metadata.csv')): | |
| print('Loading previous metadata...') | |
| metadata = pd.read_csv(os.path.join(opt.output_dir, 'metadata.csv')) | |
| else: | |
| metadata = dataset_utils.get_metadata(**opt) | |
| metadata.set_index('sha256', inplace=True) | |
| # merge downloaded | |
| df_files = [f for f in os.listdir(opt.output_dir) if f.startswith('downloaded_') and f.endswith('.csv')] | |
| df_parts = [] | |
| for f in df_files: | |
| try: | |
| df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f))) | |
| except: | |
| pass | |
| if len(df_parts) > 0: | |
| df = pd.concat(df_parts) | |
| df.set_index('sha256', inplace=True) | |
| if 'local_path' in metadata.columns: | |
| metadata.update(df, overwrite=True) | |
| else: | |
| metadata = metadata.join(df, on='sha256', how='left') | |
| for f in df_files: | |
| shutil.move(os.path.join(opt.output_dir, f), os.path.join(opt.output_dir, 'merged_records', f'{timestamp}_{f}')) | |
| # detect models | |
| image_models = [] | |
| if os.path.exists(os.path.join(opt.output_dir, 'features')): | |
| image_models = os.listdir(os.path.join(opt.output_dir, 'features')) | |
| latent_models = [] | |
| if os.path.exists(os.path.join(opt.output_dir, 'latents')): | |
| latent_models = os.listdir(os.path.join(opt.output_dir, 'latents')) | |
| ss_latent_models = [] | |
| if os.path.exists(os.path.join(opt.output_dir, 'ss_latents')): | |
| ss_latent_models = os.listdir(os.path.join(opt.output_dir, 'ss_latents')) | |
| print(f'Image models: {image_models}') | |
| print(f'Latent models: {latent_models}') | |
| print(f'Sparse Structure latent models: {ss_latent_models}') | |
| if 'rendered' not in metadata.columns: | |
| metadata['rendered'] = [False] * len(metadata) | |
| if 'voxelized' not in metadata.columns: | |
| metadata['voxelized'] = [False] * len(metadata) | |
| if 'num_voxels' not in metadata.columns: | |
| metadata['num_voxels'] = [0] * len(metadata) | |
| if 'cond_rendered' not in metadata.columns: | |
| metadata['cond_rendered'] = [False] * len(metadata) | |
| for model in image_models: | |
| if f'feature_{model}' not in metadata.columns: | |
| metadata[f'feature_{model}'] = [False] * len(metadata) | |
| for model in latent_models: | |
| if f'latent_{model}' not in metadata.columns: | |
| metadata[f'latent_{model}'] = [False] * len(metadata) | |
| for model in ss_latent_models: | |
| if f'ss_latent_{model}' not in metadata.columns: | |
| metadata[f'ss_latent_{model}'] = [False] * len(metadata) | |
| # merge rendered | |
| df_files = [f for f in os.listdir(opt.output_dir) if f.startswith('rendered_') and f.endswith('.csv')] | |
| df_parts = [] | |
| for f in df_files: | |
| try: | |
| df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f))) | |
| except: | |
| pass | |
| if len(df_parts) > 0: | |
| df = pd.concat(df_parts) | |
| df.set_index('sha256', inplace=True) | |
| metadata.update(df, overwrite=True) | |
| for f in df_files: | |
| shutil.move(os.path.join(opt.output_dir, f), os.path.join(opt.output_dir, 'merged_records', f'{timestamp}_{f}')) | |
| # merge voxelized | |
| df_files = [f for f in os.listdir(opt.output_dir) if f.startswith('voxelized_') and f.endswith('.csv')] | |
| df_parts = [] | |
| for f in df_files: | |
| try: | |
| df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f))) | |
| except: | |
| pass | |
| if len(df_parts) > 0: | |
| df = pd.concat(df_parts) | |
| df.set_index('sha256', inplace=True) | |
| metadata.update(df, overwrite=True) | |
| for f in df_files: | |
| shutil.move(os.path.join(opt.output_dir, f), os.path.join(opt.output_dir, 'merged_records', f'{timestamp}_{f}')) | |
| # merge cond_rendered | |
| df_files = [f for f in os.listdir(opt.output_dir) if f.startswith('cond_rendered_') and f.endswith('.csv')] | |
| df_parts = [] | |
| for f in df_files: | |
| try: | |
| df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f))) | |
| except: | |
| pass | |
| if len(df_parts) > 0: | |
| df = pd.concat(df_parts) | |
| df.set_index('sha256', inplace=True) | |
| metadata.update(df, overwrite=True) | |
| for f in df_files: | |
| shutil.move(os.path.join(opt.output_dir, f), os.path.join(opt.output_dir, 'merged_records', f'{timestamp}_{f}')) | |
| # merge features | |
| for model in image_models: | |
| df_files = [f for f in os.listdir(opt.output_dir) if f.startswith(f'feature_{model}_') and f.endswith('.csv')] | |
| df_parts = [] | |
| for f in df_files: | |
| try: | |
| df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f))) | |
| except: | |
| pass | |
| if len(df_parts) > 0: | |
| df = pd.concat(df_parts) | |
| df.set_index('sha256', inplace=True) | |
| metadata.update(df, overwrite=True) | |
| for f in df_files: | |
| shutil.move(os.path.join(opt.output_dir, f), os.path.join(opt.output_dir, 'merged_records', f'{timestamp}_{f}')) | |
| # merge latents | |
| for model in latent_models: | |
| df_files = [f for f in os.listdir(opt.output_dir) if f.startswith(f'latent_{model}_') and f.endswith('.csv')] | |
| df_parts = [] | |
| for f in df_files: | |
| try: | |
| df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f))) | |
| except: | |
| pass | |
| if len(df_parts) > 0: | |
| df = pd.concat(df_parts) | |
| df.set_index('sha256', inplace=True) | |
| metadata.update(df, overwrite=True) | |
| for f in df_files: | |
| shutil.move(os.path.join(opt.output_dir, f), os.path.join(opt.output_dir, 'merged_records', f'{timestamp}_{f}')) | |
| # merge sparse structure latents | |
| for model in ss_latent_models: | |
| df_files = [f for f in os.listdir(opt.output_dir) if f.startswith(f'ss_latent_{model}_') and f.endswith('.csv')] | |
| df_parts = [] | |
| for f in df_files: | |
| try: | |
| df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f))) | |
| except: | |
| pass | |
| if len(df_parts) > 0: | |
| df = pd.concat(df_parts) | |
| df.set_index('sha256', inplace=True) | |
| metadata.update(df, overwrite=True) | |
| for f in df_files: | |
| shutil.move(os.path.join(opt.output_dir, f), os.path.join(opt.output_dir, 'merged_records', f'{timestamp}_{f}')) | |
| # build metadata from files | |
| if opt.from_file: | |
| with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor, \ | |
| tqdm(total=len(metadata), desc="Building metadata") as pbar: | |
| def worker(sha256): | |
| try: | |
| if need_process('rendered') and metadata.loc[sha256, 'rendered'] == False and \ | |
| os.path.exists(os.path.join(opt.output_dir, 'renders', sha256, 'transforms.json')): | |
| metadata.loc[sha256, 'rendered'] = True | |
| if need_process('voxelized') and metadata.loc[sha256, 'rendered'] == True and metadata.loc[sha256, 'voxelized'] == False and \ | |
| os.path.exists(os.path.join(opt.output_dir, 'voxels', f'{sha256}.ply')): | |
| try: | |
| pts = utils3d.io.read_ply(os.path.join(opt.output_dir, 'voxels', f'{sha256}.ply'))[0] | |
| metadata.loc[sha256, 'voxelized'] = True | |
| metadata.loc[sha256, 'num_voxels'] = len(pts) | |
| except Exception as e: | |
| pass | |
| if need_process('cond_rendered') and metadata.loc[sha256, 'cond_rendered'] == False and \ | |
| os.path.exists(os.path.join(opt.output_dir, 'renders_cond', sha256, 'transforms.json')): | |
| metadata.loc[sha256, 'cond_rendered'] = True | |
| for model in image_models: | |
| if need_process(f'feature_{model}') and \ | |
| metadata.loc[sha256, f'feature_{model}'] == False and \ | |
| metadata.loc[sha256, 'rendered'] == True and \ | |
| metadata.loc[sha256, 'voxelized'] == True and \ | |
| os.path.exists(os.path.join(opt.output_dir, 'features', model, f'{sha256}.npz')): | |
| metadata.loc[sha256, f'feature_{model}'] = True | |
| for model in latent_models: | |
| if need_process(f'latent_{model}') and \ | |
| metadata.loc[sha256, f'latent_{model}'] == False and \ | |
| metadata.loc[sha256, 'rendered'] == True and \ | |
| metadata.loc[sha256, 'voxelized'] == True and \ | |
| os.path.exists(os.path.join(opt.output_dir, 'latents', model, f'{sha256}.npz')): | |
| metadata.loc[sha256, f'latent_{model}'] = True | |
| for model in ss_latent_models: | |
| if need_process(f'ss_latent_{model}') and \ | |
| metadata.loc[sha256, f'ss_latent_{model}'] == False and \ | |
| metadata.loc[sha256, 'voxelized'] == True and \ | |
| os.path.exists(os.path.join(opt.output_dir, 'ss_latents', model, f'{sha256}.npz')): | |
| metadata.loc[sha256, f'ss_latent_{model}'] = True | |
| pbar.update() | |
| except Exception as e: | |
| print(f'Error processing {sha256}: {e}') | |
| pbar.update() | |
| executor.map(worker, metadata.index) | |
| executor.shutdown(wait=True) | |
| # statistics | |
| metadata.to_csv(os.path.join(opt.output_dir, 'metadata.csv')) | |
| num_downloaded = metadata['local_path'].count() if 'local_path' in metadata.columns else 0 | |
| with open(os.path.join(opt.output_dir, 'statistics.txt'), 'w') as f: | |
| f.write('Statistics:\n') | |
| f.write(f' - Number of assets: {len(metadata)}\n') | |
| f.write(f' - Number of assets downloaded: {num_downloaded}\n') | |
| f.write(f' - Number of assets rendered: {metadata["rendered"].sum()}\n') | |
| f.write(f' - Number of assets voxelized: {metadata["voxelized"].sum()}\n') | |
| if len(image_models) != 0: | |
| f.write(f' - Number of assets with image features extracted:\n') | |
| for model in image_models: | |
| f.write(f' - {model}: {metadata[f"feature_{model}"].sum()}\n') | |
| if len(latent_models) != 0: | |
| f.write(f' - Number of assets with latents extracted:\n') | |
| for model in latent_models: | |
| f.write(f' - {model}: {metadata[f"latent_{model}"].sum()}\n') | |
| if len(ss_latent_models) != 0: | |
| f.write(f' - Number of assets with sparse structure latents extracted:\n') | |
| for model in ss_latent_models: | |
| f.write(f' - {model}: {metadata[f"ss_latent_{model}"].sum()}\n') | |
| f.write(f' - Number of assets with captions: {metadata["captions"].count()}\n') | |
| f.write(f' - Number of assets with image conditions: {metadata["cond_rendered"].sum()}\n') | |
| with open(os.path.join(opt.output_dir, 'statistics.txt'), 'r') as f: | |
| print(f.read()) |