Spaces:
Sleeping
Sleeping
| import numpy as np | |
| from src.datasets.batched_sampler import DynamicBatchedMultiFeatureRandomSampler | |
| class BaseDataset: | |
| """Abstract base class for multi-resolution training datasets.""" | |
| def __init__(self, config): | |
| """ | |
| Args: | |
| config: OmegaConf config containing training and data settings. | |
| """ | |
| self.num_views = config.training.num_views | |
| self._set_resolutions(config.training.res_dict) | |
| self._set_input_views(config.training.num_views) | |
| self._set_target_views(config.training.target_views) | |
| self.seed = config.data.seed | |
| self.max_num_retries = 10 | |
| def set_epoch(self, epoch): | |
| """Hook for epoch-level state updates; subclasses may override. | |
| Args: | |
| epoch: The current training epoch index. | |
| """ | |
| pass | |
| def make_sampler( | |
| self, | |
| batch_size_per_gpu, | |
| shuffle: bool = True, | |
| world_size: int = 1, | |
| rank: int = 0, | |
| drop_last: bool = True | |
| ): | |
| """Construct a DynamicBatchedMultiFeatureRandomSampler for this dataset. | |
| Args: | |
| batch_size_per_gpu: List of per-GPU batch sizes, one per view-count bucket. | |
| shuffle: Must be True; non-shuffled iteration is not implemented. | |
| world_size: Total number of distributed processes. | |
| rank: This process's distributed rank. | |
| drop_last: Whether to drop the final incomplete batch. | |
| Returns: | |
| A DynamicBatchedMultiFeatureRandomSampler instance. | |
| """ | |
| if not shuffle: | |
| raise NotImplementedError("Only shuffle=True is supported for now.") | |
| num_of_num_views = len(self.num_views) | |
| feature_to_batch_size_map = {i: bs for i, bs in enumerate(batch_size_per_gpu)} | |
| return DynamicBatchedMultiFeatureRandomSampler( | |
| dataset=self, | |
| pool_sizes=[num_of_num_views], | |
| scaling_feature_idx=0, | |
| feature_to_batch_size_map=feature_to_batch_size_map, | |
| world_size=world_size, | |
| rank=rank, | |
| drop_last=drop_last, | |
| ) | |
| def _scene_len(self): | |
| """Initialize num_of_scenes from data_path; subclasses must set data_path first.""" | |
| self.data_path = [] # set in subclass | |
| self.num_of_scenes = len(self.data_path) | |
| def __len__(self): | |
| """Return the total number of scenes in the dataset.""" | |
| return self.num_of_scenes | |
| def _get_views(self, idx, resolution, num_views_to_input, num_views_to_target): | |
| """Load and return a scene sample. Must be implemented by subclasses. | |
| Args: | |
| idx: Scene index into data_path. | |
| resolution: (height, width) tuple for image loading. | |
| num_views_to_input: Number of input views to select. | |
| num_views_to_target: Number of target views to select. | |
| Returns: | |
| A dict containing image tensors, intrinsics, and pose tensors. | |
| """ | |
| raise NotImplementedError() | |
| def _set_resolutions(self, resolutions): | |
| """Build an index-keyed dict of (height, width) resolution tuples. | |
| Args: | |
| resolutions: List of [height, width] pairs from config. | |
| """ | |
| self._resolutions = {i: tuple(res) for i, res in enumerate(resolutions)} | |
| def _set_input_views(self, num_views): | |
| """Build an index-keyed dict of input view counts. | |
| Args: | |
| num_views: List of input view counts from config. | |
| """ | |
| self._input_views = dict(enumerate(num_views)) | |
| def _set_target_views(self, target_views): | |
| """Build an index-keyed dict of target view counts. | |
| Args: | |
| target_views: List of target view counts from config. | |
| """ | |
| self._target_views = dict(enumerate(target_views)) | |
| def _getitem_fn(self, idx): | |
| """Dispatch a sampler-produced index tuple to _get_views. | |
| Args: | |
| idx: A tuple of (sample_idx, dict_idx) where dict_idx selects | |
| the resolution/view-count bucket. | |
| Returns: | |
| The result of _get_views for the resolved parameters. | |
| """ | |
| idx, dict_idx = idx | |
| resolution = self._resolutions[dict_idx] | |
| input_view = self._input_views[dict_idx] | |
| target_view = self._target_views[dict_idx] | |
| return self._get_views(idx, resolution, input_view, target_view) | |
| def __getitem__(self, idx): | |
| """Return a dataset sample, retrying with a random index on failure. | |
| Args: | |
| idx: A (sample_idx, dict_idx) tuple as yielded by the sampler. | |
| Returns: | |
| A valid sample dict from _getitem_fn. | |
| Raises: | |
| RuntimeError: If all retries are exhausted. | |
| """ | |
| for _ in range(self.max_num_retries + 1): | |
| try: | |
| return self._getitem_fn(idx) | |
| except Exception: | |
| if isinstance(idx, tuple): | |
| idx_list = list(idx) | |
| idx_list[0] = np.random.randint(0, len(self)) | |
| idx = tuple(idx_list) | |
| else: | |
| idx = np.random.randint(0, len(self)) | |
| raise RuntimeError(f"Failed to load a valid sample after {self.max_num_retries} retries.") | |