# Copyright 2026 Toyota Research Institute. All rights reserved. import os import json from typing import Any, Dict, List import torch import random import numpy as np from abc import ABC from tqdm import tqdm from torch.utils.data import Dataset from copy import deepcopy from anydata.utils.geometry import invert_extrinsics from anydata.utils.misc import update_dict, update_dict_nested, get_local_root from anydata.utils.data import shuffle_dict, merge_dicts from anydata.utils.read import read_numpy, read_json, read_video_seq, read_zarr_seq, read_seq, read_npz_node from anydata.utils.types import is_dict, is_list, is_str, is_double_list from anydata.utils.folder_tree import FolderTree from anydata.utils.decorators import iterate1 from anydata.sync.sync_utils import aws_s3_cp from anydata.geometry.depth import decode_dense class InsufficientCamerasError(RuntimeError): """Raised when camera mode is 'skip' and a sample has too few cameras.""" pass ############################################### def remext(file): """Remove extension from filename""" return os.path.splitext(file)[0] def splitjoin(file: str, st=0, fn=None): """Split filename based on / and join back a subset""" return '/'.join(file.split('/')[st:] if fn is None else file.split('/')[st:-fn]) def replace_sagemaker_local_path(path: str | list): """Change local path to be Sagemaker-compliant""" if is_list(path): return [replace_sagemaker_local_path(p) for p in path] elif is_str(path): return path.replace(get_local_root(), '/opt/ml/input/data/training') else: return path def replace_sagemaker_s3_path(path: str | list): """Change s3 path to be Sagemaker-compliant""" if is_list(path): return [replace_sagemaker_s3_path(p) for p in path] elif is_str(path): return path.replace('s3://tri-ml-datasets', 's3://tri-ml-sandbox-16011-us-west-2-datasets') else: return path ############################################### def normalize_metadata(metadata: dict) -> dict: """Normalize metadata from on-disk format to batch format. On-disk: info:{name,tags,raw_id}, rgb:{ext}, depth:{sparse,metric}, ... Batch: name, tags, raw_id (flat), conventions:{rgb,depth,...} """ # Flatten 'info' to top level (on-disk has info:{}, batch has flat fields) if 'info' in metadata: info: dict = metadata.pop('info') for key, val in info.items(): metadata[key] = val # Canonical keys are 'cameras' and 'labels'; fall back to '_available' variants if 'cameras' not in metadata and 'cameras_available' in metadata: metadata['cameras'] = metadata.pop('cameras_available') else: metadata.pop('cameras_available', None) if 'labels' not in metadata and 'labels_available' in metadata: metadata['labels'] = metadata.pop('labels_available') else: metadata.pop('labels_available', None) # Wrap label convention fields under 'conventions' if 'conventions' not in metadata: convention_keys = ['rgb', 'lowdim', 'intrinsics', 'extrinsics', 'depth', 'semantic', 'action'] conventions: dict = {} for key in convention_keys: if key in metadata and isinstance(metadata[key], dict): conventions[key] = metadata.pop(key) # Language is special: promote task/prompt to top level (no convention remainder) if 'language' in metadata and isinstance(metadata['language'], dict): lang: dict = metadata.pop('language') if 'task' in lang: metadata.setdefault('task', lang.pop('task')) if 'prompt' in lang: metadata.setdefault('prompt', lang.pop('prompt')) if 'lowdim' not in conventions: conventions['lowdim'] = {"extension": "npz"} if conventions: metadata['conventions'] = conventions return metadata ############################################### # NOTE(bvh): shared base for Unified and Webbed, but BaseWebbed either overrides or ignores # some methods (such as __getitem__ and get_sample_cameras). class BaseDataset(Dataset, ABC): # Hard-coded for now, we can change it later RGB_FOLDER = 'rgb' DEPTH_FOLDER = 'depth' SEMANTIC_FOLDER = 'semantic' LOWDIM_FOLDER = 'lowdim' OPTFLOW_FOLDER = 'optflow' SCNFLOW_FOLDER = 'scnflow' MASK_RGB_FOLDER = 'mask_rgb' MASK_DEPTH_FOLDER = 'mask_depth' def __init__(self, path, # Path to dataset path_prefix='', # Prefix to add to all paths split=None, # Which split to use tag=None, # Dataset tag base_tag=None, # Tag to return if there is no dataset tag labels=(), # Which labels to use dummy_labels=(), # Which dummy labels to use invert_extrinsics=False, # Invert extrinsics or not cameras_core=(0,), # Core cameras to use cameras_core_sample=None, # Number of core cameras to sample cameras_aux=(), # Auxiliary cameras (used alongside core cameras) cameras_aux_sample=None, # Number of auxiliary cameras to sample cameras_core_mode='strict', # Camera mode: 'strict' (crash on missing) or 'flexible' (skip missing) or 'skip' (skip if insufficient) cameras_aux_mode='strict', # Aux camera mode: same as above length=(), # Which temporal length to use length_sample=None, # Sample temporal length length_stride=1, # Apply stride to length fixed_idx=None, # Always return the same index (for debugging) # ego_camera=None, # DEPRECATED(bvh): use reference_camera in dataset YAML instead. # Set reference_camera to the forward-facing camera name (e.g. 'CAMERA_01') # to make all extrinsics relative to that camera pre-collate. flatten_keys=None, # Flatten (time,cam) keys in spatial or temporal dimension target_flow=False, # Return target optical/scene flow as well mask_range=None, # Mask depth range beyond a certain value clip_depth=None, # Clip depth range beyond a certain value data_transform=None, # Data transformations post_process=True, # Should we post-process the sample or not webbing=False, # Whether we are webbing the dataset or not webbed=False, # Whether the dataset has been webbed or not custom=None, # Custom metadata from config (added to batch metadata) rgb_decoder='torchcodec', # Which video decoder to use for rgb videos **kwargs): super().__init__() # Fix path if using sagemaker self.sagemaker = os.environ.get('SAGEMAKER') == 'enabled' if self.sagemaker: path = replace_sagemaker_local_path(path) # Include path prefix if path_prefix != '': for i in range(len(path)): # Loop over all paths if isinstance(path[i], list): # Webbed datasets can have repeat [path,repeat] path[i][0] = f'{path_prefix}/{path[i][0]}' else: path[i] = f'{path_prefix}/{path[i]}' self.tag = tag if tag is not None else base_tag # Dataset tag # Convert certain entries into lists if not already if not isinstance(length, list): length = [length] if not isinstance(labels, list): labels = [labels] if not isinstance(cameras_core, list): cameras_core = [cameras_core] if isinstance(cameras_aux, tuple): cameras_aux = list(cameras_aux) elif not isinstance(cameras_aux, list): cameras_aux = [cameras_aux] # If 'ALL' is requested in cameras or labels, get from split (webbed case) if self.tag == 'web' and ( (cameras_core == [['ALL']] or '*' in cameras_core[0][0]) or \ (len(cameras_aux) > 0 and (cameras_aux == [['ALL']] or '*' in cameras_aux[0][0])) or \ labels == ['ALL'] or \ length == ['ALL'] ): split_load = path[0][0] if isinstance(path[0], list) else path[0] shared_load = os.path.dirname(split_load) + '/metadata_shared.json' if split_load.startswith('s3://'): # Download split local_root = get_local_root() local_split = os.path.join(local_root, *split_load.split('/')[3:]) local_shared = os.path.join(local_root, *shared_load.split('/')[3:]) os.makedirs(os.path.dirname(local_split), exist_ok=True) os.makedirs(os.path.dirname(local_shared), exist_ok=True) aws_s3_cp(split_load, local_split) aws_s3_cp(shared_load, local_shared) split_load = json.load(open(local_split, 'r')) shared_load = json.load(open(local_shared, 'r')) else: # Split is local already split_load = json.load(open(split_load, 'r')) shared_load = json.load(open(shared_load, 'r')) if cameras_core == [['ALL']]: # Use all cameras cameras_core = [split_load['info']['cameras']] if cameras_aux == [['ALL']]: # Use all aux cameras cameras_aux = [split_load['info']['cameras']] if '*' in cameras_core[0][0]: # Wildcards in cameras split = cameras_core[0][0].split('*') cameras_core = split_load['info']['cameras'] if split[-1] == '' and split[0] != '': cameras_core = [c for c in cameras_core if c.startswith(split[0])] if split[0] == '' and split[-1] != '': cameras_core = [c for c in cameras_core if c.endswith(split[-1])] if split[0] == '' and split[-1] == '': cameras_core = [c for c in cameras_core if split[1] in c] cameras_core = [cameras_core] if len(cameras_aux) > 0 and '*' in cameras_aux[0][0]: # Wildcards in aux cameras split = cameras_aux[0][0].split('*') cameras_aux = split_load['info']['cameras'] if split[-1] == '' and split[0] != '': cameras_aux = [c for c in cameras_aux if c.startswith(split[0])] if split[0] == '' and split[-1] != '': cameras_aux = [c for c in cameras_aux if c.endswith(split[-1])] if split[0] == '' and split[-1] == '': cameras_aux = [c for c in cameras_aux if split[1] in c] cameras_aux = [cameras_aux] if labels == ['ALL']: # Use all labels labels = split_load['info']['labels'] if length == ['ALL']: # Use all length length = [int(split_load['info']['webbed'].split('b')[0])] self.path: List[str] | str = path # Dataset path self.split: str | None = split # Dataset split self.labels: List[str] = labels # Dataset labels self.dummy_labels: List[str] = dummy_labels # Dataset dummy labels self.data_transform = data_transform # Data transforms if rgb_decoder not in ['decord', 'torchcodec']: raise ValueError(f'Invalid RGB video decoder {rgb_decoder}') self.rgb_decoder = rgb_decoder # RGB video decoder self.fixed_idx = fixed_idx # Always return the same index if not None self.flatten_keys = flatten_keys # Flatten (time,cam) keys (spatial/temporal) self.cameras_core = cameras_core # Core cameras self.cameras_core_sample = cameras_core_sample # Number of core cameras to sample self.ego_camera = None # DEPRECATED(bvh): use reference_camera instead # If 'ALL' is requested in cameras or labels, get from split (non-webbed case) if self.cameras_core == ['ALL'] or self.labels == ['ALL']: split_load = json.load(open(self.path[0], 'r')) if self.cameras_core == ['ALL']: self.cameras_core = split_load['info']['cameras'] if self.labels == ['ALL']: self.labels = split_load['info']['labels'] assert len(self.cameras_core) > 0, '### THERE ARE NO VALID CAMERAS, EXITING...' if invert_extrinsics: import warnings warnings.warn( "\n\n### DEPRECATION WARNING ###\n" "invert_extrinsics is deprecated and will be removed.\n" "Extrinsics are now auto-converted to cam2world based on " "metadata.conventions.extrinsics.transform.\n" "Please fix the converter's metadata instead.\n" "###########################\n", DeprecationWarning, stacklevel=2 ) self.invert_extrinsics = invert_extrinsics # Invert extrinsics or not self.post_process = post_process # Post-process sample or not self.custom = custom # Custom metadata from config self.webbing = webbing # Are we webbing or not if self.webbing: self.labels.append('lowdim') # Add lowdim if we are webbing # Auxiliary cameras self.with_cameras_aux = len(cameras_aux) > 0 self.cameras_aux: List[str] = cameras_aux if self.with_cameras_aux else [] # List of all cameras (core + aux) self.all_cameras = [c for c in self.cameras_core] if self.cameras_aux is not None: self.all_cameras.extend(self.cameras_aux) # Sampling parameters (temporal and spatial) self.context_sample = length_sample self.cameras_aux_sample = cameras_aux_sample self.cameras_core_mode = cameras_core_mode self.cameras_aux_mode = cameras_aux_mode self.context_stride = length_stride self.subsample = None # Initialize dataset subsampling # Normalize length values (accept 'all'/'ALL' etc.) length = [l.upper() if isinstance(l, str) else l for l in length] self.single_slice = len(length) == 1 and length[0] == 'ALL' if self.single_slice: length = [] # Validate that remaining length values are integers for l in length: if not isinstance(l, int): raise ValueError(f'Invalid length value: {l} (expected integers or "ALL")') # If length only has 1 number it starts from 0 if len(length) == 1: length = [0, length[0]] # [41] -> [0, 41] # Subtract 1 from forward: frame 0 (anchor) is always included fwd_context = length[1] - 1 if len(length) == 2 and length[1] > 0 else 0 # 41 -> 40 # Take context stride into consideration if self.context_stride is not None: fwd_context *= self.context_stride # Context list: frame indices relative to anchor (excluding anchor itself) self.fwd_context = fwd_context # 40 self.bwd_context = 0 # NOTE: legacy self.context = list(range(1, fwd_context + 1)) # [1, 2, ..., 40] self.num_context = fwd_context # 40 self.with_context = self.num_context > 0 or self.single_slice # Flow modifiers self.target_flow = target_flow # Depth modifiers self.mask_range = mask_range self.clip_depth = clip_depth # Initialize lowdim dict as None self.lowdim = None self.rgb_tree: FolderTree | Dict[str, FolderTree] # Get base metadata if not webbed: metadata = os.path.dirname(self.path[0]) + '/metadata_shared.json' self.base_metadata = read_json(metadata) if os.path.exists(metadata) else None def __len__(self): """Return dataset length""" if self.subsample is not None: # Return subsample length return len(self.subsample) if is_dict(self.rgb_tree): # Multiple cameras, return length of the first return len(list(self.rgb_tree.values())[0]) else: # Single camera, return length return len(self.rgb_tree) def fix_sagemaker(self, filename: str): """Change filename path to follow sagemaker convention""" if self.sagemaker: filename = filename.replace(f'{get_local_root()}/datasets', '/opt/ml/input/data/training/datasets') return filename def get_filename(self, idx: int): """Return filename for a give dataset index""" # Get FolderTree from first camera rgb_tree = self.rgb_tree[self.cameras_core[0]] # Get filename from the requested index filename: str = rgb_tree.get_item(idx, return_loc=False)[0] filename = '__'.join(filename.split('/')[4:]) return filename def __getitem__(self, idx: int, force_camera=None): """Get dataset sample given an index""" sample, idx = self.initialize_sample(idx) # Initialize empty sample cameras = self.get_sample_cameras(idx, force_camera) # Select valid cameras ### METADATA filename, _, loc = self.get_filename_target(idx, cam_idx=0, cam=cameras[0]) metadata_file = '/'.join(filename.split('/')[:-3]) + '/metadata.json' metadata = merge_dicts(json.load(open(metadata_file)), self.base_metadata) self.metadata = sample['metadata'] = normalize_metadata(metadata) self.metadata['path'] = os.path.dirname(filename).replace(self.path[0] + '/', '') # Check if information is stored per-frame or per-sequence. # NOTE(bvh): storage keys were renamed 'frame'->'frames' and 'sequence'->'videos'. # Old metadata.json files may still use 'frame'/'sequence'; both paths handle # them correctly ('frame' falls into the per-frame branch via the default). if self.metadata.get('storage', 'frames') in ['sequence', 'videos']: # Add sequence information sample = self.add_video(idx, sample, cameras) else: # Loop over cameras for cam_idx, cam in enumerate(cameras): # Add target information filename, time_cam, _ = self.get_filename_target(idx, cam_idx=cam_idx, cam=cam) filename = {time_cam: filename} if self.with_context: filename.update(self.get_filename_context(idx, cam_idx=cam_idx, cam=cam)) sample = self.add_frames(sample, filename) # DEPRECATED(bvh): ego-camera system. Use reference_camera= in dataset YAML instead. # if self.ego_camera is not None: # # Add target information (only extrinsics) # filename, time_cam, loc = self.get_filename_target(idx, cam_idx='ego', cam=self.ego_camera) # extrinsics = self.get_extrinsics(filename, time_cam=time_cam, cam=self.ego_camera, loc=loc) # update_dict(sample, 'extrinsics', time_cam, extrinsics) # self.lowdim = None # Reset lowdim # if self.with_context: # # Add context information (only extrinsics) # filename_context = self.get_filename_context(idx, cam_idx='ego', cam=self.ego_camera, loc=loc) # for time_cam, filename in filename_context.items(): # iloc = loc + time_cam[0] # extrinsics = self.get_extrinsics(filename, time_cam=time_cam, cam=self.ego_camera, loc=iloc) # update_dict(sample, 'extrinsics', time_cam, extrinsics) # self.lowdim = None # Reset lowdim # Store camera and index information in metadata sample['metadata']['cameras'] = cameras sample['metadata']['idx'] = idx # Post-process and return sample return self.post_process_sample(sample) if self.post_process else sample def update_filename(self, filename): """Fix filename for sagemaker if required""" return self.fix_sagemaker(filename) def get_filename_target(self, idx: int, cam_idx: int, cam: str): """Get filename for target timestep (0)""" # Pick the FolderTree for specific camera rgb_tree = self.rgb_tree[cam] if is_dict(self.rgb_tree) else self.rgb_tree # Get target filename (plus location in the sequence) filename, loc = rgb_tree.get_item(idx, return_loc=True) # Filename fix for sagemaker filename = self.update_filename(filename[0]) # Return filename, (time,cam), sequence location return filename, (0, cam_idx), loc def get_filename_context(self, idx, cam_idx, cam): """Get filename for context timesteps""" # Pick the FolderTree for specific camera rgb_tree = self.rgb_tree[cam] if is_dict(self.rgb_tree) else self.rgb_tree # Get context filenames filename = rgb_tree.get_context(idx) # Shuffle if sampling context if self.context_sample is not None: filename = shuffle_dict(filename, n=self.context_sample, interval=None) # Apply stride if requested if self.context_stride is not None: bwd = [key for key in filename.keys() if key < 0] fwd = [key for key in filename.keys() if key > 0][::-1] bwd = bwd[::self.context_stride] fwd = fwd[::self.context_stride] filename = {key: val for key, val in filename.items() if key in fwd + bwd} # Filename fix for sagemaker filename = {key: self.update_filename(val) for key, val in filename.items()} # Prepare (time,cam) keys filename = {(time, cam_idx): val for time, val in filename.items()} # Return filename dictionary return filename def get_sample_cameras(self, idx, force_camera=None): ''' [UNIFIED only] Sample core and auxiliary cameras. NOTE(bvh): recommended to only apply shuffle when subsampling, otherwise it scrambles camera->view mapping per sample, causing collate padding at bs>=2. Keep behavior in sync with tariterators.url_opener(). ''' # Sample core cameras cameras = list(self.cameras_core) if self.cameras_core_sample is not None: random.shuffle(cameras) cameras = cameras[:self.cameras_core_sample] # Sample auxiliary cameras cameras_aux = list(self.cameras_aux) cameras_aux = [cam for cam in cameras_aux if cam not in cameras] if self.cameras_aux_sample is not None: random.shuffle(cameras_aux) cameras_aux = cameras_aux[:self.cameras_aux_sample] # Return core + aux cameras return cameras + cameras_aux def initialize_sample(self, idx): """Initialize and return empty sample""" # Use always the same index if resquested if self.fixed_idx is not None: idx = self.fixed_idx # Return empty sample return {}, idx if self.subsample is None else self.subsample[idx] def prepare_language(self, lowdim=None, language=None): if lowdim is not None and 'language' in lowdim: language = lowdim['language'].item() elif language is None: language = dict() # Gather language info from metadata (task/prompt promoted to top level by normalize_metadata) meta_language = {} for k in ['task', 'prompt']: if k in self.metadata: meta_language[k] = self.metadata[k] for key, val in meta_language.items(): if key in language: continue if isinstance(val, list): language[key] = val[0] if len(val) > 0 else '' else: language[key] = val for key, val in language.items(): if isinstance(val, list) and len(val) > 1: language[key] = random.sample(val, 1)[0] return language def prepare_metadata(self, sample, webbed=False): """Prepare batch metadata""" # # Isolate metadata dict for convenience metadata = sample['metadata'] # Remove language-related fields if not required if not self.has_label('language'): sample['metadata'].pop('task', None) sample['metadata'].pop('prompt', None) # Prepare metadata labels and dummy labels if webbed: # Dataset is webbed, take initial list from dataset config metadata['labels'] = [lab for lab in self.labels if lab in sample.keys()] else: # Dataset is not webbed, take initial list from metadata metadata['labels'] = [lab for lab in metadata['labels'] if lab in sample.keys()] metadata['dummy_labels'] = self.added_dummies # Include stride info if not webbed: sample['metadata']['stride'] = self.context_stride # Set sequence length for each camera metadata['length'] = {cam_idx: len([ key for key in sample['rgb'].keys() if key[1] == cam_idx ]) for cam_idx, cam in enumerate(metadata['cameras'])} # Include custom metadata from config if self.custom is not None: metadata['custom'] = self.custom # Dataset has been webbed already, stop here if webbed: return sample # Expand labels so they are per-camera (important for collating) for expand_label in ['resolution','framerate','stride']: metadata[expand_label] = { cam_idx: metadata[expand_label][cam] if isinstance(metadata[expand_label], dict) else metadata[expand_label] for cam_idx, cam in enumerate(metadata['cameras']) if cam in self.all_cameras } # Scale timestep by framerate timestep_start = deepcopy(sample['timestep'][(0,0)]) framerate_start = sample['metadata']['framerate'][0] if isinstance(timestep_start, list): timestep_start = np.array(timestep_start) for key, val in sample['timestep'].items(): if isinstance(val, list): val = np.array(val) sample['timestep'][key] = val - timestep_start framerate = sample['metadata']['framerate'][key[1]] if framerate > 0: sample['timestep'][key] = sample['timestep'][key] / framerate if framerate_start > 0: timestep_start = timestep_start / framerate_start sample['metadata']['timestep_start'] = timestep_start # Remove non-relevant keys for key in ['num_frames']: sample['metadata'].pop(key, None) # Return updated sample return sample def post_process_sample(self, sample, webbed=False): """Post-process sample (augmentations, to tensor, metadata)""" # Legacy: invert extrinsics if explicitly requested (deprecated) if self.invert_extrinsics and 'extrinsics' in sample.keys(): for key in sample['extrinsics'].keys(): sample['extrinsics'][key] = invert_extrinsics(sample['extrinsics'][key]) # Auto-convert extrinsics to cam2world if needed if 'extrinsics' in sample.keys(): conventions = sample['metadata'].get('conventions', {}) ext_conv = conventions.get('extrinsics', {}) transform = ext_conv.get('transform', None) if transform == 'world2cam': for key in sample['extrinsics'].keys(): sample['extrinsics'][key] = invert_extrinsics(sample['extrinsics'][key]) ext_conv['transform'] = 'cam2world' # Flatten dictionary keys if requested (temporal or spatial) if self.flatten_keys is not None: tgt = (0, 0) tuples = list([v for v in sample['rgb'].keys() if v != tgt]) for key in sample.keys(): if is_dict(sample[key]): new_sample_key = {tgt: sample[key][tgt]} for i, tup in enumerate(tuples): # Temporal (cam = 0 always) if self.flatten_keys.startswith('temporal'): new_sample_key[(i + 1, 0)] = sample[key][tup] # Spatial (time = 0 always) elif self.flatten_keys.startswith('spatial'): new_sample_key[(0, i + 1)] = sample[key][tup] # Invalid else: raise ValueError('Invalid flatten key') sample[key] = new_sample_key # Check whether random sampling is requested as well (e.g., temporal-5, or spatial-3) if '-' in self.flatten_keys: num = int(self.flatten_keys.split('-')[-1]) - 1 rnd_keys = [l for l in list(sample['rgb'].keys()) if l != tgt] random.shuffle(rnd_keys) rnd_keys = rnd_keys[:num] for key in sample.keys(): if is_dict(sample[key]): new_sample_key = {tgt: sample[key][tgt]} for i, rnd_key in enumerate(rnd_keys): if self.flatten_keys.startswith('temporal'): new_sample_key[(i + 1, 0)] = sample[key][rnd_key] elif self.flatten_keys.startswith('spatial'): new_sample_key[(0, i + 1)] = sample[key][rnd_key] else: raise ValueError('Invalid flatten key') sample[key] = new_sample_key # Apply data transformation if self.data_transform: sample = self.data_transform(sample, no_to_tensor=self.webbing, webbed=webbed) # Mask depth range if requested if self.mask_range is not None: for key in ['depth']: for tgt in sample[key].keys(): invalid = (sample[key][tgt] < self.mask_range[0]) | (sample[key][tgt] > self.mask_range[1]) sample[key][tgt][invalid] = 0.0 # Clip maximum depth if requested if self.clip_depth is not None: for key in ['depth']: for tgt in sample[key].keys(): invalid = sample[key][tgt] > self.clip_depth sample[key][tgt][invalid] = self.clip_depth # Sample a given context if requested if self.context_sample is not None: tgt = 0 ctx = [tgt] + list(set([v[0] for v in sample['rgb'].keys() if v[0] != tgt])) if is_str(self.context_sample): if '/' in self.context_sample: previous = [c for c in ctx if c < tgt] forwards = [c for c in ctx if c > tgt] new_ctx_previous = list(range(- len(previous), 0)) new_ctx_forwards = list(range(1, len(forwards) + 1)) map_dict_previous = {c: nc for c, nc in zip(previous, new_ctx_previous)} map_dict_forwards = {c: nc for c, nc in zip(forwards, new_ctx_forwards)} map_dict = {tgt: tgt, **map_dict_previous, **map_dict_forwards} else: new_ctx = [tgt] + list(range(1, len(ctx) + 1)) map_dict = {c: nc for c, nc in zip(ctx, new_ctx)} for key in sample.keys(): if is_dict(sample[key]): sample[key] = {(map_dict[k[0]], k[1]): v for k, v in sample[key].items()} # Add dummy labels sample = self.add_dummy_labels(sample) # Certain labels are not dependent on camera (especially lowdim), so remove "cam" from the (time,cam) keys. # NOTE(bvh): collapsing (time,cam)->time is only safe because these keys are typically # loaded only at cam_idx=0 (see add_frames / add_video). If ever populated at cam>0, # same-timestep entries silently collapse to last-inserted-wins. strip_labels = ['action', 'language'] + [k for k in sample if k.startswith('ego')] for time_label in strip_labels: if time_label in sample and isinstance(sample[time_label], dict): sample[time_label] = {key[0]: val for key, val in sample[time_label].items()} # Prepare metadata sample = self.prepare_metadata(sample, webbed=webbed) # Write labels back to lowdim so we can store in webdataset if self.webbing: for label in ['timestep','intrinsics','extrinsics']: if label in sample.keys(): for key in sample['lowdim'].keys(): sample['lowdim'][key][label] = sample[label][key] else: # Remove label from lowdim if not requested for key in sample['lowdim'].keys(): sample['lowdim'][key].pop(label, None) for label in ['action','language']: if label in sample.keys(): for key in sample['lowdim'].keys(): sample['lowdim'][key][label] = sample[label][key[0]] else: # Remove label from lowdim if not requested for key in sample['lowdim'].keys(): sample['lowdim'][key].pop(label, None) # Return sample return sample def add_video(self, idx, sample, cameras): # Conventions for each label labels = self.metadata['labels'] conventions: Dict[str, Dict[str, Any]] = self.metadata['conventions'] # Loop over cameras for cam_idx, cam in enumerate(cameras): # Get filename and start/finish locations length = np.diff(self.rgb_tree[cam].get_slice(idx)).item() filename, _, loc = self.get_filename_target(idx, cam_idx=cam_idx, cam=cam) st, fn = (0, length) if self.single_slice else (loc, loc + self.fwd_context + 1) indexer = slice(st, fn, self.context_stride) container_path = os.path.dirname(filename) lowdim_path = container_path.format(label=self.LOWDIM_FOLDER) lowdim_ext = conventions.get('lowdim', {}).get('extension', 'npz') lowdim = read_seq(lowdim_path, lowdim_ext, index=indexer, sto=self.metadata.get('storage', 'frames')) indices = list(range(st, fn, self.context_stride)) num_indices = len(indices) ## LOWDIM for webbing if self.has_label('lowdim'): for i in range(num_indices): time_cam = (i * self.context_stride, cam_idx) frame_dict = {"camera": cam, **read_npz_node(lowdim, index=i)} update_dict(sample, 'lowdim', time_cam, frame_dict) ### RGB if 'rgb' in labels and self.has_label('rgb'): path = container_path.format(label=self.RGB_FOLDER) ext = conventions.get('rgb', {}).get('extension', 'mp4') array = read_seq(path, ext, index=indexer, pil=True, mp4_decoder=self.rgb_decoder) for i, rgb in enumerate(array): time_cam = (i * self.context_stride, cam_idx) update_dict(sample, 'rgb', time_cam, rgb) ### INTRINSICS if 'intrinsics' in labels and self.has_label('intrinsics'): for i in range(num_indices): time_cam = (i * self.context_stride, cam_idx) intrinsics = read_npz_node(lowdim["intrinsics"], index=i) update_dict(sample, 'intrinsics', time_cam, intrinsics) ### EXTRINSICS if 'extrinsics' in labels and self.has_label('extrinsics'): for i in range(num_indices): time_cam = (i * self.context_stride, cam_idx) extrinsics = read_npz_node(lowdim["extrinsics"], index=i) update_dict(sample, 'extrinsics', time_cam, extrinsics) ### BBOX2D if 'bbox2d' in labels and self.has_label('bbox2d'): bbox2d = read_npz_node(lowdim["bbox2d"]) for i, ind in enumerate(indices): time_cam = (i * self.context_stride, cam_idx) update_dict(sample, 'bbox2d', time_cam, bbox2d[i]) ### BBOX3D if 'bbox3d' in labels and self.has_label('bbox3d'): bbox3d = read_npz_node(lowdim["bbox3d"]) for i, ind in enumerate(indices): time_cam = (i * self.context_stride, cam_idx) update_dict(sample, 'bbox3d', time_cam, bbox3d[i]) ### MAP if 'map' in labels and self.has_label('map'): map_features = read_npz_node(lowdim["map"]) for i, ind in enumerate(indices): time_cam = (i * self.context_stride, cam_idx) update_dict(sample, 'map', time_cam, map_features[i]) ### MASK_DEPTH if 'mask_depth' in labels and (self.has_label('mask_depth') or self.apply_mask_depth): path = container_path.format(label=f'mask_{self.DEPTH_FOLDER}') ext = conventions.get('mask_depth', {}).get('extension', 'zarr') array = read_seq(path, ext, key="mask_depth", index=indexer) for i, mask_depth in enumerate(array): time_cam = (i * self.context_stride, cam_idx) update_dict(sample, 'mask_depth', time_cam, mask_depth) ### DEPTH if 'depth' in labels and self.has_label('depth'): path = container_path.format(label=self.DEPTH_FOLDER) ext = conventions.get('depth', {}).get('extension', 'zarr') array = read_seq(path, ext, key="depth", index=indexer) for i, depth in enumerate(array): time_cam = (i * self.context_stride, cam_idx) if self.apply_mask_depth and 'mask_depth' in labels: depth *= sample['mask_depth'][time_cam] update_dict(sample, 'depth', time_cam, depth) ### SEMANTIC if 'semantic' in labels and self.has_label('semantic'): path = container_path.format(label=self.SEMANTIC_FOLDER) ext = conventions.get('semantic', {}).get('extension', 'zarr') array = read_seq(path, ext, key="semantic", index=indexer) for i, semantic in enumerate(array): time_cam = (i * self.context_stride, cam_idx) update_dict(sample, 'semantic', time_cam, semantic) ### OPTFLOW for mode in ['bwd', 'fwd']: if f'optflow_{mode}' in labels and self.has_label('optflow'): path = container_path.format(label=f'{self.OPTFLOW_FOLDER}_{mode}') ext = conventions.get('optflow', {}).get('extension', 'zarr') array = read_seq(path, ext, key=f'optflow_{mode}', index=indexer) array = np.transpose(array, (0, 2, 3, 1)) for i, optflow in enumerate(array): if optflow.mean() != -1: # if it's valid time_cam = (i * self.context_stride, cam_idx) optflow[..., 0] *= optflow.shape[0] optflow[..., 1] *= optflow.shape[1] update_dict_nested(sample, 'optflow', time_cam, mode, optflow) ### ACTION (only camera 0) if 'action' in labels and self.has_label('action') and cam_idx == 0: for i in range(num_indices): time_cam = (i * self.context_stride, cam_idx) action = read_npz_node(lowdim["action"], index=i) update_dict(sample, 'action', time_cam, action) ### EGO_* (only camera 0, opt-in via 'action' label) if self.has_label('action') and cam_idx == 0: for ego_key in lowdim: if not ego_key.startswith('ego'): continue for i in range(num_indices): time_cam = (i * self.context_stride, cam_idx) ego_val = read_npz_node(lowdim[ego_key], index=i) update_dict(sample, ego_key, time_cam, ego_val) ### LANGUAGE (only camera 0) if 'language' in labels and self.has_label('language') and cam_idx == 0: for i in range(num_indices): time_cam = (i * self.context_stride, cam_idx) language = read_npz_node(lowdim["language"], index=i) if "language" in lowdim else {} language = self.prepare_language(language=language) update_dict(sample, 'language', time_cam, language) ### TIMESTEP for i in range(num_indices): time_cam = (i * self.context_stride, cam_idx) timestep = read_npz_node(lowdim["timestep"], index=i) update_dict(sample, 'timestep', time_cam, np.array([float(timestep)])) # Return sample return sample def add_frames(self, sample, filename_context): """Add context information to sample""" # Conventions for each label labels = self.metadata['labels'] conventions = self.metadata['conventions'] # Loop over all context for time_cam, filename in filename_context.items(): ### RGB if 'rgb' in labels and self.has_label('rgb'): update_dict(sample, 'rgb', time_cam, self.get_rgb(filename)) ### INTRINSICS if 'intrinsics' in labels and self.has_label('intrinsics'): update_dict(sample, 'intrinsics', time_cam, self.get_intrinsics(filename)) ### EXTRINSICS if 'extrinsics' in labels and self.has_label('extrinsics'): update_dict(sample, 'extrinsics', time_cam, self.get_extrinsics(filename)) ### DEPTH if 'depth' in labels and self.has_label('depth'): update_dict(sample, 'depth', time_cam, self.get_depth(filename)) ### SEMANTIC if 'semantic' in labels and self.has_label('semantic'): update_dict(sample, 'semantic', time_cam, self.get_semantic(filename)) ### LOWDIM if self.has_label('lowdim'): update_dict(sample, 'lowdim', time_cam, self.get_lowdim(filename)) ### MASK RGB if self.has_label('mask_rgb'): update_dict(sample, 'mask_rgb', time_cam, self.get_mask_rgb(filename)) ### MASK DEPTH if self.has_label('mask_depth'): update_dict(sample, 'mask_depth', time_cam, self.get_mask_depth(filename)) ### OPTICAL FLOW if self.has_label('optflow'): update_dict_nested(sample, 'optflow', time_cam, 'bwd', self.get_optflow(filename, 'bwd')) update_dict_nested(sample, 'optflow', time_cam, 'fwd', self.get_optflow(filename, 'fwd')) ### SCENE FLOW if self.has_label('scnflow'): update_dict_nested(sample, 'scnflow', time_cam, 'bwd', self.get_scene_flow(filename, 'bwd')) update_dict_nested(sample, 'scnflow', time_cam, 'fwd', self.get_scene_flow(filename, 'fwd')) ### BBOX2D if self.has_label('bbox2d'): update_dict(sample, 'bbox2d', time_cam, self.get_bbox2d(filename)) ### BBOX3D if self.has_label('bbox3d'): update_dict(sample, 'bbox3d', time_cam, self.get_bbox3d(filename)) ### ACTION (only camera 0) if self.has_label('action') and time_cam[1] == 0: update_dict(sample, 'action', time_cam, self.get_action(filename)) ### EGO_* (only camera 0, opt-in via 'action' label) if self.has_label('action') and time_cam[1] == 0: for ego_key, ego_val in self.get_ego_fields(filename).items(): update_dict(sample, ego_key, time_cam, ego_val) ### LANGUAGE (only camera 0) if self.has_label('language') and time_cam[1] == 0: update_dict(sample, 'language', time_cam, self.get_language(filename)) ### TIMESTEP update_dict(sample, 'timestep', time_cam, self.get_timestep(filename)) # Reset lowdim self.lowdim = None # Return sample return sample def relative_path(self, filename): """Make filename path relative to dataset path""" return {key: os.path.splitext(val.replace(self.path + '/', ''))[0] for key, val in filename.items()} def within_context(self, time, direction): """Check if a time key is within the desired context""" if len(self.context) == 0: return False if direction == 'fwd': return time < self.context[-1] if direction == 'bwd': return time > 0 return False def add_dummy_labels(self, sample): """Add dummy labels if they are not presented""" self.added_dummies = [] for label in self.dummy_labels: # Create label key if not present if label not in sample: sample[label] = dict() # Loop over all keys for key, val in sample['rgb'].items(): if key not in sample[label].keys(): # Add dummy label to list for metadata if label not in self.added_dummies: self.added_dummies.append(label) ### DUMMY DEPTH if label == 'depth': sample[label][key] = torch.zeros(1, *val.shape[1:]) ### DUMMY INTRINSICS elif label == 'intrinsics': sample[label][key] = torch.tensor([ [val.shape[2] / 2, 0.0, val.shape[2] / 2], [0.0, val.shape[1] / 2, val.shape[1] / 2], [0.0, 0.0, 1.0], ]) ### DUMMY EXTRINSICS elif label == 'extrinsics': sample[label][key] = torch.eye(4) ### DUMMY LANGUAGE elif label == 'language': sample[label][key] = {'prompt': 'do something'} else: raise ValueError(f'Invalid dummy label {label}') return sample def has_label(self, label: str): """Check if a label is used""" return any([l.startswith(label) for l in self.labels]) def get_timestep(self, filename): """Get timestep given a filename and (time,cam) key""" raise NotImplementedError def get_rgb(self, filename): """Get RGB image given a filename and (time,cam) key""" raise NotImplementedError def get_intrinsics(self, filename): """Get intrinsics given a filename and (time,cam) key""" raise NotImplementedError def get_extrinsics(self, filename): """Get extrinsics given a filename and (time,cam) key""" raise NotImplementedError def get_depth(self, filename): """Get depth given a filename and (time,cam) key""" raise NotImplementedError def get_semantic(self, filename): """Get semantic segmentation given a filename and (time,cam) key""" raise NotImplementedError def get_lowdim(self, filename): """Get low-dimensional labels given a filename and (time,cam) key""" raise NotImplementedError def get_mask_rgb(self, filename): """Get RGB mask given a filename and (time,cam) key""" raise NotImplementedError def get_mask_depth(self, filename): """Get depth mask given a filename and (time,cam) key""" raise NotImplementedError def get_optflow(self, filename, mode): """Get optical flow given a filename and (time,cam) key""" raise NotImplementedError def get_scene_flow(self, filename, mode): """Get scene flow given a filename and (time,cam) key""" raise NotImplementedError def get_action(self, filename): """Get action label given a filename and (time,cam) key""" raise NotImplementedError def get_bbox2d(self, filename): """Get 2D bounding boxes given a filename and (time,cam) key""" raise NotImplementedError def get_bbox3d(self, filename): """Get 3D bounding boxes given a filename and (time,cam) key""" raise NotImplementedError def get_language(self, filename): """Get language instruction given a filename and (time,cam) key""" raise NotImplementedError