# Copyright 2026 Toyota Research Institute. All rights reserved. import os import sys import numpy as np import multiprocessing import argparse from tqdm import tqdm from glob import glob from functools import partial from anydata.converters.utils import list_s3_recursive, list_s3, get_splits from anydata.sync.sync_utils import ( get_seqs_subset, aws_s3_cp, aws_s3_sync, extract_tar, strip_s3_prefix, ) from anydata.utils.misc import get_local_root from anydata.utils.read import read_json from rich.console import Console console = Console() ####################################################### LOWDIM_LABELS = ['intrinsics','extrinsics','action','language'] ####################################################### def parse_args(dataset=None): parser = argparse.ArgumentParser() parser.add_argument("dataset", type=str) parser.add_argument("--num_procs", type=int, default=16) parser.add_argument("--s3_bucket", type=str, default='s3://tri-ml-sandbox-16011-us-west-2-datasets') parser.add_argument("--data_folder", type=str, default='cv_unified') parser.add_argument("--local_bucket", type=str, default=get_local_root()) parser.add_argument('--metadata_only', action='store_true') parser.add_argument('--labels', type=str, nargs="+", default=None) parser.add_argument('--subset', type=str, default=None) parser.add_argument("--src", type=str, default=None) parser.add_argument("--dst", type=str, default=None) parser.add_argument("--verbose", action='store_true') parser.add_argument("--skip_existing", action='store_true', help='Skip episodes already present locally (idempotent resume)') mode = parser.add_mutually_exclusive_group(required=True) mode.add_argument('--frames', action='store_const', dest='mode', const='frames') mode.add_argument('--videos', action='store_const', dest='mode', const='videos') args = parser.parse_args() if dataset is not None: args.dataset = dataset args.dataset = f'{args.mode}/{args.dataset}' if args.src is None: args.src = f'{args.s3_bucket}/{args.data_folder}/{args.dataset}' if args.dst is None: args.dst = f'{args.local_bucket}/{args.data_folder}/{args.dataset}' if args.dataset.endswith('.json'): args.dst = os.path.dirname(args.dst) return args ####################################################### def download_files(files, dst): if files is None: return # Do nothing if there are no files os.makedirs(dst, exist_ok=True) for file in files: # Download each file separately aws_s3_cp(file, f'{dst}/{os.path.basename(file)}') def download_sequences(i, seqs, args): n_downloaded = 0 n_skipped = 0 progress = tqdm(seqs, ncols=96, leave=False) for seq in progress: result = download_sequence(i, seq, args) if result == 'skipped': n_skipped += 1 if hasattr(args, 'verbose') and args.verbose: console.print(f'[dim] [Thread {i+1}] SKIP {seq}[/dim]') else: n_downloaded += 1 if hasattr(args, 'verbose') and args.verbose: console.print(f'[green] [Thread {i+1}] DONE {seq}[/green]') progress.set_description(f'Thread {i+1}/{args.num_procs} downloaded={n_downloaded} skipped={n_skipped}') def download_sequence(i, seq, args): ### Prepare paths (seqs are relative to args.src / args.dst) s3_path = f'{args.src}/{seq}' local_path = f'{args.dst}/{seq}' if seq.endswith('.json') or seq.endswith('.txt'): ### Skip if already present locally if args.skip_existing and os.path.exists(local_path): return 'skipped' ### Download json/txt file os.makedirs(os.path.dirname(local_path), exist_ok=True) aws_s3_cp(s3_path, local_path) elif seq.endswith('.tar.gz'): ### Skip label if not in requested label list mid_folder = local_path.replace('.tar.gz', '') label = os.path.basename(mid_folder) if args.labels is not None and label not in args.labels: return 'skipped' ### Skip if already extracted # NOTE(bvh): extracted folder presence is reliable — extract_tar runs after download. if args.skip_existing and os.path.exists(mid_folder): return 'skipped' ### Download and extract tar file os.makedirs(os.path.dirname(local_path), exist_ok=True) aws_s3_cp(s3_path, local_path) extract_tar(local_path, mid_folder) else: if args.skip_existing and os.path.exists(local_path): files = glob(f'{local_path}/*') files_tar = glob(f'{local_path}/*.tar.gz') if len(files) > 0 and len(files_tar) == 0: # Check if there are leftover tarfiles return 'done' ### Sync entire episode directory (aws s3 sync is already idempotent) os.makedirs(local_path, exist_ok=True) aws_s3_sync(s3_path, f'{local_path}/', robust=True) ### Extract and delete all tarfiles from sequence for tf in glob(f'{local_path}/*.tar.gz'): tarfolder = tf.replace(".tar.gz", "") extract_tar(tf, tarfolder) return 'done' ####################################################### def download(args): if args.src.endswith('.json'): ### Download split JSON to local dst split_name = os.path.basename(args.src) shared_name = "metadata_shared.json" local_split = f'{args.dst}/{split_name}' local_shared = f'{args.dst}/{shared_name}' if not os.path.exists(local_split): os.makedirs(args.dst, exist_ok=True) aws_s3_cp(args.src, local_split) aws_s3_cp(args.src.replace(split_name, shared_name), local_shared, check=False) ### Get labels and sequences data = read_json(local_split) raw_seqs = list(data['sequences'].keys()) # labels = [l for l in labels if l not in LOWDIM_LABELS] # labels.append('lowdim') # Check if split requires certain labels if args.labels is None: if "filters" in data and "labels" in data['filters']: args.labels = data['filters']['labels'] # Always download lowdim (maybe change that later) if args.labels is not None: args.labels.append('lowdim') ### Seqs from split are relative to parent of dataset dir ### So move src up one level to match; dst stays as-is (user controls it) args.src = os.path.dirname(args.src) seqs = list(raw_seqs) files = None else: ### Get sequences (list_s3_recursive returns bucket-relative paths, ### so strip prefix to make them relative to args.src) suffix = ['metadata.json'] # All sequences have metadata.json raw = list_s3_recursive(args.src, suffix=suffix) # List all files recursively seqs = strip_s3_prefix(raw, args.src) if not args.metadata_only: # If it's only metadata, target these files directly seqs = sorted(list(set([os.path.dirname(s) for s in seqs]))) ### Include root files (splits and shared metadata) files = list_s3(args.src, suffix='.json') files = [f'{args.src}/{f}' for f in files] seqs = get_seqs_subset(seqs, args.subset) from anydata.sync.sync_utils import multi_thread multi_thread = partial(multi_thread, name='DOWNLOADING UNIFIED', fn_seqs=download_sequences, fn_files=download_files, tarfiles=True) multi_thread(seqs, files, args) if args.subset is not None: command = f'python anydata/converters/misc/cleanup.py {os.path.basename(args.dst)} --official' if args.mode == 'frames': command += ' --frames' if args.mode == 'videos': command += ' --videos' os.system(command) ####################################################### if __name__ == '__main__': args = parse_args() if '*' in args.dataset: # Loop over available datasets if wildcard "*" is used dataset = os.path.dirname(args.dataset) split = os.path.basename(args.dataset).split('*') folders = list_s3(os.path.dirname(args.src)) if split[-1] == '': folders = [f for f in folders if os.path.basename(f).startswith(split[-2])] if split[0] == '': folders = [f for f in folders if os.path.basename(f).endswith(split[1])] for folder in folders: args = parse_args(dataset=f'{dataset}/{folder}') download(args) else: download(args) #######################################################