import os import warnings from copy import deepcopy from functools import partial from itertools import islice from typing import Dict, Iterable, List, Any import numpy as np import webdataset as wbd from anydata.geometry.depth import decode_dense from anydata.dataloaders.Base import BaseDataset, replace_sagemaker_s3_path, normalize_metadata from anydata.sync.sync_utils import aws_s3_cp from anydata.utils.data import make_list from anydata.utils.read import read_npz_node from anydata.utils.webdataset import register_s3_gopen def pytorch_worker_info(group=None): """Return node and worker info for PyTorch and some distributed environments.""" rank, world_size, worker, num_workers = 0, 1, 0, 1 if "RANK" in os.environ and "WORLD_SIZE" in os.environ: rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) else: try: import torch.distributed if torch.distributed.is_available() and torch.distributed.is_initialized(): group = group or torch.distributed.group.WORLD rank = torch.distributed.get_rank(group=group) world_size = torch.distributed.get_world_size(group=group) except ModuleNotFoundError: pass if "WORKER" in os.environ and "NUM_WORKERS" in os.environ: worker = int(os.environ["WORKER"]) num_workers = int(os.environ["NUM_WORKERS"]) else: try: import torch.utils.data worker_info = torch.utils.data.get_worker_info() if worker_info is not None: worker = worker_info.id num_workers = worker_info.num_workers except ModuleNotFoundError: pass return rank, world_size, worker, num_workers def split_by_worker_and_node(src, single_gpu=False): rank, world_size, worker, num_workers = pytorch_worker_info() if int(os.environ['PARALLELISM']) > 1: rank = rank // int(os.environ['PARALLELISM']) if single_gpu: rank, world_size = 0, 1 if num_workers > 1 or world_size > 1: src = islice( src, worker + rank * num_workers, None, num_workers * world_size, ) try: while True: yield next(src) except: pass class BaseWebbedDataset(BaseDataset): def __init__(self, pipe_mode='cp', split_mode='worker_and_node', shard_type='simple', parallelism=None, single_gpu=False, multi_tarfiles='auto', multi_tarfiles_multiplier=1.25, subsample=None, **kwargs, ): kwargs['webbed'] = True super().__init__(**kwargs, base_tag='web') self.path = make_list(self.path) self.pipe_mode = pipe_mode self.initialized = False if parallelism is not None: os.environ["PARALLELISM"] = str(parallelism) else: os.environ["PARALLELISM"] = str(0) from anydata.utils.read import read_json tarfiles = [] cameras_core = dict() cameras_aux = dict() for i, path in enumerate(self.path): if isinstance(path, list): path, repeat = path else: repeat = 1 # Check if tarfiles are stored locally on on s3 is_s3 = isinstance(path, str) and path.startswith('s3://') if is_s3: # Get tarfiles from s3 register_s3_gopen() if self.sagemaker: path = replace_sagemaker_s3_path(path) local_manifest = f'/tmp/anydata_manifest_{os.getpid()}_{i}.json' aws_s3_cp(path, local_manifest) data = read_json(local_manifest) dirname = os.path.dirname(path) else: # Get tarfiles locally data = read_json(path) dirname = os.path.dirname(path) if data is None: raise FileNotFoundError(f'Could not read webdataset manifest: {path}') # NOTE(bvh): this stride refers to the FRAME stride used during webbing # (e.g. stride=6 means frames 0,6,12,... were stored). Not to be confused # with the buffer/context stride used for sub-clip sampling. self.stride = data.get('info', {}).get('stride', 1) # Get files from split files = data['sequences'] # Sort properly based on int values for key in files.keys(): files[key] = sorted(files[key], key=lambda x: int(x.split('_')[0])) # if it's a dict, stack all values with keys as prefix if isinstance(files, dict): files_flatten = [] for key, val in files.items(): if multi_tarfiles != 1: # Create list of multiple tarfiles to load simultaneously if multi_tarfiles == 'auto': # Automatically determine the number of tarfiles # Episode bounds parsed from tarfile names like '0000_0017.tar'. # per_tarfile = avg raw frames per tarfile across the episode. start, finish = int(val[0][:-4].split('_')[0]), int(val[-1][:-4].split('_')[1]) length = self.fwd_context * self.stride + 1 per_tarfile = (finish - start) / len(val) # Multiply by multi_tarfiles_multiplier to ensure diversity when randomly sampling multi = int(np.ceil(length * multi_tarfiles_multiplier / per_tarfile)) multi = max(1, min(multi, len(val))) if finish - start < length: # Remove episode if it's too short warnings.warn( f"BaseWebbed: skipping episode {key!r} (path={dirname}) -- " f"only {finish - start} raw frames available but " f"fwd_context*stride={self.fwd_context}*{self.stride}+1={length} " f"are required.") continue else: # Use the provided value multi = multi_tarfiles val = [f'{key}/%s/{v}' for v in val] # %s is a placeholder for camera name val = [[val[j+k] for k in range(multi)] for j in range(len(val) - multi + 1)] else: val = [f'{key}/%s/{v}' for v in val] # %s is a placeholder for camera name files_flatten.extend(val) files = files_flatten # Get urls from files if multi_tarfiles != 1: # TODO: should this not be `if multi != 1` instead ?? urls = [[f'{dirname}/{f}' for f in ff] for ff in files] else: urls = [f'{dirname}/{f}' for f in files] tarfiles.extend(urls * repeat) # TODO(bvh): implement smarter sampling that oversamples tarfiles from shorter # episodes for balancing (currently uniform across all snippets, biased toward # longer episodes). Could use inverse-length weights with ResampledShards. # Set cameras core and aux from split cameras_core[dirname] = self.cameras_core[i] cameras_aux[dirname] = None if len(self.cameras_aux) == 0 else self.cameras_aux[i] # Uniform subsample (or supersample): select N evenly spaced tarfiles, # same as Unified's np.linspace. When N > len(tarfiles), duplicates entries # to reach the target count (supersample for small datasets). if subsample is not None and subsample > 0 and subsample != len(tarfiles): indices = [int(n) for n in np.linspace(0, len(tarfiles) - 1, subsample)] tarfiles = [tarfiles[i] for i in indices] self.num_tarfiles = len(tarfiles) self.split_mode = { 'worker_and_node': partial(split_by_worker_and_node, single_gpu=single_gpu), }[split_mode] self.shard_type_name = shard_type self.shard_type = { 'simple': wbd.SimpleShardList(tarfiles), 'resampled': wbd.ResampledShards(tarfiles), }[shard_type] self.dataset = wbd.DataPipeline( self.shard_type, self.split_mode, wbd.tarfile_to_samples( handler=wbd.warn_and_continue, cameras=dict( multi_tarfiles=multi_tarfiles, core=cameras_core, aux=cameras_aux, core_sample=self.cameras_core_sample, aux_sample=self.cameras_aux_sample, core_mode=self.cameras_core_mode, aux_mode=self.cameras_aux_mode, dirname=dirname, ), ), wbd.decode('pil', handler=wbd.warn_and_continue), self.decode_sample, ) # NOTE(bvh): keep a reference to the pipeline so __getitem__ can re-iterate # it after exhaustion (setting initialized=False). Without this, the line # `self.dataset = iter(self.dataset)` in __getitem__ overwrites the pipeline # with a single-use generator, and subsequent iter() calls return the same # exhausted generator -- breaking repeated validation passes. self._pipeline = self.dataset self.dict_keys = {'resolution','framerate','stride','length','resolution_original','augmentations'} def _parse_time_cam(self, key: str): """Parse (time, cam_idx) from a data key like 'rgb.(3,1).96.jpg'.""" tc_str = key.split('.')[-3] # '(3,1)' or (3:10,1) times, cam_idx = tc_str[1:-1].split(',') time_st, time_fn = map(int, times.split(':')) if ':' in times else (int(times), None) return time_st, time_fn, int(cam_idx) def _decode_value(self, key: str, data: dict | np.ndarray, meta: Dict[str, Any | Dict[str, Any | dict]]): """Decode a data entry based on its label prefix.""" label = key.split('.')[-4] if label == 'lowdim': return {k: v[()] for k, v in data.items()} if label == 'rgb' and self.has_label('rgb'): return data if label == 'depth' and self.has_label('depth'): encode_info = data.get('encode_info', None) if encode_info is not None: encode_info = encode_info.item() # recover dict return decode_dense(data['depth'], encode_info=encode_info) return None def decode_sample(self, data: Iterable[Dict[str, Any | dict]]): for data_key in data: info = data_key['__info__'] core_aux_loaded: list = info['core_aux_loaded'] # core+aux cameras that opened successfully # Separate data keys from internal/metadata keys entry_keys = [k for k in data_key.keys() if not k.startswith(('__', 'metadata'))] # Build cam_idx -> camera name mapping from the lowdim 'camera' field. # This is ground truth: each lowdim entry stores the camera name it # was written for during webbing. idx_to_cam = {} for k in entry_keys: if not k.startswith('lowdim'): continue _, time_fn, cam_idx = self._parse_time_cam(k) if cam_idx not in idx_to_cam: cam_name = data_key[k]['camera'] idx_to_cam[cam_idx] = str(cam_name if time_fn is None else cam_name[0]) # Build the output camera mapping: only keep cameras in core_aux_loaded, # and remap their indices to 0, 1, 2, ... in core_aux_loaded order. cam_remap = {} # old cam_idx -> new cam_idx for old_idx, cam_name in idx_to_cam.items(): if cam_name in core_aux_loaded: cam_remap[old_idx] = core_aux_loaded.index(cam_name) # Build metadata from the first shard's metadata entry meta_names = [k for k in data_key.keys() if k.startswith('metadata')] meta_all = [data_key[meta_name]['metadata'].item() for meta_name in meta_names] metadata = normalize_metadata(deepcopy(meta_all[0])) # Set tarfile path (with %s camera placeholder) for traceability metadata['path'] = data_key.get('__url__', '') # Set cameras to the core_aux_loaded selection metadata['cameras'] = list(core_aux_loaded) # Expand per-camera dict_keys (framerate, stride, etc.) to be # indexed by new camera index. For per-camera metadata, the values # are scalars; for pre-merged, they're already dicts. for dk in self.dict_keys.intersection(metadata.keys()): metadata[dk] = { i: v[dk][0] if isinstance(v[dk], dict) else v[dk] for i, v in enumerate(meta_all) } # Find time offset so time starts at 0 min_time = min(self._parse_time_cam(k)[0] for k in entry_keys) # Build the output sample sample = {'metadata': metadata} for key in entry_keys: time, time_end, cam_idx = self._parse_time_cam(key) if cam_idx not in cam_remap: continue # camera not in core_aux_loaded, skip t_start, cam = (time - min_time, cam_remap[cam_idx]) val = self._decode_value(key, data_key[key], metadata) if val is None: continue if key.startswith('lowdim'): # Unpack lowdim into separate labels labels_to_unpack = ['intrinsics', 'extrinsics', 'timestep', 'language', 'action', 'bbox2d', 'bbox3d'] # NOTE(bvh): auto-pickup of ego_* fields when 'action' is requested if self.has_label('action'): labels_to_unpack += [k for k in val.keys() if k.startswith('ego')] for label in labels_to_unpack: if label not in val: continue if not label.startswith('ego') and not self.has_label(label) and label != 'timestep': continue if label not in sample: sample[label] = {} if time_end is None: sample[label][(t_start, cam)] = val[label] else: for t in range(0, time_end - time, self.stride): sample[label][(t_start+t, cam)] = read_npz_node(val[label], index=t // self.stride) else: label = key.split('.')[-4] if label not in sample: sample[label] = {} if time_end is None: sample[label][(t_start, cam)] = val else: for t in range(0, time_end - time, self.stride): sample[label][(t_start+t, cam)] = val[t // self.stride] # Sort (time, cam) keys for deterministic ordering sort_labels = [l for l in sample if l != 'metadata' and isinstance(sample[l], dict)] if sort_labels: sorted_keys = sorted(sample[sort_labels[0]].keys(), key=lambda x: (x[0], x[1])) for label in sort_labels: sample[label] = {k: sample[label][k] for k in sorted_keys if k in sample[label]} yield sample # NOTE(bvh): __len__ must be implemented in subclasses def get_sample(self, idx): """Get sample from webdataset""" web_sample = next(self.dataset) sample, idx = self.initialize_sample(idx) sample.update(web_sample) return sample def reset_iter(self): """Rebuild the iterator over self._pipeline for a fresh pass. Called by the trainer at the start of each validation round so repeated passes deterministically re-consume the same leading subsequence of shards. Safe to call any time; cheap (does not re-download shards). """ self.dataset = iter(self._pipeline) self.initialized = True def __getitem__(self, idx): """Get dataset sample given an index""" if not self.initialized: self.reset_iter() while True: sample = self.get_sample(idx) return sample