| from dataset import DEFAULTDataset | |
| from torch.utils.data import WeightedRandomSampler | |
| def get_dataset(cfg): | |
| if cfg.dataset.name == 'DEFAULT': | |
| train_dataset = DEFAULTDataset( | |
| root_dir=cfg.dataset.root_dir, internal_resolution=cfg.model.internal_resolution) | |
| val_dataset = DEFAULTDataset( | |
| root_dir=cfg.dataset.root_dir, internal_resolution=cfg.model.internal_resolution) | |
| sampler = None | |
| return train_dataset, val_dataset, sampler | |
| raise ValueError(f'{cfg.dataset.name} Dataset is not available') | |