Spaces:
Running
Running
| """ Compose multiple datasets in a single loader. """ | |
| import numpy as np | |
| from copy import deepcopy | |
| from torch.utils.data import Dataset | |
| from .wireframe_dataset import WireframeDataset | |
| from .holicity_dataset import HolicityDataset | |
| class MergeDataset(Dataset): | |
| def __init__(self, mode, config=None): | |
| super(MergeDataset, self).__init__() | |
| # Initialize the datasets | |
| self._datasets = [] | |
| spec_config = deepcopy(config) | |
| for i, d in enumerate(config['datasets']): | |
| spec_config['dataset_name'] = d | |
| spec_config['gt_source_train'] = config['gt_source_train'][i] | |
| spec_config['gt_source_test'] = config['gt_source_test'][i] | |
| if d == "wireframe": | |
| self._datasets.append(WireframeDataset(mode, spec_config)) | |
| elif d == "holicity": | |
| spec_config['train_split'] = config['train_splits'][i] | |
| self._datasets.append(HolicityDataset(mode, spec_config)) | |
| else: | |
| raise ValueError("Unknown dataset: " + d) | |
| self._weights = config['weights'] | |
| def __getitem__(self, item): | |
| dataset = self._datasets[np.random.choice( | |
| range(len(self._datasets)), p=self._weights)] | |
| return dataset[np.random.randint(len(dataset))] | |
| def __len__(self): | |
| return np.sum([len(d) for d in self._datasets]) | |