# Copyright 2026 Toyota Research Institute. All rights reserved. import os import copy import json import random import numpy as np from typing import Dict from argparse import Namespace from anydata.converters.utils import S3_PATH from anydata.dataloaders.Base import BaseDataset, InsufficientCamerasError from anydata.geometry.depth import decode_dense from anydata.utils.folder_tree import FolderTree from anydata.utils.data import make_list from anydata.utils.misc import get_local_root from anydata.utils.read import read_image, read_depth, read_json, read_config, read_png8, read_png32, read_npz from anydata.utils.types import is_dict, is_list from anydata.sync.download_unified import download from anydata.utils.write import frame_name ############################################### class UnifiedDataset(BaseDataset): def __init__(self, root=None, # Root folder to apply to all paths language_dict=None, # Use language dictionary to get prompts break_sequences=None, # Break sequences into smaller chunks return_all_prompts=False, # Return all language prompts apply_mask_depth=None, # Apply valid mask to rgb images subsample=None, # Subsample sequences ### FolderTree parameters stride=1, folders_start=None, folders_end=None, folders_with=None, folders_without=None, first_only=False, keep_first=None, keep_folders=None, remove_folders=None, pad_numbers=None, pad_first=None, pad_last=None, crop_first=None, crop_last=None, subset=None, filter_camera=None, ### End of FolderTree parameters **kwargs ): super().__init__(**kwargs, base_tag='Unified') # Flag if we are filtering cameras for webbing self.filter_camera = filter_camera # If there are multiple paths (it is a list), prepare each separately if is_list(self.path): if self.split is None: self.split = [None] * len(self.path) for i in range(len(self.path)): if self.path[i].endswith('.json'): path = self.path[i].split('/') self.path[i] = '/'.join(path[:-1]) self.split[i] = path[-1].replace('.json', '') # There is only one path else: if self.path.endswith('.json'): path = self.path.split('/') self.path = '/'.join(path[:-1]) self.split = path[-1].replace('.json', '') # If root is provided, append it to all paths if root is not None: if is_list(self.path): self.path = [f'{root}/{path}' for path in self.path] else: self.path = f'{root}/{self.path}' # Get all used cameras cameras = list(set(self.cameras_core + self.cameras_aux)) # Account for ego camera all_cameras = cameras if self.ego_camera is not None and self.ego_camera not in all_cameras: all_cameras += [self.ego_camera] # Set relevant variables self.subset = subset self.return_all_prompts = return_all_prompts self.kwargs = kwargs # Select masks to use self.apply_mask_depth = apply_mask_depth # Initialize one FolderTree for each camera self.rgb_tree: Dict[str, FolderTree] = {} for i, cam in enumerate(all_cameras): # Get dataset paths and splits paths, splits = self.path, self.split if not is_list(paths): paths = [paths] # If splits are provided if not is_list(splits): # Repeat same split for all cameras splits = [splits] * len(paths) assert len(splits) == len(paths) # Get path for each FolderTree folder_trees = [] for path, split in zip(paths, splits): tree_path = self.get_rgb_tree_path(path, split, cam, folders_start, folders_end) folder_tree = FolderTree( tree_path, stride=stride, context=self.context, folders_start=folders_start, folders_end=folders_end, keep_folders=keep_folders, remove_folders=remove_folders, folders_with=folders_with, folders_without=folders_without, first_only=first_only, keep_first=keep_first, pad_first=pad_first, pad_last=pad_last, crop_first=crop_first, crop_last=crop_last, subset=subset, single_slice=self.single_slice, filter_camera=filter_camera, ) folder_trees.append(folder_tree) # Create camera FolderTree (anydata/utils/folder_tree.py) self.rgb_tree[cam] = FolderTree.from_list(folder_trees) # Option to break sequences into smaller sequences if break_sequences is not None: tree = self.rgb_tree[cam].folder_tree n = len(tree) for j in range(n): if len(tree[j]) > break_sequences: idxs = list(range(0,len(tree[j]) + 1, break_sequences)) for k in range(1, len(idxs) - 1): tree.append(tree[j][idxs[k]:idxs[k+1]]) tree[j] = tree[j][idxs[0]:idxs[1]] self.rgb_tree[cam].folder_tree = tree self.rgb_tree[cam].prepare() # Validate / filter episodes based on cameras_core_mode and cameras_aux_mode self._validate_camera_episodes(cameras) # Check if all cameras have the same number of sequences and same sequence lengths for cam in cameras: assert len(self.rgb_tree[cameras[0]].folder_tree) == len(self.rgb_tree[cam].folder_tree) assert len(self.rgb_tree[cameras[0]]) == len(self.rgb_tree[cam]) # Subsample sequences if subsample is not None and subsample > 0: length = len(self) - 1 self.subsample = [int(n) for n in np.linspace(0, length, subsample)] # NOTE(bvh): ^ removed set() to allow for duplication of samples for small datasets. # this means subsample can now also mean supersample! # (a bit similar to the previous vidar trimsample and tile_minor functionality) # Add language labels if self.has_label('language') and language_dict is not None: self.language = None for path in [f'{language_dict}', f'externals/vidar/{language_dict}']: if os.path.exists(path): self.language = read_config(path).language_dict.dict break assert self.language is not None, 'Couldnt find language dictionary!' ############################################### def _validate_camera_episodes(self, cameras): """Validate or filter episodes based on cameras_core_mode / cameras_aux_mode. For each episode, reads metadata.json to check which cameras are available. - strict: raises ValueError if any episode is missing a requested camera. - flexible: builds a per-episode available cameras map so that at load time only the available cameras are used (missing cameras are skipped, not entire episodes). - skip: like flexible, but raises InsufficientCamerasError at load time if available cameras fall below the required count. """ from rich import print as rprint # Use first camera's tree to enumerate episodes ref_cam = cameras[0] ref_tree = self.rgb_tree[ref_cam] num_episodes = len(ref_tree.folder_tree) # Per-episode available cameras (episode_idx -> set of camera names) # None means "all requested cameras are available" (the common case) self._episode_cameras = None if num_episodes == 0: return requested_core = set(self.cameras_core) requested_aux = set(self.cameras_aux) episode_cameras = {} missing_core_summary = {} missing_aux_summary = {} any_missing = False for ep_idx in range(num_episodes): # File path looks like: {path}/{episode}/rgb/{cam}/{frame} first_file = ref_tree.folder_tree[ep_idx][0] episode_dir = os.path.dirname(os.path.dirname(os.path.dirname(first_file))) metadata_path = os.path.join(episode_dir, 'metadata.json') if not os.path.exists(metadata_path): continue metadata = json.load(open(metadata_path, 'r')) if 'info' in metadata and 'cameras' in metadata['info']: available = set(metadata['info']['cameras']) elif 'cameras' in metadata: available = set(metadata['cameras']) elif 'cameras_available' in metadata: available = set(metadata['cameras_available']) else: continue # Check core cameras: strict crashes only if fewer available than required sample count missing_core = requested_core - available if missing_core: avail_core = len(requested_core) - len(missing_core) core_required = self.cameras_core_sample if self.cameras_core_sample is not None else len(requested_core) if avail_core < core_required and self.cameras_core_mode == 'strict': raise ValueError( f"cameras_core_mode='strict': episode '{episode_dir}' " f"has only {avail_core}/{core_required} core cameras. " f"Missing: {sorted(missing_core)}, available: {sorted(available)}." ) any_missing = True for cam in missing_core: missing_core_summary[cam] = missing_core_summary.get(cam, 0) + 1 # Check aux cameras: same logic missing_aux = requested_aux - available if missing_aux: avail_aux = len(requested_aux) - len(missing_aux) aux_required = self.cameras_aux_sample if self.cameras_aux_sample is not None else len(requested_aux) if avail_aux < aux_required and self.cameras_aux_mode == 'strict': raise ValueError( f"cameras_aux_mode='strict': episode '{episode_dir}' " f"has only {avail_aux}/{aux_required} aux cameras. " f"Missing: {sorted(missing_aux)}, available: {sorted(available)}." ) any_missing = True for cam in missing_aux: missing_aux_summary[cam] = missing_aux_summary.get(cam, 0) + 1 # Store available cameras for this episode (only if something is missing) if missing_core or missing_aux: episode_cameras[ep_idx] = available # Store per-episode camera map so get_sample_cameras can filter at load time if any_missing: self._episode_cameras = episode_cameras n_affected = len(episode_cameras) rprint(f"[yellow]{n_affected}/{num_episodes} episodes " f"have missing cameras (will sample from available ones)[/yellow]") if missing_core_summary: rprint(f"[yellow] missing core cameras: {dict(missing_core_summary)}[/yellow]") if missing_aux_summary: rprint(f"[yellow] missing aux cameras: {dict(missing_aux_summary)}[/yellow]") def _episode_dir_from_idx(self, ep_idx, ref_cam): """Get the episode directory path from an episode index.""" first_file = self.rgb_tree[ref_cam].folder_tree[ep_idx][0] return os.path.dirname(os.path.dirname(os.path.dirname(first_file))) def get_sample_cameras(self, idx, force_camera=None): """Sample cameras, filtering out unavailable ones in flexible/skip mode.""" cameras = super().get_sample_cameras(idx, force_camera) # If no per-episode camera map, all cameras are available if self._episode_cameras is None: return cameras # Get the episode index from the flat dataset idx ref_cam = self.cameras_core[0] actual_idx = idx if self.subsample is None else self.subsample[idx] ep_idx, _ = self.rgb_tree[ref_cam].get_idxs(actual_idx) # If this episode has no missing cameras, return as-is if ep_idx not in self._episode_cameras: return cameras available = self._episode_cameras[ep_idx] cameras = [cam for cam in cameras if cam in available] # In "skip" mode, raise if filtered cameras don't meet the required count if self.cameras_core_mode == 'skip' or self.cameras_aux_mode == 'skip': episode_dir = self._episode_dir_from_idx(ep_idx, ref_cam) if self.cameras_core_mode == 'skip': actual_core = len([c for c in cameras if c in set(self.cameras_core)]) core_required = self.cameras_core_sample if self.cameras_core_sample is not None else len(self.cameras_core) if actual_core < core_required: raise InsufficientCamerasError( f"only {actual_core}/{core_required} core cameras in {episode_dir}" ) if self.cameras_aux_mode == 'skip': actual_aux = len([c for c in cameras if c in set(self.cameras_aux)]) aux_required = self.cameras_aux_sample if self.cameras_aux_sample is not None else len(self.cameras_aux) if actual_aux < aux_required: raise InsufficientCamerasError( f"only {actual_aux}/{aux_required} aux cameras in {episode_dir}" ) return cameras def get_rgb_tree_path(self, path, split, cam, folders_start, folders_end): """Prepare sequences for FolderTree""" if split is None: ### Return path directly if there is no split return path else: ### Download split if not available split = f'{path}/{split}.json' split_available = os.path.exists(split) if not split_available: src, dst = split.replace(get_local_root(), S3_PATH), path # Download split download(Namespace( dataset=os.path.basename(path), src=src, dst=dst, num_procs=16, s3_bucket='s3://tri-ml-sandbox-16011-us-west-2-datasets', # TODO @vitor remove hardcoded path data_folder='cv_unified', local_bucket=get_local_root(), subset=self.subset, labels=None, )) # Get base metadata if available metadata = f'{self.path[0]}/metadata_shared.json' self.base_metadata = read_json(metadata) if os.path.exists(metadata) else None elif not os.path.exists(split): raise ValueError(f'Split {split} does not exist') ### Load split sequences: Dict[str, dict | int | list] = json.load(open(split, 'r'))['sequences'] ### Prepare sequences for key, _ in list(sequences.items()): if is_dict(sequences[key]): if str(cam) in sequences[key]: sequences[key] = sequences[key][str(cam)] elif self.cameras_core_mode in ('flexible', 'skip') or self.cameras_aux_mode in ('flexible', 'skip'): # Camera not in this episode's per-camera dict; use a fallback # frame count from another camera so the FolderTree stays aligned fallback = next(iter(sequences[key].values())) if isinstance(fallback, int): sequences[key] = fallback elif isinstance(fallback, list): sequences[key] = len(fallback) else: sequences[key] = 0 else: sequences[key] = sequences[key][str(cam)] # Let it raise KeyError in strict if isinstance(sequences[key], int): frames = [frame_name(i) for i in range(sequences[key])] elif isinstance(sequences[key], list): if isinstance(sequences[key][0], int): frames = [frame_name(i) for i in sequences[key]] elif isinstance(sequences[key][0], str): frames = sequences[key] else: raise ValueError(f'Invalid split frame format, {type(sequences[key][0])}') else: raise ValueError(f'Invalid split frame format, {type(sequences[key])}') sequences[key] = [f'{path}/{key}/{{label}}/{cam}/{frame}' for frame in frames] ### Filter sequences based on starting string if folders_start is not None: new_sequences = {} for key, val in sequences.items(): for start in make_list(folders_start): if key.startswith(str(start)): new_sequences[key] = val sequences = new_sequences ### Filter sequences based on ending string if folders_end is not None: new_sequences = {} for key, val in sequences.items(): for start in make_list(folders_end): if key.endswith(str(start)): new_sequences[key] = val sequences = new_sequences return sequences def load_lowdim(self, filename): """Load lowdim dict (only if hasn't been loaded yet)""" if self.lowdim is None: ext = self.metadata['conventions']['lowdim']['extension'] filename = filename.format(label=self.LOWDIM_FOLDER) self.lowdim = read_npz(f'{filename}.{ext}') ############################################### def get_rgb(self, filename): """Return RGB label from file""" filename = filename.format(label=self.RGB_FOLDER) ext = self.metadata['conventions']['rgb']['extension'] image = read_image(f'{filename}.{ext}') return image def get_timestep(self, filename): """Return timestep label from lowdim""" self.load_lowdim(filename) return np.array([self.lowdim['timestep']]) def get_intrinsics(self, filename): """Return intrinsics label from lowdim""" self.load_lowdim(filename) return self.lowdim['intrinsics'] def get_lowdim(self, filename): """Return full lowdim dict from file""" self.load_lowdim(filename) lowdim = copy.deepcopy(self.lowdim) if 'language' not in lowdim: language = self.get_language(filename) if language: lowdim['language'] = np.array(language, dtype=object) return lowdim def get_extrinsics(self, filename): """Return extrinsics label from lowdim""" self.load_lowdim(filename) return self.lowdim['extrinsics'] def get_depth(self, filename): """Return depth label from file""" filename = filename.format(label=self.DEPTH_FOLDER) ext = self.metadata['conventions']['depth'].get('extension', 'npz') if ext == 'npz': depth = read_depth(f'{filename}.npz', key='depth') elif ext == 'png8': depth = read_png8(f'{filename}.png') elif ext == 'png32': depth = read_png32(f'{filename}.png') else: raise ValueError(f'Invalid depth extension {ext}') if self.apply_mask_depth and 'mask_depth' in self.labels: depth *= self.get_mask_depth(filename) if len(depth.shape) == 2: depth = np.expand_dims(depth, 0) return depth def get_optflow(self, filename, mode): """Return optical flow label from file""" label = f"{self.OPTFLOW_FOLDER}_{mode}" filename = filename.format(label=label) + '.npz' optflow = read_depth(filename, key=label) optflow = np.transpose(optflow, (1, 2, 0)) if optflow.mean() != -1: # if it's valid optflow[..., 0] *= optflow.shape[0] optflow[..., 1] *= optflow.shape[1] return optflow else: # if it's invalid return None def get_scene_flow(self, filename, mode): """Return scene flow label from file""" label = f"{self.SCNFLOW_FOLDER}_{mode}" filename = filename.format(label=label) + '.npz' if os.path.exists(filename): scene_flow = np.load(filename)['data'] return np.transpose(scene_flow, (1, 2, 0)) else: return None def get_action(self, filename): """Return action label from lowdim""" self.load_lowdim(filename) action = dict() if 'action' in self.lowdim: val = self.lowdim['action'] val = val.item() if isinstance(val, np.ndarray) and val.dtype == np.object_ else val if isinstance(val, dict): action.update(val) else: action['action'] = val # TODO: proprio is probably deprecated now? if 'proprio' in self.lowdim: val = self.lowdim['proprio'] action['proprio'] = val.item() if isinstance(val, np.ndarray) and val.dtype == np.object_ else val return action def get_ego_fields(self, filename): ''' Return all lowdim keys starting with 'ego' as a {key: value} dict ''' self.load_lowdim(filename) result = {} for key in self.lowdim: if not key.startswith('ego'): continue val = self.lowdim[key] if isinstance(val, np.ndarray) and val.dtype == np.object_: val = val.item() result[key] = val return result def get_language(self, filename): """Return language label from metadata or lowdim""" self.load_lowdim(filename) return self.prepare_language(lowdim=self.lowdim) def get_mask_rgb(self, filename): """Return mask rgb label from file""" filename = filename.format(label=self.MASK_RGB_FOLDER) + '.png' return read_image(filename, '1') def get_mask_depth(self, filename): """Return mask depth label from file""" filename = filename.format(label=self.MASK_DEPTH_FOLDER) + '.png' return read_image(filename, '1') def get_bbox2d(self, filename): """Return 2d bounding boxes from lowdim""" self.load_lowdim(filename) return self.lowdim['bbox2d'] def get_bbox3d(self, filename): """Return 3d bounding boxes from lowdim""" self.load_lowdim(filename) return self.lowdim['bbox3d'] def get_semantic(self, filename): """Return semantic label from file""" filename_semantic = filename.format(label='semantic') + '.npz' return read_depth(filename_semantic, key='semantic')