Spaces:
Running on Zero
Running on Zero
| """DataPipe wraps datasets to share the prepossessing pipeline.""" | |
| from __future__ import annotations | |
| import bisect | |
| import random | |
| from collections.abc import Callable, Iterable | |
| from torch.utils.data import ConcatDataset, Dataset | |
| from .reference import MultiViewDataset | |
| from .transforms.base import TFunctor | |
| from .typing import DictData, DictDataOrList | |
| class DataPipe(ConcatDataset[DictDataOrList]): | |
| """DataPipe class. | |
| This class wraps one or multiple instances of a PyTorch Dataset so that the | |
| preprocessing steps can be shared across those datasets. Composes dataset | |
| and the preprocessing pipeline. | |
| """ | |
| def __init__( | |
| self, | |
| datasets: Dataset[DictDataOrList] | Iterable[Dataset[DictDataOrList]], | |
| preprocess_fn: Callable[ | |
| [list[DictData]], list[DictData] | |
| ] = lambda x: x, | |
| ): | |
| """Creates an instance of the class. | |
| Args: | |
| datasets (Dataset | Iterable[Dataset]): Dataset(s) to be wrapped by | |
| this data pipeline. | |
| preprocess_fn (Callable[[list[DictData]], list[DictData]]): | |
| Preprocessing function of a single sample. It takes a list of | |
| samples and returns a list of samples. Defaults to identity | |
| function. | |
| """ | |
| if isinstance(datasets, Dataset): | |
| datasets = [datasets] | |
| super().__init__(datasets) | |
| self.preprocess_fn = preprocess_fn | |
| self.has_reference = any( | |
| _check_reference(dataset) for dataset in datasets | |
| ) | |
| if self.has_reference and not all( | |
| _check_reference(dataset) for dataset in datasets | |
| ): | |
| raise ValueError( | |
| "All datasets must be MultiViewDataset / has reference if " | |
| + "one of them is." | |
| ) | |
| def __getitem__(self, idx: int) -> DictDataOrList: | |
| """Wrap getitem to apply augmentations.""" | |
| samples = super().__getitem__(idx) | |
| if isinstance(samples, list): | |
| return self.preprocess_fn(samples) | |
| return self.preprocess_fn([samples])[0] | |
| class MultiSampleDataPipe(DataPipe): | |
| """MultiSampleDataPipe class. | |
| This class wraps DataPipe to support augmentations that require multiple | |
| images (e.g., Mosaic and Mixup) by sampling additional indices for each | |
| image. NUM_SAMPLES needs to be defined as a class attribute for transforms | |
| that require multi-sample augmentation. | |
| """ | |
| def __init__( | |
| self, | |
| datasets: Dataset[DictDataOrList] | Iterable[Dataset[DictDataOrList]], | |
| preprocess_fn: list[list[TFunctor]], | |
| ): | |
| """Creates an instance of the class. | |
| Args: | |
| datasets (Dataset | Iterable[Dataset]): Dataset(s) to be wrapped by | |
| this data pipeline. | |
| preprocess_fn (list[list[TFunctor]]): Preprocessing functions of a | |
| single sample. Different than DataPipe, this is a list of lists | |
| of transformation functions. The inner list is for transforms | |
| that needs to share the same sampled indices (e.g., | |
| GenMosaicParameters and MosaicImages), and the outer list is | |
| for different transforms. | |
| """ | |
| super().__init__(datasets) | |
| self.preprocess_fns = preprocess_fn | |
| def _sample_indices(self, idx: int, num_samples: int) -> list[int]: | |
| """Sample additional indices for multi-sample augmentation.""" | |
| indices = [idx] | |
| for _ in range(1, num_samples): | |
| indices.append(random.randint(0, len(self) - 1)) | |
| return indices | |
| def __getitem__(self, idx: int) -> DictDataOrList: | |
| """Wrap getitem to apply augmentations.""" | |
| samples = super(DataPipe, self).__getitem__(idx) | |
| if not isinstance(samples, list): | |
| samples = [samples] | |
| single_view = True | |
| else: | |
| single_view = False | |
| for preprocess_fn in self.preprocess_fns: | |
| if hasattr(preprocess_fn[0], "NUM_SAMPLES"): | |
| num_samples = preprocess_fn[0].NUM_SAMPLES | |
| aug_inds = self._sample_indices(idx, num_samples) | |
| add_samples = [ | |
| super(DataPipe, self).__getitem__(ind) | |
| for ind in aug_inds[1:] | |
| ] | |
| prep_samples = [] | |
| for i, samp in enumerate(samples): | |
| prep_samples.append(samp) | |
| prep_samples += [ | |
| s[i] if isinstance(s, list) else s for s in add_samples | |
| ] | |
| else: | |
| num_samples = 1 | |
| prep_samples = samples | |
| for prep_fn in preprocess_fn: | |
| prep_samples = prep_fn.apply_to_data(prep_samples) # type: ignore # pylint: disable=line-too-long | |
| samples = prep_samples[::num_samples] | |
| return samples[0] if single_view else samples | |
| def _check_reference(dataset: Dataset[DictDataOrList]) -> bool: | |
| """Check if the datasets have reference.""" | |
| has_reference = ( | |
| dataset.has_reference if hasattr(dataset, "has_reference") else False | |
| ) | |
| return has_reference or isinstance(dataset, MultiViewDataset) | |
| class TupleDataPipe(ConcatDataset): | |
| """A custom ConcatDataset that supports indexing with a tuple. | |
| Modified from VGGT's TupleConcatDataset. | |
| Standard PyTorch ConcatDataset only accepts an integer index. This class | |
| extends that functionality to allow passing a tuple like (sample_idx, | |
| num_images, aspect_ratio), where the first element is used to determine | |
| which sample to fetch, and the full tuple is passed down to the selected | |
| dataset's __getitem__ method. | |
| It also supports an option to randomly sample across all datasets, ignoring | |
| the provided index. This is useful during training when shuffling the | |
| entire dataset might cause memory issues due to duplicating dictionaries. | |
| If doing this, you can set PyTorch's dataloader shuffle to False. | |
| """ | |
| def __init__(self, datasets, inside_random: bool = False): | |
| """Initialize the TupleConcatDataset. | |
| Args: | |
| datasets (iterable): An iterable of PyTorch Dataset objects to | |
| concatenate. | |
| common_config (dict): Common configuration dict, used to check for | |
| random sampling. | |
| """ | |
| super().__init__(datasets) | |
| # If True, ignores the input index and samples randomly across all datasets | |
| # This provides an alternative to dataloader shuffling for large datasets | |
| self.inside_random = inside_random | |
| def __getitem__(self, idx): | |
| """Retrieves an item using either an integer index or a tuple index. | |
| Args: | |
| idx (int or tuple): The index. If tuple, the first element is the | |
| sequence index across the concatenated datasets, and the rest | |
| are passed down. If int, it's treated as the sequence index. | |
| Returns: | |
| The item returned by the underlying dataset's __getitem__ method. | |
| Raises: | |
| ValueError: If the index is out of range or the tuple doesn't have | |
| exactly three elements. | |
| """ | |
| idx_tuple = None | |
| if isinstance(idx, tuple): | |
| idx_tuple = idx | |
| idx = idx_tuple[0] # Extract the sequence index | |
| # Override index with random value if inside_random is enabled | |
| if self.inside_random: | |
| total_len = self.cumulative_sizes[-1] | |
| idx = random.randint(0, total_len - 1) | |
| # Handle negative indices | |
| if idx < 0: | |
| if -idx > len(self): | |
| raise ValueError( | |
| "absolute value of index should not exceed dataset length" | |
| ) | |
| idx = len(self) + idx | |
| # Find which dataset the index belongs to | |
| dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx) | |
| if dataset_idx == 0: | |
| sample_idx = idx | |
| else: | |
| sample_idx = idx - self.cumulative_sizes[dataset_idx - 1] | |
| # Create the tuple to pass to the underlying dataset | |
| if len(idx_tuple) == 3: | |
| idx_tuple = (sample_idx,) + idx_tuple[1:] | |
| else: | |
| raise ValueError("Tuple index must have exactly three elements") | |
| # Pass the modified tuple to the appropriate dataset | |
| return self.datasets[dataset_idx][idx_tuple] | |