|
""" 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__() |
|
|
|
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]) |
|
|
