Spaces:
Running on Zero
Running on Zero
File size: 8,447 Bytes
0122a25 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 | """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]
|