import os import sys import shutil import traceback import argparse import multiprocessing import numpy as np import torch from typing import Any, Dict, List from PIL import Image from copy import deepcopy from glob import glob, escape as glob_escape from tqdm import tqdm from threading import Thread from functools import partial from webdataset import TarWriter from rich.console import Console from anydata.converters.utils import write_stats, get_shared_info, get_num_frames from anydata.utils.misc import stack_sample from anydata.utils.read import read_config, read_json, read_npz_node from anydata.utils.write import create_folder, write_json, write_video_seq from anydata.geometry.depth import encode_dense from anydata.dataloaders.instantiate import dataset_from_config from anydata.dataloaders.Unified import UnifiedDataset from anydata.webdataset.utils import encode_latents, rgb_to_byte_array, mask_to_byte_array from anydata.sync.sync_utils import remove_path, aws_s3_cp, aws_s3_rm, rm_sync, mv_sync, Paginator # from anydata.augmentations.resize import parse_resize_params console = Console() def _episode_cameras(unified_path, split_path): """Return {episode_key: set_of_cameras} by reading each episode's metadata.json.""" split_data = read_json(split_path) sequences = split_data.get('sequences', {}) result = {} skipped = [] for seq_key in sequences: meta_file = os.path.join(unified_path, seq_key, 'metadata.json') if not os.path.exists(meta_file): skipped.append(seq_key) continue meta = read_json(meta_file) if meta is not None: result[seq_key] = set(meta.get('cameras', [])) if skipped: console.print(f'[yellow]WARNING: {len(skipped)}/{len(sequences)} episodes missing metadata.json ' f'(not downloaded?), skipping[/yellow]') return result def _filter_split(split_path, keep_keys, out_path): """Write a filtered copy of a split file keeping only the given episode keys.""" split_data = read_json(split_path) orig_seqs = split_data.get('sequences', {}) filtered = {k: v for k, v in orig_seqs.items() if k in keep_keys} split_data['sequences'] = filtered split_data['size'] = dict(split_data.get('size', {}), seqs=len(filtered)) write_json(out_path, split_data) return out_path def _rank_token(): """Stable token to avoid cross-rank temp split filename collisions. Returns empty string for single-process/local runs to keep filenames unchanged. """ world_size = int(os.environ.get("WORLD_SIZE", "1")) if world_size <= 1: return "" rank = os.environ.get("RANK") local_rank = os.environ.get("LOCAL_RANK") if rank is not None: return f"__r{rank}" if local_rank is not None: return f"__lr{local_rank}" return "__multi" def _generate_split_all(run_path, raw_config, args): """Generate a combined split JSON from all per-camera or flex-group manifests.""" source_split = os.path.basename(raw_config['path'][0]) prefix = 'split_all' # os.path.splitext(source_split)[0] if source_split.endswith('.json') else 'split_all' suffix = '' if args.subset is None else f'__{args.subset.replace("/","_")}' split_name = f'{prefix}{suffix}.json' split_all_path = f'{run_path}/{split_name}' # Find all per-camera/flex-group manifest JSONs (exclude metadata_shared, tmp_, split_all itself) json_files = sorted(glob(f'{glob_escape(run_path)}/{prefix}_*.json')) # Get intermediate json to merge and delete json_files_all = [f for f in json_files if os.path.basename(f).startswith(prefix) and \ os.path.basename(f).endswith(f'{suffix}.json') and \ os.path.basename(f) != split_name ] # Remove subset files (they have been merged already) json_files = [jf for jf in json_files_all if len(jf.split('__')) < 4] if len(json_files) == 0: # Revert to all json files if they haven't been merged json_files = json_files_all # Find all per-camera/flex-group manifest JSONs (exclude metadata_shared, tmp_, split_all itself) error_files = sorted(glob(f'{run_path}/errors_all_*.json')) # Get intermediate json to merge and delete error_files = [f for f in error_files if os.path.basename(f).endswith(f'{suffix}.json') ] if not json_files: console.print(f'[yellow]No split manifests found for prefix {prefix}, skipping {split_name} generation') return # Merge all manifests all_sequences = {} all_tags = [] all_labels = [] all_cameras = [] all_metadata = {} splits = [] errors = {} shared_metadata_path = f'{args.path}/metadata_shared.json' shared_metadata = read_json(shared_metadata_path) invalids = [] # Loop over tarfile sequences for jf in tqdm(json_files, ncols=96, leave=False, desc='JSON Loop'): data = read_json(jf) split = os.path.basename(jf)[:-len('.json')].split('__')[1] splits.append(split) if data is None: continue seqs = data.get('sequences', {}) for seq_key, tars in tqdm(seqs.items(), ncols=96, leave=False, desc='JSON Sequence Loop'): if seq_key in invalids: continue # Seq has been flagged as invalid tars = sorted(tars, key=lambda x: int(x.split('_')[0])) metadata_path = f'{args.path}/{seq_key}/metadata.json' # Get metadata path if not os.path.exists(metadata_path): errors[seq_key] = f'MISSING METADATA [*FILTER*] : {metadata_path}' invalids.append(seq_key) continue metadata = read_json(metadata_path) # Get metadata # Check if length is the same across cameras num_frames = get_num_frames(metadata['num_frames'], metadata['cameras']) # NOTE(yams_any4d 2026-07-09): writer names the last tar by raw end frame, # which under stride>1 can exceed num_frames//stride*stride by up to # stride-1. Tolerate that instead of invalidating every sequence. _expected_end = num_frames // args.stride * args.stride _last_end = int(tars[-1][:-len('.tar')].split('_')[-1]) valid_length = int(tars[0].split('_')[0]) == 0 and \ abs(_last_end - _expected_end) < 2 * max(1, args.stride) if not valid_length: msg = f'{tars[0]} & {tars[-1]} != {num_frames}' errors[seq_key] = f'INCONSISTENT LENGTH [*FILTER*] : {msg}' invalids.append(seq_key) continue # Skip if frame length is wrong if seq_key not in all_sequences: all_sequences[seq_key] = {} # Merge tar files (deduplicate) if split not in all_sequences[seq_key]: all_sequences[seq_key][split] = tars else: valid_tars = all_sequences[seq_key][split] == tars if not valid_tars: msg = f'{len(all_sequences[seq_key][split])} != {len(tars)}' errors[seq_key] = f'INCONSISTENT NUMBER OF TARFILES [*FILTER*] : {msg}' invalids.append(seq_key) continue # Skip if tarfile length is wrong all_metadata[seq_key] = metadata all_tags.extend(metadata['info']['tags']) all_tags = list(set(all_tags)) all_labels.extend(metadata['labels']) all_labels = list(set(all_labels)) all_cameras.extend(metadata['cameras']) all_cameras = list(set(all_cameras)) # Deduplicate and sort valid sequences keys = sorted(list(set(all_sequences.keys()))) keys = [key for key in keys if key not in invalids] all_sequences = {key: all_sequences[key] for key in keys} # Check and consolidate tarfiles from multiple cameras valid_sequences = {} for seq_key, val in all_sequences.items(): if isinstance(val, dict): val_keys = list(val.keys()) valid_tars = all([set(val[val_keys[0]]) == set(val[val_key]) for val_key in val_keys]) if args.cameras is None: valid_cameras = set(val_keys) == set(all_metadata[seq_key]['cameras']) else: valid_cameras = set(val_keys).issubset(args.cameras) if valid_tars and valid_cameras: valid_sequences[seq_key] = val[val_keys[0]] elif not valid_tars: errors[seq_key] = 'INCONSISTENT NUMBER OF TARFILES [*FILTER*]' elif not valid_cameras: msg = f'{set(val_keys)} != {all_metadata[seq_key]["cameras"]} | {args.cameras}' errors[seq_key] = f'INCONSISTENT NUMBER OF CAMERAS [*FILTER*] : {msg}' valid_sequences[seq_key] = val[val_keys[0]] # Flag but keep since other cameras are still useful else: valid_sequences[seq_key] = val all_sequences = valid_sequences all_metadata = {key: all_metadata[key] for key in all_sequences.keys()} num_files = sum(len(v) for v in all_sequences.values()) shared_tags, shared_labels, shared_cameras = get_shared_info(all_metadata) info = dict( webbed=args.store, stride=args.stride, num_sequences=len(all_sequences), num_files=num_files, tags=sorted(all_tags), labels=sorted(all_labels), cameras=sorted(all_cameras), shared_tags=sorted(shared_tags), shared_labels=sorted(shared_labels), shared_cameras=sorted(shared_cameras), source_manifests=[os.path.basename(f) for f in json_files], ) write_json(split_all_path, dict( command=f'auto-generated {split_name} combining all per-camera/flex-group manifests', config=raw_config, info=info, sequences=all_sequences, )) console.print(f'[bold green]Combined {split_name} written:[/bold green] {split_all_path} ' f'({len(all_sequences)} sequences, {num_files} files)') for error_path in error_files: errors.update(read_json(error_path)['errors']) errors_keys = sorted(list(set(errors.keys()))) errors = {key: errors[key] for key in errors_keys} # Write webdataset statistics stats_name = write_stats(args, shared_metadata_path, all_metadata, shared_metadata, errors) ### Log errors in the webbed folder errors_path = f'{run_path}/errors_all.json' if args.subset is not None: errors_path = errors_path.replace('.json', f'__{args.subset.replace("/","_")}.json') errors_dict = {'total': len(errors), 'errors': {}} for key, val in errors.items(): errors_dict['errors'][key] = val write_json(errors_path, errors_dict) # Dispose of intermediate splits + errors (per-camera/flex manifests) per --intermediate if args.intermediate == 'move': intermediate_dir = f'{run_path}/intermediate' for json_file in json_files_all: mv_sync(json_file, intermediate_dir, args) for error_file in error_files: mv_sync(error_file, intermediate_dir, args) elif args.intermediate == 'delete': for json_file in json_files_all: rm_sync(json_file, args) for error_file in error_files: rm_sync(error_file, args) # 'keep': leave intermediates in place # Upload if requested if args.upload: run_path_s3 = run_path.replace(args.local_path, args.s3_path) aws_s3_cp(split_all_path, f'{run_path_s3}/{split_name}', robust=True) stats_name_s3 = stats_name.replace(args.local_path, args.s3_path) aws_s3_cp(stats_name, stats_name_s3, robust=True) errors_name_s3 = errors_path.replace(args.local_path, args.s3_path) aws_s3_cp(errors_path, errors_name_s3, robust=True) console.print(f'[bold green]Combined {split_name} uploaded to S3') def parse_args(): parser = argparse.ArgumentParser() parser.add_argument('dataset', type=str, nargs='+') # e.g. "DROID/DROID12 41b8" parser.add_argument('--subset', type=str, default=None) parser.add_argument("--num_procs", type=int, default=12) parser.add_argument("--max_threads", type=int, default=32) parser.add_argument('--official', action='store_true') parser.add_argument('--upload', action='store_true') parser.add_argument('--delete', action='store_true') parser.add_argument('--intermediate', type=str, default='move', choices=['move', 'delete', 'keep'], help='What to do with per-camera/flex split + error manifests after consolidating into ' 'split_all.json. "move" (default): into an intermediate/ subfolder (keeps their real ' 'commands). "delete": remove local + S3. "keep": leave in place.') parser.add_argument('--restart', action='store_true') parser.add_argument('--overwrite', action='store_true') parser.add_argument('--cleanup', action='store_true') parser.add_argument('--per_camera', action='store_true') parser.add_argument('--global_cam_idx', type=int, default=None) parser.add_argument('--split_cam', type=str, default='split_all') parser.add_argument('--quiet', action='store_true') parser.add_argument('--local_path', type=str, default=None) parser.add_argument('--s3_path', type=str, default='s3://tri-ml-sandbox-16011-us-west-2-datasets') parser.add_argument('--folder', type=str, default='cv_webbed') parser.add_argument('--folder_unified', type=str, default='cv_unified') parser.add_argument("--subfolder", type=str, default=None) parser.add_argument("--unified_base", type=str, default=None) parser.add_argument("--metadata_shared", type=str, default=None, help='Path to metadata_shared JSON (default: auto-detect from unified dir)') parser.add_argument("--name_suffix", type=str, default=None, help='Suffix appended to the webbed output folder name') parser.add_argument('--stride', type=int, default=1) # Apply frame stride when creating the context parser.add_argument('--resize', type=float, nargs='+', default=None) # Resize sample to desired value parser.add_argument('--crop_first', type=int, default=None) parser.add_argument('--pad_last', type=int, default=None) parser.add_argument('--labels', type=str, nargs='+', default=['rgb','intrinsics','extrinsics','language','action']) parser.add_argument('--web_only', action='store_true') # Filter cameras/labels from split filters parser.add_argument('--depth', action='store_true') # Include depth as a label parser.add_argument('--depth_bits', type=int, default=None) # Include depth as label and define how many bits to encode parser.add_argument('--decord', action='store_true') # Use decord as mp4 decoder parser.add_argument('--cameras', type=str, nargs='+', default=None, help='Optional list of camera names to include. ' 'Overrides/restricts the default or detected camera list; ' 'when per-camera export is enabled, only these cameras are exported.') parser.add_argument('--camera_mode', type=str, default='flexible', choices=['strict', 'flexible'], help='How to handle episodes with missing cameras. ' '"strict": skip episodes missing any requested camera. ' '"flexible": include each episode with whatever requested cameras it has.') parser.add_argument('--snippet_formula', type=float, default=1.0, help='Exponent controlling how snippet count scales with sequence length. ' '1.0 (default/linear): parts ~ length/buffer. ' '0.5 (sqrt): parts ~ sqrt(length/buffer). ' 'Values like 0.8 give intermediate density between sqrt and linear.') 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') mode.add_argument('--latents', action='store_const', dest='mode', const='latents') from_mode = parser.add_mutually_exclusive_group(required=True) from_mode.add_argument('--from_frames', action='store_const', dest='from_mode', const='frames') from_mode.add_argument('--from_videos', action='store_const', dest='from_mode', const='videos') args = parser.parse_args() # Cleanup requires no subset if args.cleanup: args.subset = None args.restart = True # Remove subset if it's a single part if args.subset == '1/0': args.subset = None args.folder_unified += f'/{args.from_mode}' # Include depth as label if requested if args.depth or args.depth_bits is not None: args.labels.append('depth') # Ensure lowdim is loaded when derived labels (intrinsics, extrinsics, action, # language) are requested -- these live inside lowdim NPZ for some datasets. _LOWDIM_DERIVED = {'intrinsics', 'extrinsics', 'action', 'language'} if _LOWDIM_DERIVED & set(args.labels) and 'lowdim' not in args.labels: args.labels.append('lowdim') # Add debug to folder if not official (be careful when using official!) if not args.official: args.folder += '_debug' console.print(f'[orange3]DEBUG mode enabled') # Unpack positional args assert len(args.dataset) == 2, \ "Expected dataset spec and store as two positional arguments, e.g. DROID/DROID12 41b8" args.dataset_spec, args.store = args.dataset # "DROID/DROID12", "41b8" if 'b' in args.store: args.store_context, args.store_buffer = [int(v) for v in args.store.split('b')] else: args.store_context, args.store_buffer = int(args.store), 0 # Auto-derive local_path from absolute dataset_spec when not provided if args.local_path is None: if os.path.isabs(args.dataset_spec): # e.g. /datasets/basile/cv_unified_debug/RH20T_d9/split_all.json # -> base = /datasets/basile (strip cv_unified* and everything after) unified_dir = os.path.dirname(args.dataset_spec) parent = os.path.dirname(unified_dir) # parent is e.g. /datasets/basile/cv_unified_debug -> strip cv_unified* if os.path.basename(parent).startswith(args.folder_unified): args.local_path = os.path.dirname(parent) else: args.local_path = parent console.print(f'[cyan]Auto-derived --local_path from dataset spec: {args.local_path}') else: console.print('[red]ERROR: --local_path is required when using a relative dataset spec[/red]') sys.exit(1) # Save original base path before mangling (used for resolving unified paths) args.local_path_base = args.local_path # Prepare local and s3 paths args.local_path = f'{args.local_path}/{args.folder}/{args.mode}' args.s3_path = f'{args.s3_path}/{args.folder}/{args.mode}' console.print(f'[orange3]Local output path: {args.local_path}') console.print(f'[orange3]S3 output path: {args.s3_path}') return args def get_name(dataset, recipe, store, stride, config, snippet_formula='linear', camera_mode='flexible', name_suffix=None): # Record labels (using the famous IKEDOCLABSN convention; see README naming table) labels_str = '' if 'rgb' in dataset.labels: labels_str += 'I' if 'intrinsics' in dataset.labels: labels_str += 'K' if 'extrinsics' in dataset.labels: labels_str += 'E' if 'depth' in dataset.labels: labels_str += 'D' if 'optflow' in dataset.labels: labels_str += 'O' if 'scnflow' in dataset.labels: labels_str += 'C' if 'language' in dataset.labels: labels_str += 'L' if 'action' in dataset.labels: labels_str += 'A' if 'bbox2d' in dataset.labels or 'bbox3d' in dataset.labels: labels_str += 'B' if 'semantic' in dataset.labels: labels_str += 'S' if 'normals' in dataset.labels: labels_str += 'N' # Initialize name if recipe.endswith('.json'): # Remove .json extension recipe = recipe[:-len('.json')] name = f'{recipe}__{labels_str}_{store}' # Context stride if stride > 1: name += f'_str{stride}' # Snippet formula if snippet_formula != 1.0: name += f'_snip{snippet_formula:g}' # Camera mode if camera_mode == 'strict': name += '_strict' ### Record relevant configuration parameters # Augmentations augmentation = config.get('augmentation', None) if augmentation is not None: if 'resize' in augmentation: # Round float values to int when possible for cleaner names for i in range(len(augmentation["resize"])): if round(augmentation["resize"][i]) == augmentation["resize"][i]: augmentation["resize"][i] = int(augmentation["resize"][i]) # Add resize parameters to name (removing problematic characters) name += f'_res{augmentation["resize"]}'.replace('[','').replace(']','').replace(',','x') # Append user suffix if name_suffix: name += f'_{name_suffix}' # Return final name return name.replace(' ', '').replace("'",'') def write_tarfile(args, filename, sample): # Write tarfile tarfile = TarWriter(filename, compress=False, encoder=True, keep_meta=True) tarfile.write(sample) tarfile.close() if args.upload: # Upload tarfile to S3 immediately (incremental) filename_s3 = filename.replace(args.local_path, args.s3_path) if not aws_s3_cp(filename, filename_s3, robust=True): return # skip delete if upload failed if args.delete: remove_path(filename) with open(filename.replace('.tar', '.txt'), 'w') as fp: pass def _frame_sample_for_cam(args: argparse.Namespace, data: Dict[str, dict], cam, time_st, time_fn, store_raw): """Build a sample dict with raw frame data for one camera in a time window.""" sample = {} # NOTE(bvh): fix bug that was recently introduced; # cam is the dataloader's local index (for data access), # typically =0 only for --per_camera mode. # global_cam_idx is only for output key strings (tar sample naming). out_cam = args.global_cam_idx if args.global_cam_idx is not None else cam for key, key_data in data.items(): for time in range(time_st, time_fn, args.stride): tc_key = f'({time},{out_cam})' if store_raw: if key.startswith('rgb'): sample[f'{key}.{tc_key}.jpg'] = rgb_to_byte_array(key_data[cam][time]) elif key.startswith('depth'): encoded_depth, encode_info = encode_dense(key_data[cam][time], bits=args.depth_bits) sample[f'{key}.{tc_key}.npz'] = dict(depth=encoded_depth, encode_info=encode_info) elif key.startswith('mask'): sample[f'{key}.{tc_key}.png'] = mask_to_byte_array(key_data[cam][time]) if key.startswith('lowdim'): sample[f'{key}.{tc_key}.npz'] = read_npz_node(key_data[cam], index=time) return sample def _video_sample_for_cam(args: argparse.Namespace, data: Dict[str, dict], cam, time_st, time_fn, store_raw): """Build a sample dict with raw sequence data for one camera in a time window.""" sample = {} # NOTE(bvh): cam is the dataloader's local index (for data access), # typically =0 only for --per_camera mode. # global_cam_idx is only for output key strings (tar sample naming). out_cam = args.global_cam_idx if args.global_cam_idx is not None else cam time_cam_str = f'({time_st}:{time_fn},{out_cam})' indexer = slice(time_st, time_fn, args.stride) for key, key_data in data.items(): if store_raw: if key.startswith('rgb'): rgb_ext = 'mp4' rgb_frames = key_data[cam][indexer] sample[f'{key}.{time_cam_str}.{rgb_ext}'] = write_video_seq(None, rgb_frames, ext=rgb_ext) if key.startswith('depth'): encoded_depth, encode_info = encode_dense(key_data[cam][indexer], bits=args.depth_bits) sample[f'{key}.{time_cam_str}.npz'] = dict(depth=encoded_depth, encode_info=encode_info) if key.startswith('mask'): sample[f'{key}.{time_cam_str}.npy'] = key_data[cam][indexer] if key.startswith('lowdim'): sample[f'{key}.{time_cam_str}.npz'] = read_npz_node(key_data[cam], index=indexer) return sample def stack_features(feature_list: List[dict | np.ndarray | Image.Image], label: str = None): """Stack a list of features (dicts or arrays) along a new time dimension.""" return stack_sample(feature_list) if isinstance(feature_list[0], dict) else np.stack(feature_list) def _consolidate_data_dict(data: Dict[str, Dict[str, Any]]): """ Consolidate a data dict with keys like {label: {(time, cam): value}} into {label: {cam: time_array}}. """ consolidated = {k: {} for k in data.keys()} for name, modality in data.items(): if name == 'metadata': consolidated[name] = modality continue keys = list(modality.keys()) if not isinstance(keys[0], tuple): # action and language are just cam 0 # If keys are not (time, cam) tuples, we assume they're already consolidated consolidated[name][0] = stack_features([modality[k] for k in keys]) else: per_cam_features = {cam: [] for cam in range(max(k[1] for k in keys) + 1)} for k in keys: per_cam_features[k[1]].append(modality[k]) for cam, cam_features in per_cam_features.items(): consolidated[name][cam] = stack_features(cam_features) return consolidated def _latent_sample_for_cam(args: argparse.Namespace, cam, time_st, time_fn, lat_rgb, lat_depth, lat_language): """Build a sample dict with VAE-encoded latents for one camera in a time window.""" sample = {} for time in range(time_st, time_fn, args.stride): time_cam = (time, cam) if args.global_cam_idx is None: tc_key = f'({time},{cam})' else: tc_key = f'({time},{args.global_cam_idx})' if time_cam in lat_rgb: sample[f'lat_rgb.{tc_key}.npz'] = dict(data=lat_rgb[time_cam].float().cpu().numpy()) if time_cam in lat_depth: sample[f'lat_depth.{tc_key}.npz'] = dict(data=lat_depth[time_cam].float().cpu().numpy()) if lat_language is not None and cam == 0: sample[f'lat_language.({time_st},0).npz'] = lat_language return sample def prepare(args: argparse.Namespace, seq, i, data, consolidated, time_st, time_fn, tokenizers, device): """Build per-camera samples for one snippet. Returns list of [args, filename, sample].""" metadata = data['metadata'] meta_name = f'{args.name}/{args.split}' meta_path = metadata['path'].split('/') meta_path = f'{meta_name}/{"/".join(meta_path[:-2])}' meta_cameras = metadata['cameras'] # Store stride information in metadata metadata['stride'] = args.stride metadata['length'] = args.length cams = list(set([tc[1] for tc in data['rgb'].keys()])) # Encode latents if needed lat_rgb, lat_depth = {}, {} lat_language = None if args.mode in ['latents']: args1, args2 = (data, time_st, time_fn, cams), (tokenizers, args.buffer, args.stride) if 'rgb' in data: lat_rgb = encode_latents(*args1, 'rgb', *args2, device=device) if 'depth' in data: lat_depth = encode_latents(*args1, 'depth', *args2, device=device) if getattr(tokenizers, "text_encoder", None) is not None and 'language' in data: lang = data['language'] prompt = None if isinstance(lang, dict) and len(lang) > 0: key0 = next(iter(lang.keys())) if isinstance(key0, tuple): prompt = lang.get((0, 0), {}).get('prompt', None) else: prompt = lang.get(0, {}).get('prompt', None) if prompt is not None: if isinstance(prompt, list): prompt_in = prompt[0] if len(prompt) > 0 else None else: prompt_in = prompt if prompt_in is not None: encoded = tokenizers.encode_text(prompt_in) lat_language = {} for key, val in encoded.items(): if isinstance(val, torch.Tensor): lat_language[key] = val[0].float().cpu().numpy() else: lat_language[key] = val lat_language['prompt'] = np.array(prompt_in, dtype=object) samples = [] for cam in cams: sample = {} if args.mode == 'frames': sample.update(_frame_sample_for_cam(args, consolidated, cam, time_st, time_fn, store_raw=True)) elif args.mode == 'videos': sample.update(_video_sample_for_cam(args, consolidated, cam, time_st, time_fn, store_raw=True)) elif args.mode in ['latents']: sample.update(_latent_sample_for_cam(args, cam, time_st, time_fn, lat_rgb, lat_depth, lat_language)) else: raise ValueError(f'Unsupported mode: {args.mode}') sample['metadata.npz'] = dict(metadata=np.array(data['metadata'], dtype=object)) key = f'{meta_path}/{meta_cameras[cam]}' filename = f'{args.local_path}/{key}/%04d_%04d.tar' % (time_st, time_fn) create_folder(filename) if args.subset is not None: if '//' in args.subset: # Apply subset of subsets (for faster parallel processing) subset = args.subset.split('//')[0].split('/') else: subset = args.subset.split('/') seq = int(subset[0]) * seq + int(subset[1]) # Tag sample with length length = len(sample.keys()) - 1 if length <= 0: continue sample_with_length = {'__key__': f'{seq}_{i}'} for k, val in sample.items(): if not k.startswith('__'): k = f'{k[:-4]}.{length}{k[-4:]}' sample_with_length[k] = val sample = sample_with_length samples.append([args, filename, sample]) return samples def webdatasetize_sequence(i, seq, args: argparse.Namespace, tokenizers, device): """Web a specific sequence""" dataset_obj: UnifiedDataset = args.dataset_obj cameras = dataset_obj.get_sample_cameras(seq, None) # Select valid cameras filename, time_cam0, loc = dataset_obj.get_filename_target(seq, cam_idx=0, cam=cameras[0]) context = dataset_obj.get_filename_context(seq, cam_idx=0, cam=cameras[0]) time_cams = [time_cam0] + list(context.keys()) sample_folder = '/'.join(filename.split('/')[:-3]) # sample_folder/label/cam/container.ext -> sample_folder sample_id = sample_folder[len(args.unified_path)+1:] path = f"{sample_id}/{cameras[0]}" length = len(set(tc[0] for tc in time_cams)) start = 0 # Start from the beginning my_length, buffer = args.store_context, args.store_buffer context = my_length - 1 # 41b4 -> 40s4 args.length = my_length # Save for storing in metadata if buffer == 0: # Sequence too short if length <= 1: if not args.quiet: console.print(f'[yellow][SKIP][/yellow] seq {seq} ({path}) too short ({length} frames < {my_length} required)') return 'too_short' # Create starts and finishs list by linearly sampling values = np.arange(0, length + context, context + 1) values = [int(np.round(v / values[-1] * length)) for v in values] if args.stride > 1: # Account for stride and re-normalize for length values = [v * args.stride for v in values] values = [v for v in values if v <= length] values = [int(np.round(v / values[-1] * length)) for v in values] # Enforce the same stride across tarfiles values = [int(v // args.stride * args.stride) for v in values] values = [v for v in values if v <= length] # Separate starts and finishs starts, finishs = values[:-1], values[1:] else: # End point given context and stride finish = length - (context + buffer) * args.stride # Try smaller buffer sizes if sequence is too short while finish < start and buffer > 1: buffer -= 1 finish = length - (context + buffer) * args.stride # Sequence too short for this buffer if finish < start: if not args.quiet: console.print(f'[yellow][SKIP][/yellow] seq {seq} ({path}) too short ({length} frames < {(context + buffer) * args.stride} required)') return 'too_short' if args.snippet_formula != 1.0: # Exponent scaling: parts ~ (finish / (2*buffer))^exponent # exponent=0.5 is sqrt, exponent=0.8 is intermediate, exponent=1.0 is linear. parts = max(1, int(np.round((finish / (2 * buffer)) ** args.snippet_formula))) else: parts = int(np.round(finish / buffer)) + 1 # Linear: parts ~ length/buffer starts = np.linspace(start, finish, parts) # List of starts starts = [int(start) for start in starts] finishs = [start + (context + buffer) * args.stride for start in starts] # List of finishs # Use the whole sequence if there is only one part if len(finishs) == 1: finishs[0] = length # Check if sequence already exists in s3, and skip if it's there already if not args.overwrite: name = args.name # NOTE(bvh): metadata['name'] can differ from the S3 directory name (e.g. DROID vs DROID101) local_path = f'{args.local_path}/{name}/{args.split}/{sample_id}' already_processed = os.path.exists(f'{local_path}/metadata.json') if args.upload: s3_path = f'{args.s3_path}/{name}/{args.split}/{sample_id}/' s3_files = args.paginator.ls_seq(s3_path) if already_processed and args.upload: already_processed = '' in s3_files and len(s3_files['']) == 1 and s3_files[''][0] == 'metadata.json' if already_processed: cam_valids = {cam: False for cam in cameras} for cam in cameras: if not already_processed: continue local_path = f'{args.local_path}/{name}/{args.split}/{sample_id}/{cam}' if args.upload: files = s3_files[cam] if cam in s3_files else [] else: files = sorted(os.listdir(local_path)) if os.path.exists(local_path) else [] files = sorted(files, key=lambda x: int(x.split('_')[0])) if len(files) == len(starts): cam_valids[cam] = True for basename, st, fn in zip(files, starts, finishs): st_fn = basename[:-len('.tar')].split('_') if int(st_fn[0]) != st or int(st_fn[1]) != fn: cam_valids[cam] = False break if cam_valids[cam]: if not args.quiet: console.print(f'[yellow][SKIP][/yellow] seq {seq} ({path}) already processed') tarfiles = sorted(list(glob(f'{local_path}/*.tar'))) txtfiles = sorted(list(glob(f'{local_path}/*.txt'))) tarfiles = [os.path.basename(t) for t in tarfiles] txtfiles = [os.path.basename(t).replace('.txt', '.tar') for t in txtfiles] available_files = tarfiles + txtfiles missing_files = [f for f in files if f not in available_files] for f in missing_files: os.makedirs(local_path, exist_ok=True) filename = f'{local_path}/{f}' with open(filename.replace('.tar', '.txt'), 'w') as fp: pass already_processed = all(cam_valids.values()) if already_processed: return 'already_processed' # Get metadata information data = args.dataset_obj[seq] metadata = data['metadata'] # Prepare each snippet write_threads = [] consolidated = _consolidate_data_dict(data) for i in range(len(starts)): time_st, time_fn = starts[i], finishs[i] data_process = prepare(args, seq, i, data, consolidated, time_st, time_fn, tokenizers, device) for dp in data_process: if args.upload: thread = Thread(target=write_tarfile, args=dp) thread.start() write_threads.append(thread) else: write_tarfile(*dp) # Keep number of threads in check to avoid blowing up memory while len(write_threads) > args.max_threads: write_threads[0].join() write_threads = write_threads[1:] # Finish all threads before moving on for thread in write_threads: thread.join() # Read+enrich per-episode metadata and write to webbed output meta_path = metadata['path'].split('/')[:-2] unified_root = args.unified_path if isinstance(unified_root, str) and unified_root.endswith(".json"): unified_root = os.path.dirname(unified_root) meta_src = os.path.join(unified_root, *meta_path) meta_dst = os.path.join(args.local_path, args.name, args.split, *meta_path) meta = read_json(f'{meta_src}/metadata.json') tarfiles = sorted(list(glob(f'{meta_dst}/{metadata["cameras"][0]}/*.tar'))) if len(tarfiles) == 0: tarfiles = sorted(list(glob(f'{meta_dst}/{metadata["cameras"][0]}/*.txt'))) tarfiles = [f.replace('.txt', '.tar') for f in tarfiles] meta['tarfiles'] = [os.path.basename(t) for t in tarfiles] if hasattr(args, 'metadata_shared_data') and args.metadata_shared_data is not None: shared_cams = args.metadata_shared_data.get('cameras', args.metadata_shared_data.get('cameras_available', [])) shared_labels = args.metadata_shared_data.get('labels', args.metadata_shared_data.get('labels_available', [])) meta['cameras'] = [c for c in meta['cameras'] if c in shared_cams] meta['labels'] = [c for c in meta['labels'] if c in shared_labels] if args.stride > 1: meta['stride'] = args.stride ### METADATA FIX (not applied yet, keeping for future incorporation) # if args.resize is not None and 'resolution' in meta: # meta['resolution'] = parse_resize_params(meta['resolution'][::-1], args.resize) # if args.stride > 1 and 'framerate' in meta: # meta['framerate'] = meta['framerate'] / args.stride # if 'labels' in meta: # for label in meta['labels']: # if label not in args.labels: # meta.pop(label) # meta['labels'] = [l for l in meta['labels'] if l in args.labels] write_json(f'{meta_dst}/metadata.json', meta) # Upload tarfiles if requested if args.upload: meta_dst = f'{meta_dst}/metadata.json' meta_dst_s3 = meta_dst.replace(args.local_path, args.s3_path) aws_s3_cp(meta_dst, meta_dst_s3, robust=True) # meta_dst_s3 = meta_dst.replace(args.local_path, args.s3_path) # aws_s3_sync(meta_dst, meta_dst_s3, extras='--exclude "*" --include "*.tar" --include "*.json"', robust=True) # #### ADD CHECK TO SEE IF EVERYTHING WAS UPLOADED PROPERLY # s3_meta_dst = os.path.dirname(meta_dst.replace('/data/', '')) + '/' # s3_files = args.paginator.ls_seq(s3_meta_dst) # s3_cameras = [key for key in s3_files if key != ''] # local_cameras = glob(f'{os.path.dirname(meta_dst)}/*') # local_cameras = [os.path.basename(cam) for cam in local_cameras if not cam.endswith('.json')] # local_tarfiles = meta['tarfiles'] # # Consider only current cameras # s3_cameras = [cam for cam in s3_cameras if cam in cameras] # local_cameras = [cam for cam in local_cameras if cam in cameras] # check_metadata = '' in s3_files and len(s3_files['']) == 1 and s3_files[''][0] == 'metadata.json' # check_cameras = set(local_cameras) == set(s3_cameras) # check_tarfiles = all([set(s3_files[cam]) == set(local_tarfiles) for cam in s3_cameras]) # if not (check_metadata and check_cameras and check_tarfiles): # if repeat < 5: # webdatasetize_sequence(i, seq, args, tokenizers, device, repeat + 1) # else: # raise ValueError(f'Incomplete upload: {check_metadata}/{check_cameras}/{check_tarfiles}') del data def webdatasetize_sequences(i, seqs, args: argparse.Namespace): """Web each sequence in a subset""" if args.mode in ["frames", "sequence", "videos"]: tokenizers = None device = "cpu" else: console.print(f'[orange3]mode={args.mode}, initializing VAE tokenizers...') if torch.cuda.is_available(): local_rank = int(os.environ.get("LOCAL_RANK", "0")) local_cuda_idx = local_rank % torch.cuda.device_count() torch.cuda.set_device(local_cuda_idx) device = f"cuda:{local_cuda_idx}" else: device = "cpu" from anydata.webdataset.vae import VAE tokenizers = VAE( device, vae_path="s3://tri-ml-sandbox-16011-us-west-2-datasets/cosmos-predict-2/checkpoints/nvidia/Cosmos-Predict2-2B-Video2World/tokenizer/tokenizer.pth", text_encoder_path="s3://tri-ml-sandbox-16011-us-west-2-datasets/cosmos-predict-2/checkpoints/google-t5/t5-11b", load_text_encoder=getattr(args, "need_language_latents", False), ) args.paginator = Paginator() # Plain dict fallback for direct / single-process calls. if not hasattr(args, '_skipped'): args._skipped = {} suc, err, don = 0, 0, 0 progress = tqdm(seqs, ncols=96, leave=False) for seq in progress: # Get filename path for logging purposes camera = args.cam filename = args.dataset_obj.get_filename_target(seq, cam_idx=0, cam=camera)[0] # filename = clean_string(filename) filename = filename.replace(f'{args.unified_path}/', '') filename = '/'.join(filename.split('/')[:-3]) filename_tmp = filename.replace('/', '__') + f'__{camera}' filename_tmp = f'{args.local_path}/{args.name}/{args.split}/tmp/{filename_tmp}.txt' try: description = f'| SED: {suc}/{err}/{don}' if args.subset is not None: description = f': {args.subset} {description}' progress.set_description(f'# {args.name}/{camera} {description}') reason = webdatasetize_sequence(i, seq, args, tokenizers, device) if reason is None: # Run again to make sure it's done for _ in range(3): reason_check = webdatasetize_sequence(i, seq, args, tokenizers, device) if reason_check == 'already_processed': break # Successful, no reason to try again if reason_check != 'already_processed': raise ValueError(f'Incomplete upload') # Raise error if not processed # if reason is not None: # args._skipped[filename] = reason rm_sync(filename_tmp, args) if reason == 'already_processed': don += 1 elif reason is None: suc += 1 except Exception as e: args._skipped[filename] = f'error: {type(e).__name__}: {e}' # Create tmp file to indicate the failure create_folder(filename_tmp) with open(filename_tmp, 'w') as fp: pass if args.upload: # Upload tmp file to s3 filename_tmp_s3 = filename_tmp.replace(args.local_path, args.s3_path) aws_s3_cp(filename_tmp, filename_tmp_s3, robust=True) # Display error if not running quietly if not args.quiet: console.print(f'[red][ERROR][/red] seq {seq} failed, skipping: {type(e).__name__}: {e}') console.print(traceback.format_exc()) err += 1 def _resolve_metadata_shared(unified_dir, override_path=None): """Resolve the metadata_shared JSON path. Uses --metadata_shared override if provided, otherwise default.""" if override_path is not None: if not os.path.isabs(override_path): override_path = os.path.join(unified_dir, override_path) return override_path return os.path.join(unified_dir, 'metadata_shared.json') def webdatasetize(args: argparse.Namespace, dataset_config=None): args.buffer = max(args.store_context - 1, 0) # Include additional parameters to dataset config added = {'webbing': True} if args.subset is not None: added['subset'] = args.subset if args.resize is not None: if 'augmentation' not in added: added['augmentation'] = dict() added['augmentation']['resize'] = args.resize if args.decord: added['rgb_decoder'] = 'decord' # Resolve dataset spec: support absolute paths, relative JSON paths, or YAML recipe names spec = args.dataset_spec if os.path.isabs(spec): # Absolute path to split JSON: /path/to/unified/RH20T_d9/split_all.json split_json_path = spec unified_path = os.path.dirname(split_json_path) dataset_name = os.path.basename(unified_path) recipe = os.path.basename(split_json_path) elif '/' not in spec: spec = f'{spec}/split_all.json' dataset_name, recipe = spec.split('/', maxsplit=1) else: dataset_name, recipe = spec.split('/', maxsplit=1) # Track metadata_shared and filters for later enrichment metadata_shared = None filters = {} if not recipe.endswith('.json'): # YAML branch console.print('[red][DEPRECATED] YAML config path is deprecated. ' 'Use a JSON split path instead (absolute or relative)') config_yaml = f'anydata/webdataset/configs/{dataset_name}.yaml' raw_config = read_config(config_yaml)[recipe] else: # JSON branch # Resolve unified_path from the split JSON path in raw_config, # or from absolute spec, or from --unified_base / --local_path fallback. if not os.path.isabs(spec): # Relative spec like "DROID/DROID17/split_all.json" -- resolve unified_path unified_base = getattr(args, 'unified_base', None) if unified_base is None: # Fallback: try local_path base (before webbed mangling) with cv_unified unified_base = f'{args.local_path_base}/{args.folder_unified}' if not args.official: unified_base += '_debug' unified_path = f'{unified_base}/{dataset_name}' # else: unified_path already set from absolute spec above split_json_path = f'{unified_path}/{recipe}' if not os.path.exists(split_json_path): console.print(f'[red]Split JSON not found: {split_json_path}[/red]') return else: console.print(f'[green]Split JSON found: {split_json_path}[/green]') metadata_shared_path = _resolve_metadata_shared(unified_path, args.metadata_shared) metadata_shared = read_json(metadata_shared_path) cameras_available = metadata_shared.get('cameras', metadata_shared.get('cameras_available', [])) labels_available = metadata_shared.get('labels', metadata_shared.get('labels_available', [])) # Remove labels that were not requested labels_available = [label for label in labels_available if label in args.labels] # Get split filters filters = read_json(split_json_path).get('filters', {}) # Set cameras based on filters if 'with_cameras' in filters and args.web_only: cameras_available = filters['with_cameras'] if 'with_only_cameras' in filters: cameras_available = filters['with_only_cameras'] if 'without_cameras' in filters: cameras_available = [f for f in cameras_available if f not in filters['without_cameras']] # Set labels based on filters if 'with_labels' in filters and args.web_only: labels_available = filters['with_labels'] if 'with_only_labels' in filters: labels_available = filters['with_only_labels'] if 'without_labels' in filters: labels_available = [f for f in labels_available if f not in filters['without_labels']] # Create config dict raw_config = dict( name='Unified', path=[split_json_path], labels=labels_available, cameras=cameras_available, length='ALL', ) # Camera mode handling (strict/flexible) for non-per_camera mode if not args.per_camera and not getattr(args, '_camera_mode_applied', False): split_path = raw_config['path'][0] unified_dir = os.path.dirname(split_path) cams_cfg = raw_config.get('cameras', []) # Flatten nested list (JSON path produces [[cam1, cam2, ...]]) if cams_cfg and isinstance(cams_cfg[0], list): cams_cfg = cams_cfg[0] requested_cams = set(cams_cfg) if requested_cams: ep_cams = _episode_cameras(unified_dir, split_path) if args.camera_mode == 'strict': # Keep only episodes that have ALL requested cameras keep = {k for k, cams in ep_cams.items() if requested_cams.issubset(cams)} skipped = len(ep_cams) - len(keep) if skipped > 0: console.print(f'[yellow]camera_mode=strict: keeping {len(keep)}/{len(ep_cams)} episodes ' f'({skipped} skipped for missing cameras)[/yellow]') if len(keep) == 0: console.print('[red]No episodes have all requested cameras, aborting[/red]') return # Write filtered split filtered_split = split_path.replace('.json', f'_strict{_rank_token()}.json') _filter_split(split_path, keep, filtered_split) raw_config['path'] = [filtered_split] elif args.camera_mode == 'flexible': # Group episodes by which requested cameras they have from collections import defaultdict groups = defaultdict(set) for ep, cams in ep_cams.items(): available = tuple(sorted(requested_cams & cams)) if available: groups[available].add(ep) if len(groups) == 0: console.print('[red]camera_mode=flexible: no episodes have any requested cameras, aborting[/red]') return elif len(groups) == 1: # All episodes have the same cameras -- just use strict-like path cam_set = next(iter(groups)) if set(cam_set) == requested_cams: console.print(f'[green]camera_mode=flexible: all {len(ep_cams)} episodes have all cameras[/green]') else: console.print(f'[yellow]camera_mode=flexible: all episodes share cameras {list(cam_set)}[/yellow]') raw_config['cameras'] = list(cam_set) else: # Multiple camera groups -- run each group as a separate webdatasetize call console.print(f'[cyan]camera_mode=flexible: {len(groups)} camera groups found[/cyan]') for gi, (cam_set, eps) in enumerate(sorted(groups.items(), key=lambda x: -len(x[1]))): cam_label = '_'.join(cam_set) console.print(f' Group {gi}/{len(groups)} [{cam_label}]: {len(eps)} episodes, {len(cam_set)} cameras') # Write filtered split for this group group_split = split_path.replace('.json', f'_flex_{cam_label}{_rank_token()}.json') _filter_split(split_path, eps, group_split) # Build config for this group group_config = deepcopy(raw_config) group_config['path'] = [group_split] group_config['cameras'] = list(cam_set) group_args = deepcopy(args) group_args._camera_mode_applied = True group_args.split_cam = f'split_all_flex_{cam_label}' group_args.restart = True # Run sequentially (multi_thread forks processes internally, # which deadlocks inside ThreadPoolExecutor threads) dataset, config = dataset_from_config( config=group_config, name=recipe, added=added, subfolder=group_args.subfolder, ) webdatasetize(group_args, dataset_config=(dataset, config)) # Generate combined split_all.json from all flex-group manifests webbed_name = get_name(dataset, recipe, args.store, args.stride, config, args.snippet_formula, args.camera_mode, getattr(args, 'name_suffix', None)) run_path = f'{args.local_path}/{dataset_name}/{webbed_name}' _generate_split_all(run_path, raw_config, args) return # Per-camera mode: process each camera independently into the same output folder if args.per_camera: path = raw_config['path'][0] split_name = os.path.basename(path).split('.')[0] unified_dir = os.path.dirname(path) base_meta_path = _resolve_metadata_shared(unified_dir, args.metadata_shared) base_meta = read_json(base_meta_path) # Use config cameras (intersected with available) to respect the config cfg_cams = raw_config.get('cameras', []) if cfg_cams and isinstance(cfg_cams[0], list): cfg_cams = cfg_cams[0] if args.cameras is None: available_cams = base_meta.get('cameras', base_meta.get('cameras_available', [])) else: available_cams = args.cameras if cfg_cams: cameras = [c for c in cfg_cams if c in available_cams] else: cameras = available_cams all_cameras = cameras # Load base split to get episode list base_split = read_json(path) base_sequences = base_split.get('sequences', {}) # Build per-episode camera availability from unified metadata episode_cameras = {} for ep_name in tqdm(base_sequences, ncols=96, leave=False, desc='Camera Availability'): ep_meta_path = f'{unified_dir}/{ep_name}/metadata.json' if os.path.exists(ep_meta_path): ep_meta = read_json(ep_meta_path) episode_cameras[ep_name] = ep_meta.get('cameras', []) # Build per-camera (dataset, config, args) tuples, then run in parallel cam_tasks = [] for i, cam in enumerate(cameras): # Filter episodes to only those that have this camera cam_sequences = { ep: frames for ep, frames in base_sequences.items() if cam in episode_cameras.get(ep, []) } if not cam_sequences: console.print(f'[yellow]Skipping camera {cam}: no episodes have it') continue # Write a filtered split file for this camera cam_split_path = path.replace(f'{split_name}.json', f'{split_name}__{cam}{_rank_token()}.json') cam_split = deepcopy(base_split) cam_split['sequences'] = cam_sequences cam_split['info'] = { **cam_split.get('info', {}), 'camera': cam, 'num_sequences': len(cam_sequences), } write_json(cam_split_path, cam_split) console.print(f'[cyan]Camera {cam}: {len(cam_sequences)}/{len(base_sequences)} episodes') raw_config_cam = deepcopy(raw_config) raw_config_cam['path'] = [cam_split_path] # Override all camera fields to load only this camera raw_config_cam['cameras'] = [cam] raw_config_cam.pop('cameras_core', None) raw_config_cam.pop('cameras_core_sample', None) raw_config_cam.pop('cameras_aux', None) raw_config_cam.pop('cameras_aux_sample', None) dataset, config = dataset_from_config( config=raw_config_cam, name=recipe, added=added, subfolder=args.subfolder, ) cam_args = deepcopy(args) cam_args.per_camera = False cam_args._camera_mode_applied = True # skip flexible grouping (already handled) cam_args._skip_metadata_shared = True # parent already copied it cam_args.split_cam = f'split_all__{cam}.json' cam_args.global_cam_idx = all_cameras.index(cam) cam_args.quiet = True cam_args.dataset_obj = dataset cam_tasks.append((cam_args, dataset, config)) # Copy metadata_shared.json once from unified to webbed output if cam_tasks: first_args = cam_tasks[0][0] first_config = cam_tasks[0][2] webbed_name = get_name(cam_tasks[0][1], recipe, args.store, args.stride, first_config, args.snippet_formula, args.camera_mode, getattr(args, 'name_suffix', None)) parent_meta_dst = os.path.join(first_args.local_path, dataset_name, webbed_name) ### METADATA FIX (not applied yet, keeping for future incorporation) # meta_shared = read_json(base_meta_path) # if args.resize is not None and 'resolution' in meta_shared: # meta_shared['resolution'] = parse_resize_params(meta_shared['resolution'][::-1], args.resize) # if args.stride > 1 and 'framerate' in meta_shared: # meta_shared['framerate'] = meta_shared['framerate'] / args.stride # if 'labels' in meta_shared: # for label in meta_shared['labels']: # if label not in args.labels: # meta_shared.pop(label) # if 'labels' in meta_shared: # meta_shared['labels'] = [l for l in meta_shared['labels'] if l in args.labels] # if 'shared_labels' in meta_shared: # meta_shared['shared_labels'] = [l for l in meta_shared['shared_labels'] if l in args.labels] # write_json(f'{parent_meta_dst}/metadata_shared.json', meta_shared) os.makedirs(parent_meta_dst, exist_ok=True) meta_src_path = base_meta_path if os.path.exists(meta_src_path): shutil.copy2(meta_src_path, f'{parent_meta_dst}/metadata_shared.json') # Removing camera parallelism, there are usually more per-camera sequences than cameras outcomes = [] for task in cam_tasks: cam_args, dataset, config = task # Loop over each camera outcome = webdatasetize(cam_args, dataset_config=(dataset, config)) outcomes.append(outcome) if outcome == 'exiting': return # Generate combined split_all.json from all per-camera manifests if cam_tasks and args.subset is None and any(outcome != 'already_processed' for outcome in outcomes): split_name = get_name( dataset, recipe, args.store, args.stride, config, args.snippet_formula, args.camera_mode, getattr(args, 'name_suffix', None)) # NOTE(bvh): use dataset_name (the unified dir basename) here. The old # `args.dataset[0].split('/')[0]` yields '' for absolute split specs, which # broke args.path (double slash) and crashed _generate_split_all on metadata_shared. args.path = f'{args.local_path}/{dataset_name}/{split_name}' first_dataset, first_config = cam_tasks[0][1], cam_tasks[0][2] webbed_name = get_name(first_dataset, recipe, args.store, args.stride, first_config, args.snippet_formula, args.camera_mode, getattr(args, 'name_suffix', None)) run_path = f'{args.local_path}/{dataset_name}/{webbed_name}' _generate_split_all(run_path, raw_config, args) return if dataset_config is None: dataset, config = dataset_from_config( deepcopy(raw_config), name=recipe, added=added, subfolder=args.subfolder, ) else: dataset, config = dataset_config unified_path = config['path'][0] # e.g. /datasets/basile/cv_unified_debug/DROID14/split_all.json if isinstance(unified_path, str) and unified_path.endswith('.json'): unified_path = os.path.dirname(unified_path) # Detect if language latents are needed (for VAE mode) args.need_language_latents = False if args.mode in ["latents", "both"] and len(dataset) > 0: try: sample0 = dataset[0] args.need_language_latents = "language" in sample0 except Exception: labels = getattr(dataset, "labels", []) args.need_language_latents = "language" in labels seqs = list(range(len(dataset))) files = None if args.subfolder is not None: dataset_name = f'{dataset_name}/{args.subfolder}' split_name = get_name( dataset, recipe, args.store, args.stride, config, args.snippet_formula, args.camera_mode, getattr(args, 'name_suffix', None)) # Store derived values in args for child functions (prepare, webdatasetize_sequence) args.dataset_obj = dataset args.name = dataset_name args.split = split_name args.unified_path = unified_path args.metadata_shared_data = metadata_shared # Create split path and save a tmp copy for logging purposes path = f'{args.local_path}/{args.name}/{args.split}/{args.split_cam}' if not path.endswith('.json'): path = f'{path}.json' if args.subset is not None: path = path.replace('.json', f'__{args.subset.replace("/","_")}.json') if dataset_config is not None and args.global_cam_idx is None: # Camera name as suffix (only when NOT already set via per-camera split_cam) cam_suffix = config.get('cameras_core', config.get('cameras', None)) if isinstance(cam_suffix, list) and len(cam_suffix) == 1: path = path.replace('.json', f'__{cam_suffix[0]}.json') dirname, basename = os.path.dirname(path), os.path.basename(path) tmp_path = f'{dirname}/tmp_{basename}' # Check if split already exists, and stop if not restarting if not args.restart: if os.path.exists(path): print(f'### {path} already exists, exiting...') return 'exiting' if os.path.exists(tmp_path): print(f'### {tmp_path} already exists, exiting...') return 'exiting' # Restarting, so warn if files already exist if os.path.exists(path): print(f'### {path} already exists, restarting!!!') rm_sync(path, args) elif os.path.exists(tmp_path): print(f'### {tmp_path} already exists, restarting!!!') rm_sync(tmp_path, args) # Write temporary file to get things started write_json(tmp_path, dict()) upload_manifest = args.upload if upload_manifest: tmp_path_s3 = tmp_path.replace(args.local_path, args.s3_path) aws_s3_cp(tmp_path, tmp_path_s3, robust=True) console.print(f'[bold green]Temp Manifest uploaded:[/bold green] {tmp_path_s3}') # Write metadata_shared.json to webbed output meta_src = args.unified_path meta_dst = os.path.join(args.local_path, args.name, args.split) os.makedirs(meta_dst, exist_ok=True) # Copy metadata_shared.json from unified to webbed (never modify it) if getattr(args, '_skip_metadata_shared', False): pass # Parent per_camera caller already copied it else: meta_shared_src = _resolve_metadata_shared( meta_src, getattr(args, 'metadata_shared', None)) if os.path.exists(meta_shared_src): shutil.copy2(meta_shared_src, f'{meta_dst}/metadata_shared.json') else: print(f'### WARNING: "{meta_shared_src}" not found') if args.upload and os.path.exists(f'{meta_dst}/metadata_shared.json'): meta_dst_s3 = meta_dst.replace(args.local_path, args.s3_path) aws_s3_cp(f'{meta_dst}/metadata_shared.json', f'{meta_dst_s3}/metadata_shared.json', robust=True) # Use Manager().dict() so child processes can share _skipped with the parent mp_manager = multiprocessing.Manager() args._skipped = mp_manager.dict() # Camera name taken from split_cam args.cam = args.split_cam.split("__")[-1][:-len('.json')] if not args.cleanup: from anydata.sync.sync_utils import multi_thread multi_thread = partial(multi_thread, name=f'WEBBING : {args.name} - {args.split} - {args.split_cam}', fn_seqs=webdatasetize_sequences, fn_files=None, tarfiles=True) multi_thread(seqs, files, args) errors = dict(args._skipped) # convert proxy to regular dict mp_manager.shutdown() if seqs and len(errors) == len(seqs) and all(reason == 'already_processed' for reason in errors.values()): return 'already_processed' else: ### Cleaning up, collect errors for logging and move on errors_path = glob(f'{args.local_path}/{args.name}/{args.split}/errors_all__{args.cam}__*.json') if len(errors_path) == 0: errors_path = [f'{args.local_path}/{args.name}/{args.split}/errors_all.json'] errors = dict() for error_path in errors_path: error_json = read_json(error_path)['errors'] error_json = {key: val for key, val in error_json.items() if '[*FILTER*]' not in val} errors.update(error_json) ### Log errors in the webbed folder errors_path = f'{args.local_path}/{args.name}/{args.split}/errors_all__{args.cam}.json' if args.subset is not None: errors_path = errors_path.replace('.json', f'__{args.subset.replace("/","_")}.json') errors_dict = {'total': len(errors), 'errors': {}} for key, val in errors.items(): errors_dict['errors'][key] = val write_json(errors_path, errors_dict) cameras = args.dataset_obj.cameras_core labels = args.dataset_obj.labels # Remove "non-label" labels labels = [l for l in labels if l not in ['lowdim']] sequences = {} for seq in seqs: # Get filename path for logging purposes camera = args.dataset_obj.cameras_core[0] filename = args.dataset_obj.get_filename_target(seq, cam_idx=0, cam=camera)[0] # filename = clean_string(filename) tarfolder = filename.replace(f'{args.unified_path}/', '') tarfolder = '/'.join(tarfolder.split('/')[:-3]) tarpath = f'{args.local_path}/{args.name}/{args.split}/{tarfolder}/{camera}' tarfiles = [os.path.basename(t).replace('.txt', '.tar') for t in sorted(glob(f'{tarpath}/*.t*'))] # Handles both .tar and .txt if len(tarfiles) > 0: sequences[tarfolder] = tarfiles num_files = sum(len(v) for v in sequences.values()) manifest = dict( command='python ' + ' '.join(sys.argv), config=raw_config, info=dict( webbed=args.store, stride=args.stride, num_sequences=len(sequences), num_files=num_files, cameras=sorted(cameras), labels=sorted(labels), ), sequences=sequences, ) write_json(path, manifest) msg = f'[bold green]Manifest written:[/bold green] {path} ({len(sequences)} sequences, {num_files} files)' console.print(msg) if upload_manifest: path_s3 = path.replace(args.local_path, args.s3_path) aws_s3_cp(path, path_s3, robust=True) # Copy split to s3 aws_s3_rm(tmp_path_s3, robust=True) # Delete temporary split from s3 errors_path_s3 = errors_path.replace(args.local_path, args.s3_path) aws_s3_cp(errors_path, errors_path_s3, robust=True) # Copy split to s3 console.print(f'[bold green]Manifest uploaded:[/bold green] {path_s3}/') # Delete local temporary split remove_path(tmp_path) if __name__ == '__main__': args = parse_args() webdatasetize(args)