Spaces:
Running
Running
File size: 1,834 Bytes
ba23d94 | 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 | # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
from typing import Any, Callable, List, Tuple, Union
from sapiens.registry import DATASETS
from .base_dataset import BaseDataset
@DATASETS.register_module()
class CombinedDataset(BaseDataset):
def __init__(
self, datasets: list, pipeline: List[Union[dict, Callable]] = [], **kwargs
):
self.datasets = []
for cfg in datasets:
dataset = DATASETS.build(cfg)
self.datasets.append(dataset)
self._lens = [len(dataset) for dataset in self.datasets]
self._len = sum(self._lens)
super(CombinedDataset, self).__init__(pipeline=pipeline, **kwargs)
assert len(self.datasets) > 0
return
def __len__(self):
return self._len
def _get_subset_index(self, index: int) -> Tuple[int, int]:
if index >= len(self) or index < -len(self):
raise ValueError(f"index {index} out of bounds for length {len(self)}.")
if index < 0:
index = index + len(self)
subset_index = 0
while index >= self._lens[subset_index]:
index -= self._lens[subset_index]
subset_index += 1
return subset_index, index
def prepare_data(self, idx: int) -> Any:
data_info = self.get_data_info(idx)
if data_info is None:
return None
for transform in self.pipeline.transforms:
data_info = transform(data_info)
return data_info
def get_data_info(self, idx: int) -> dict:
subset_idx, sample_idx = self._get_subset_index(idx)
data_info = self.datasets[subset_idx][sample_idx]
return data_info
|