1

dataset_toolkits/datasets/ObjaverseXL.py

@@ -0,0 +1,92 @@
+import os
+import argparse
+from concurrent.futures import ThreadPoolExecutor
+from tqdm import tqdm
+import pandas as pd
+import objaverse.xl as oxl
+from utils import get_file_hash
+
+
+def add_args(parser: argparse.ArgumentParser):
+    parser.add_argument('--source', type=str, default='sketchfab',
+                        help='Data source to download annotations from (github, sketchfab)')
+
+
+def get_metadata(source, **kwargs):
+    if source == 'sketchfab':
+        metadata = pd.read_csv("hf://datasets/JeffreyXiang/TRELLIS-500K/ObjaverseXL_sketchfab.csv")
+    elif source == 'github':
+        metadata = pd.read_csv("hf://datasets/JeffreyXiang/TRELLIS-500K/ObjaverseXL_github.csv")
+    else:
+        raise ValueError(f"Invalid source: {source}")
+    return metadata
+        
+
+def download(metadata, output_dir, **kwargs):    
+    os.makedirs(os.path.join(output_dir, 'raw'), exist_ok=True)
+
+    # download annotations
+    annotations = oxl.get_annotations()
+    annotations = annotations[annotations['sha256'].isin(metadata['sha256'].values)]
+    
+    # download and render objects
+    file_paths = oxl.download_objects(
+        annotations,
+        download_dir=os.path.join(output_dir, "raw"),
+        save_repo_format="zip",
+    )
+    
+    downloaded = {}
+    metadata = metadata.set_index("file_identifier")
+    for k, v in file_paths.items():
+        sha256 = metadata.loc[k, "sha256"]
+        downloaded[sha256] = os.path.relpath(v, output_dir)
+
+    return pd.DataFrame(downloaded.items(), columns=['sha256', 'local_path'])
+
+
+def foreach_instance(metadata, output_dir, func, max_workers=None, desc='Processing objects') -> pd.DataFrame:
+    import os
+    from concurrent.futures import ThreadPoolExecutor
+    from tqdm import tqdm
+    import tempfile
+    import zipfile
+    
+    # load metadata
+    metadata = metadata.to_dict('records')
+
+    # processing objects
+    records = []
+    max_workers = max_workers or os.cpu_count()
+    try:
+        with ThreadPoolExecutor(max_workers=max_workers) as executor, \
+            tqdm(total=len(metadata), desc=desc) as pbar:
+            def worker(metadatum):
+                try:
+                    local_path = metadatum['local_path']
+                    sha256 = metadatum['sha256']
+                    if local_path.startswith('raw/github/repos/'):
+                        path_parts = local_path.split('/')
+                        file_name = os.path.join(*path_parts[5:])
+                        zip_file = os.path.join(output_dir, *path_parts[:5])
+                        with tempfile.TemporaryDirectory() as tmp_dir:
+                            with zipfile.ZipFile(zip_file, 'r') as zip_ref:
+                                zip_ref.extractall(tmp_dir)
+                            file = os.path.join(tmp_dir, file_name)
+                            record = func(file, sha256)
+                    else:
+                        file = os.path.join(output_dir, local_path)
+                        record = func(file, sha256)
+                    if record is not None:
+                        records.append(record)
+                    pbar.update()
+                except Exception as e:
+                    print(f"Error processing object {sha256}: {e}")
+                    pbar.update()
+            
+            executor.map(worker, metadata)
+            executor.shutdown(wait=True)
+    except:
+        print("Error happened during processing.")
+        
+    return pd.DataFrame.from_records(records)

dataset_toolkits/stat_latent.py

@@ -0,0 +1,66 @@
+import os
+import json
+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
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--output_dir', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--model', type=str, default='dinov2_vitl14_reg_slat_enc_swin8_B_64l8_fp16',
+                        help='Latent model to use')
+    parser.add_argument('--num_samples', type=int, default=50000,
+                        help='Number of samples to use for calculating stats')
+    opt = parser.parse_args()
+    opt = edict(vars(opt))
+
+    # get file list
+    if os.path.exists(os.path.join(opt.output_dir, 'metadata.csv')):
+        metadata = pd.read_csv(os.path.join(opt.output_dir, 'metadata.csv'))
+    else:
+        raise ValueError('metadata.csv not found')
+    if opt.filter_low_aesthetic_score is not None:
+        metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+    metadata = metadata[metadata[f'latent_{opt.model}'] == True]
+    sha256s = metadata['sha256'].values
+    sha256s = np.random.choice(sha256s, min(opt.num_samples, len(sha256s)), replace=False)
+
+    # stats
+    means = []
+    mean2s = []
+    with ThreadPoolExecutor(max_workers=16) as executor, \
+        tqdm(total=len(sha256s), desc="Extracting features") as pbar:
+        def worker(sha256):
+            try:
+                feats = np.load(os.path.join(opt.output_dir, 'latents', opt.model, f'{sha256}.npz'))
+                feats = feats['feats']
+                means.append(feats.mean(axis=0))
+                mean2s.append((feats ** 2).mean(axis=0))
+                pbar.update()
+            except Exception as e:
+                print(f"Error extracting features for {sha256}: {e}")
+                pbar.update()
+
+        executor.map(worker, sha256s)
+        executor.shutdown(wait=True)
+
+    mean = np.array(means).mean(axis=0)
+    mean2 = np.array(mean2s).mean(axis=0)
+    std = np.sqrt(mean2 - mean ** 2)
+
+    print('mean:', mean)
+    print('std:', std)
+
+    with open(os.path.join(opt.output_dir, 'latents', opt.model, 'stats.json'), 'w') as f:
+        json.dump({
+            'mean': mean.tolist(),
+            'std': std.tolist(),
+        }, f, indent=4)
+