Method __init__

ram/data/three_task_dataset.py:12–38

(self, opt)

@DATASET_REGISTRY.register()
class threeTaskDataset(data.Dataset):
    def __init__(self, opt):
        super(threeTaskDataset, self).__init__()
        self.opt = opt

        self.ots_dataset = DehazeOTSBETADataset(opt, lq_path=opt['ots_lq_path'], gt_path=opt['ots_gt_path'],enlarge_ratio=opt['haze_enlarge_ratio'])
        self.rain100L_dataset = Rain100LTrainDataset(opt, dataroot=opt['rain_path'],enlarge_ratio=opt['rain_enlarge_ratio'])

        augmentators1 = parse_degradations(opt['augment1'])
        noise_dataset1 = [LowCostNoiseDataset(opt, dataroot=opt['bsd_path'], augmentator=augmentator,enlarge_ratio=opt['noise_enlarge_ratio']) for augmentator in augmentators1]
        noise_dataset2 = [LowCostNoiseDataset(opt, dataroot=opt['wed_path'], augmentator=augmentator,enlarge_ratio=opt['noise_enlarge_ratio']) for augmentator in augmentators1]
        augmentators2 = parse_degradations(opt['augment2'])
        noise_dataset3 = [LowCostNoiseDataset(opt, dataroot=opt['bsd_path'], augmentator=augmentator,enlarge_ratio=opt['noise_enlarge_ratio']) for augmentator in augmentators2]
        noise_dataset4 = [LowCostNoiseDataset(opt, dataroot=opt['wed_path'], augmentator=augmentator,enlarge_ratio=opt['noise_enlarge_ratio']) for augmentator in augmentators2]
        augmentators3 = parse_degradations(opt['augment3'])
        noise_dataset5 = [LowCostNoiseDataset(opt, dataroot=opt['bsd_path'], augmentator=augmentator,enlarge_ratio=opt['noise_enlarge_ratio']) for augmentator in augmentators3]
        noise_dataset6 = [LowCostNoiseDataset(opt, dataroot=opt['wed_path'], augmentator=augmentator,enlarge_ratio=opt['noise_enlarge_ratio']) for augmentator in augmentators3]

        self.datasets = [self.ots_dataset, self.rain100L_dataset] + \
                        noise_dataset1 + noise_dataset2 + noise_dataset3 + noise_dataset4 + noise_dataset5 + noise_dataset6

        self.sample_ids = []
        for i, dataset in enumerate(self.datasets):
            length = len(dataset)
            for idx in range(length):
                self.sample_ids.append({'dataset_idx': i, 'sample_idx': idx})

        random.shuffle(self.sample_ids)

    def __getitem__(self, index):
        sample_info = self.sample_ids[index]