| from .catalog import DatasetCatalog | |
| from ldm.util import instantiate_from_config | |
| import torch | |
| class ConCatDataset(): | |
| def __init__(self, dataset_name_list, ROOT, which_embedder, train=True, repeats=None): | |
| self.datasets = [] | |
| cul_previous_dataset_length = 0 | |
| offset_map = [] | |
| which_dataset = [] | |
| if repeats is None: | |
| repeats = [1] * len(dataset_name_list) | |
| else: | |
| assert len(repeats) == len(dataset_name_list) | |
| Catalog = DatasetCatalog(ROOT, which_embedder) | |
| for dataset_idx, (dataset_name, yaml_params) in enumerate(dataset_name_list.items()): | |
| repeat = repeats[dataset_idx] | |
| dataset_dict = getattr(Catalog, dataset_name) | |
| target = dataset_dict['target'] | |
| params = dataset_dict['train_params'] if train else dataset_dict['val_params'] | |
| if yaml_params is not None: | |
| params.update(yaml_params) | |
| dataset = instantiate_from_config( dict(target=target, params=params) ) | |
| self.datasets.append(dataset) | |
| for _ in range(repeat): | |
| offset_map.append( torch.ones(len(dataset))*cul_previous_dataset_length ) | |
| which_dataset.append( torch.ones(len(dataset))*dataset_idx ) | |
| cul_previous_dataset_length += len(dataset) | |
| offset_map = torch.cat(offset_map, dim=0).long() | |
| self.total_length = cul_previous_dataset_length | |
| self.mapping = torch.arange(self.total_length) - offset_map | |
| self.which_dataset = torch.cat(which_dataset, dim=0).long() | |
| def total_images(self): | |
| count = 0 | |
| for dataset in self.datasets: | |
| print(dataset.total_images()) | |
| count += dataset.total_images() | |
| return count | |
| def __getitem__(self, idx): | |
| dataset = self.datasets[ self.which_dataset[idx] ] | |
| return dataset[ self.mapping[idx] ] | |
| def __len__(self): | |
| return self.total_length | |