import random import time import traceback import numpy as np from rich import print import torch from torch.utils.data import ConcatDataset, Dataset from vidar.utils.config import read_config, Config from vidar.utils.setup import setup_dataset from custom.utils.constants import VAE_RATIO_THW from imaginaire.utils import log # Figure out correct FPS for each dataset, leaving most as 10 for now _DEFAULT_FPS = { 'argoverse2sync': 10.0, 'assemblyhands': 30.0, 'assemblyhands_v2': 30.0, 'ddad': 10.0, 'dl3dv': 15.0, # TODO is this correct? paper says: 30 or 60? 'droidcalib': 10.0, # TODO is this correct? 'droidcalibv2': 10.0, # TODO is this correct? 'dycheck': 30.0, # seems to match webpage 'dycheck_chunks41': 30.0, # seems to match webpage 'dycheck_mosca_chunks41': 30.0, # seems to match webpage 'egoexo4d': 30.0, # TODO is this correct? 'gcdkubric4d': 12.0, # NOTE this is actually variable 'gcdkubric4deval': 12.0, # NOTE this is actually variable 'gcdpd4d': 10.0, # NOTE this is actually variable 'gcdpd4deval': 10.0, # NOTE this is actually variable 'humanoid': 30.0, # TODO is this correct? 'iodacore': 5.0, # TODO is this correct? seems very fast 'kubric4d': 24.0, 'kubric5d': 24.0, 'lbmv12': 10.0, 'lyftl5': 10.0, 'mvimgnet': 10.0, # TODO is this correct? 'omnimmt': 10.0, 'pd4d': 10.0, 'pdnvs25': 10.0, 'pdnvs89': 10.0, 'pdv1': 10.0, 'pdv2': 10.0, 'realestate10k': 30.0, # TODO is this correct? 'robocasa': 24.0, 'scannetvideo': 15.0, # TODO is this correct? ScanNet++ paper says 60 but seems high 'tartanair': 20.0, # TODO is this correct? 'waymo': 10.0, 'wildrgbd': 24.0, # TODO is this correct? '_fallback': 15.0, } # NOTE(bvh): This is used as a general placeholder whenever language is not available. Note that in # LBM or other robotics data, the prompts are IMPERATIVE rather than DESCRIPTIVE, but I feel like # mixing them might be fine since we are aiming for a flexible / general-purpose representation anyway. _DEFAULT_PROMPT = { 'argoverse2sync': 'drive the car around', 'assemblyhands': 'let the person perform a task', 'assemblyhands_v2': 'let the person perform a task', 'ddad': 'drive the car around', 'dl3dv': 'move the camera around', 'droidcalib': 'perform the requested task', 'droidcalibv2': 'perform the requested task', 'dycheck': 'let the person do something', 'dycheck_chunks41': 'let the person do something', 'dycheck_mosca_chunks41': 'let the person do something', 'egoexo4d': 'let the people perform a task', 'gcdkubric4d': 'let the objects interact', 'gcdkubric4deval': 'let the objects interact', 'gcdpd4d': 'drive the car around', 'gcdpd4deval': 'drive the car around', 'humanoid': 'perform the requested task', 'iodacore': 'drive the car around', 'kubric4d': 'let the objects interact', 'kubric5d': 'let the objects interact', 'lbmv12': 'perform the requested task', 'lyftl5': 'drive the car around', 'mvimgnet': 'move the camera around', 'omnimmt': 'drive the robot around', 'pd4d': 'drive the car around', 'pdnvs25': 'drive the car around', 'pdnvs89': 'drive the car around', 'pdv1': 'drive the car around', 'pdv2': 'drive the car around', 'realestate10k': 'move the camera around', 'robocasa': 'perform the requested task', 'scannetvideo': 'move the camera around', 'tartanair': 'fly the drone around', 'waymo': 'drive the car around', 'wildrgbd': 'move the camera around', '_fallback': 'do something', } # OLD: mostly based on size / max camera translation values # _SF_DRIVING = 1.0 / 16.0 # because scenes are physically large # _SF_KUBRIC = 1.0 / 16.0 # because synthetic objects are oversized # _SF_ROBOTICS = 1.0 / 4.0 # _SF_4D_LARGE = 1.0 / 16.0 # _SF_4D_SMALL = 1.0 / 4.0 # _SF_3D_SCENES = 1.0 / 16.0 # _SF_3D_OBJECTS = 1.0 / 16.0 # mostly accidental # NEW: more important to consider metric scale consistency and zero-shot domain generalization _SF_DRIVING = 1.0 / 16.0 # because scenes are physically large _SF_KUBRIC = 1.0 / 32.0 # because synthetic objects are oversized _SF_ROBOTICS = 1.0 / 8.0 # ideally ~4 to be more comfortable, but we need metric scale consistency _SF_4D_LARGE = 1.0 / 8.0 # ideally ~16 ^^ _SF_4D_SMALL = 1.0 / 8.0 # ideally ~4 ^^ _SF_3D_SCENES = 1.0 / 8.0 # ideally ~16 ^^ _SF_3D_OBJECTS = 1.0 / 8.0 # ideally ~4 ^^ _SF_MVIMGNET = 1.0 / 32.0 # 3D objects, but unfortunately non-metric scale # NOTE(bvh): This affects cams (plucker cross & orig), depth, and points (xyz). _DEFAULT_SCALE_FACTOR = { 'argoverse2sync': _SF_DRIVING, 'assemblyhands': _SF_4D_SMALL, 'assemblyhands_v2': _SF_4D_SMALL, 'ddad': _SF_DRIVING, 'dl3dv': _SF_3D_SCENES, 'droidcalib': _SF_ROBOTICS, 'droidcalibv2': _SF_ROBOTICS, 'dycheck': _SF_4D_SMALL, # zero-shot 'dycheck_chunks41': _SF_4D_SMALL, # zero-shot 'dycheck_mosca_chunks41': _SF_4D_SMALL, # zero-shot 'egoexo4d': _SF_4D_LARGE, 'gcdkubric4d': _SF_KUBRIC, 'gcdkubric4deval': _SF_KUBRIC, 'gcdpd4d': _SF_DRIVING, 'gcdpd4deval': _SF_DRIVING, 'humanoid': _SF_ROBOTICS, 'iodacore': _SF_DRIVING, 'kubric4d': _SF_KUBRIC, 'kubric5d': _SF_KUBRIC, 'lbmv12': _SF_ROBOTICS, 'lyftl5': _SF_DRIVING, 'mvimgnet': _SF_MVIMGNET, 'omnimmt': _SF_ROBOTICS, # NOTE: mobile robotics is a bit of an exception! 'pd4d': _SF_DRIVING, 'pdnvs25': _SF_DRIVING, 'pdnvs89': _SF_DRIVING, 'pdv1': _SF_DRIVING, 'pdv2': _SF_DRIVING, 'realestate10k': _SF_3D_SCENES, 'robocasa': _SF_ROBOTICS, 'scannetvideo': _SF_3D_SCENES, 'tartanair': _SF_3D_SCENES, 'waymo': _SF_DRIVING, 'wildrgbd': _SF_3D_OBJECTS, '_fallback': _SF_4D_SMALL, } _ZO_DRIVING = 'egot' _ZO_ROBOTICS = False _ZO_3D = True _ZO_4D = False # Figure out correct zero origin for each dataset # there are roughly 3-4 groups (robotics, driving, 3D/4D/everyday) _DEFAULT_ZERO_ORIGIN = { 'argoverse2sync': _ZO_DRIVING, 'assemblyhands': _ZO_ROBOTICS, 'assemblyhands_v2': _ZO_ROBOTICS, 'ddad': _ZO_DRIVING, 'dl3dv': _ZO_3D, 'droidcalib': _ZO_ROBOTICS, 'droidcalibv2': _ZO_ROBOTICS, 'dycheck': _ZO_4D, # zero-shot 'dycheck_chunks41': _ZO_4D, # zero-shot 'dycheck_mosca_chunks41': _ZO_4D, # zero-shot 'egoexo4d': _ZO_4D, 'gcdkubric4d': _ZO_4D, 'gcdkubric4deval': _ZO_4D, 'gcdpd4d': _ZO_DRIVING, 'gcdpd4deval': _ZO_DRIVING, 'humanoid': _ZO_ROBOTICS, 'iodacore': _ZO_DRIVING, 'kubric4d': _ZO_4D, 'kubric5d': _ZO_4D, 'lbmv12': _ZO_ROBOTICS, 'lyftl5': _ZO_DRIVING, 'mvimgnet': _ZO_3D, 'omnimmt': _ZO_DRIVING, # NOTE: mobile robotics is a bit of an exception! 'pd4d': _ZO_DRIVING, 'pdnvs25': _ZO_DRIVING, 'pdnvs89': _ZO_DRIVING, 'pdv1': _ZO_DRIVING, 'pdv2': _ZO_DRIVING, 'realestate10k': _ZO_3D, 'robocasa': _ZO_ROBOTICS, 'scannetvideo': _ZO_3D, 'tartanair': _ZO_3D, 'waymo': _ZO_DRIVING, 'wildrgbd': _ZO_3D, '_fallback': _ZO_4D, } # Figure out correct stride range for each dataset # NOTE(bvh): Prefer to include higher strides when the typical FPS values are higher. # TODO / FIXME @vitor this crashes if > 1 when metrics/evaluator is active _DEFAULT_STRIDE = { 'assemblyhands': [1,2,3], 'assemblyhands_v2': [1,2,3], 'dl3dv': [1,2], 'egoexo4d': [1,2,3], 'humanoid': [1,2,3], 'kubric4d': [1,2,3,4], 'kubric5d': [1,2,3,4], 'pd4d': [1,2], 'realestate10k': [1,2,3], 'robocasa': [1,2,3], 'scannetvideo': [1,2], 'tartanair': [1,2], 'wildrgbd': [1,2,3], '_fallback': 1, } def get_episode_index(vidar_tag): pass # TODO class VidarDataset(Dataset): def __init__( self, dataset_dir, num_frames=93, phase='train', vae_ratio=VAE_RATIO_THW, config_key='dataset', params=None, # TODO @vitor please document data_defaults=None, # accepted for API compat with AnyDataset, not used data_overrides=None, # used for zero_origin / frame_stride overrides (see _vidar_get_*) shuffle=None, # accepted for API compat with AnyDataset; VidarDataset uses DistributedSampler instead ): super().__init__() self.dataset_dir = dataset_dir self.num_frames = num_frames self.phase = phase self.vae_ratio = vae_ratio self.data_overrides = data_overrides or {} if params is None: # TODO @vitor this seems arbitrary / unclean, please rewrite self.pad_dims = (384, 384) else: self.pad_dims = (576, 576) # Instantiate vidar dataset cfg = read_config(self.dataset_dir) cfg.dataset.phase = [phase] if params is not None: if 'resize' in params: if cfg.dataset.has('augmentation'): cfg.dataset.augmentation.resize = params['resize'] else: cfg.dataset.augmentation = Config(resize=params['resize']) (vidar_datasets, vidar_configs) = setup_dataset(cfg.dict[config_key], verbose=True) assert len(vidar_datasets) >= 1, f'Expected at least one vidar_dataset, got {len(vidar_datasets)}' assert len(vidar_configs) >= 1, f'Expected at least one vidar_config, got {len(vidar_configs)}' self.vidar_dataset = ConcatDataset(vidar_datasets) self.vidar_configs = vidar_configs self.error_count = 0 def __str__(self): return f'{len(self.vidar_dataset)} samples from {self.dataset_dir}' def __len__(self): return len(self.vidar_dataset) def _vidar_get_fps(self, sample, index, dset_name): ''' :return fps (float). ''' dset_name = dset_name.lower() if 'fps' in sample.keys(): fps = sample['fps'] elif dset_name in _DEFAULT_FPS.keys(): fps = _DEFAULT_FPS[dset_name] else: fps = _DEFAULT_FPS['_fallback'] if index <= 5: log.warning(f'Unknown dataset {dset_name}, using fallback FPS {fps}', rank0_only=False) return fps def _vidar_get_prompt(self, sample, index, dset_name): ''' :return prompt (str). ''' dset_name = dset_name.lower() if 'language' in sample.keys(): prompt = sample['language'][(0, 0)]['prompt'] # NOTE(bvh): temporary fix for newer webdatasets that broke existing code # prompt = [x.split('\'')[1] for x in str(prompt[0]).split('\n')] # assuming there are no apostrophes in the dataset prompt = [x.replace('[','').replace(']','').strip()[1:-1] for x in str(prompt[0]).split('\n')] # now we have a clean list of str again elif 'lat_language' in sample.keys(): prompt = sample['lat_language'][(0, 0)]['prompt'] elif dset_name in _DEFAULT_PROMPT.keys(): prompt = _DEFAULT_PROMPT[dset_name] else: prompt = _DEFAULT_PROMPT['_fallback'] if index <= 5: log.warning(f'Unknown dataset {dset_name}, using fallback prompt {prompt}', rank0_only=False) if isinstance(prompt, list): # Multiple prompts can be provided, so we choose one at random. if index <= 5: log.warning(f'Multiple prompts ({len(prompt)}) for {dset_name}, picking random one', rank0_only=False) prompt = random.choice(prompt) return prompt def _vidar_get_scale_factor(self, sample, index, dset_name): ''' :return scale_factor (float). ''' # NOTE(bvh): Both sample-dependent and overall global scaling / clipping factors are not ideal, # so I am opting for semi-global per-dataset/domain factors instead. dset_name = dset_name.lower() if 'scale_factor' in sample.keys(): scale_factor = sample['scale_factor'] elif dset_name in _DEFAULT_SCALE_FACTOR.keys(): scale_factor = _DEFAULT_SCALE_FACTOR[dset_name] else: scale_factor = _DEFAULT_SCALE_FACTOR['_fallback'] if index <= 5: log.warning(f'Unknown dataset {dset_name}, using fallback scale factor {scale_factor}', rank0_only=False) return scale_factor def _vidar_get_zero_origin(self, sample, index, dset_name): ''' :return zero_origin (str or bool). ''' dset_name = dset_name.lower() return True # DEBUG / TEMP for AnyView ECCV rebuttal # Phase-time override (data_train_overrides.zero_origin / data_val_overrides.zero_origin). # Replaces the deprecated flat `override_zero_origin` field. Set to None to skip. ov = self.data_overrides.get('zero_origin', None) if ov is not None: return ov if 'zero_origin' in sample.keys(): zero_origin = sample['zero_origin'] elif dset_name in _DEFAULT_ZERO_ORIGIN.keys(): zero_origin = _DEFAULT_ZERO_ORIGIN[dset_name] else: zero_origin = _DEFAULT_ZERO_ORIGIN['_fallback'] if index <= 5: log.warning(f'Unknown dataset {dset_name}, using fallback zero origin {zero_origin}', rank0_only=False) return zero_origin def _vidar_get_stride(self, sample, index, dset_name): ''' :return stride (int). ''' dset_name = dset_name.lower() return 1 # DEBUG / TEMP for AnyView ECCV rebuttal # Phase-time override (data_train_overrides.frame_stride / data_val_overrides.frame_stride). # Replaces the deprecated flat `override_temporal_stride_{train,val}` fields. Set to None to skip. ov = self.data_overrides.get('frame_stride', None) if ov is not None: if isinstance(ov, list): ov = random.choice(ov) return ov if 'stride' in sample.keys(): stride = sample['stride'] elif dset_name in _DEFAULT_STRIDE.keys(): stride = _DEFAULT_STRIDE[dset_name] else: stride = _DEFAULT_STRIDE['_fallback'] if index <= 5: log.warning(f'Unknown dataset {dset_name}, using fallback stride {stride}', rank0_only=False) if isinstance(stride, list): stride = random.choice(stride) return stride def reset_iter(self): """No-op for map-style vidar datasets (nothing to reset).""" pass def __getitem__(self, index): tensor_kwargs = dict(dtype=torch.bfloat16) vidar_tag = '' try: # Get vidar sample (can be from any dataset) vidar_dict = self.vidar_dataset[index] # ^ NOTE(bvh): Entries do not yet have batch dimension, # e.g. 'rgb' -> (3, H, W) tensors of float in [0, 1]. # Get short, meaningful dataset source identifier vidar_tag = vidar_dict.get('tag', 'unknown_tag') # e.g. # 'kubric4d__scn02800' TODO new tag # 'LBMv12__BimanualPutRedBellPepperInBin__riverway__sim' # 'PDnvs25__25_camera__scene_000000' dset_name = vidar_tag.split('__')[0].lower() # e.g. 'kubric4d' or 'lbmv12' # Get long, specific video identifier vidar_filename = vidar_dict.get('filename', 'unknown_filename') sample_id = vidar_filename # e.g. # 'BimanualPutRedBellPepperInBin__riverway__sim__2025-01-02T10-49-28-05-00__000000/seq-000002-16_48.tar_0' # Get nominal dimensions in both pixel and latent space # NOTE(bvh): In Any4D, different viewpoints do NOT necessarily need to have the same # resolution or duration, but we should obtain representative values for conditioning anyway. if 'rgb' in vidar_dict: # e.g. pixel (raw) webdata h, w = list(vidar_dict['rgb'].values())[0].shape[-2:] t = len([key for key in vidar_dict['rgb'].keys() if key[1] == 0]) hh, ww = h // self.vae_ratio[1], w // self.vae_ratio[2] tt = (t - 1) // self.vae_ratio[0] + 1 else: # e.g. latent (cached) webdata hh, ww = list(vidar_dict['lat_rgb'].values())[0].shape[-2:] tt = len([key for key in vidar_dict['lat_rgb'].keys() if key[1] == 0]) h, w = int(hh * self.vae_ratio[1]), int(ww * self.vae_ratio[2]) t = (tt - 1) * self.vae_ratio[0] + 1 # Get specific dataset information fps = self._vidar_get_fps(vidar_dict, index, dset_name) prompt = self._vidar_get_prompt(vidar_dict, index, dset_name) scale_factor = self._vidar_get_scale_factor(vidar_dict, index, dset_name) # NOTE(bvh): These two options can now be overridden in the config # e.g. to avoid introducing more randomness in val, # and also because evaluator per-frame averaging breaks otherwise (for now) zero_origin = self._vidar_get_zero_origin(vidar_dict, index, dset_name) frame_stride = self._vidar_get_stride(vidar_dict, index, dset_name) # NOTE(bvh): FPS correction happens later in a4d_model because of stride alignment # + potential intra-batch differences. # Add Cosmos-related parameters (necessary for vanilla video2world) image_size = torch.tensor([h, w, h, w], **tensor_kwargs) # pixel space num_frames = torch.tensor([t], **tensor_kwargs) # pixel space padding_mask = torch.zeros((1, h, w), **tensor_kwargs) # pixel space # NOTE(bvh): prompt being weird is giving collate errors in recent webdatasets # default_collate: batch must contain tensors, numpy arrays, numbers, dicts or lists; found = 20: log.error('Too many errors, raising exception...', rank0_only=False) raise e return self[np.random.randint(len(self))]