Exemple #1
0
    def __getitem__(self, index):
        if self.data_type == 'lmdb' and self.GT_env is None:
            self._init_lmdb()

        key = self.paths_GT[index]
        img_GT = self.read_img(key, True)
        img_LQ = self.read_img(key, False)

        if self.opt['phase'] == 'train':
            H, W, _ = img_LQ.shape
            rnd_h = random.randint(0, max(0, H - self.LQ_size))
            rnd_w = random.randint(0, max(0, W - self.LQ_size))
            img_LQ = img_LQ[rnd_h:rnd_h + self.LQ_size,
                            rnd_w:rnd_w + self.LQ_size, :]
            rnd_h_HR, rnd_w_HR = int(rnd_h * self.scale), int(rnd_w *
                                                              self.scale)

            img_GT = img_GT[rnd_h_HR : rnd_h_HR + self.GT_size, \
                                rnd_w_HR : rnd_w_HR + self.GT_size, :]
            # augmentation - flip, rotate
            img_LQ, img_GT = util.augment([img_LQ, img_GT],
                                          self.opt['use_flip'],
                                          self.opt['use_rot'])

        # BGR ==> RGB
        img_GT, img_LQ = img_GT[:, :, (2, 1, 0)], img_LQ[:, :, (2, 1, 0)]
        # HWC ==> CHW
        img_GT = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
        img_LQ = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_LQ, (2, 0, 1)))).float()
        return {'LQs': img_LQ, 'GT': img_GT, 'key': key, 'scale': self.scale}
    def __getitem__(self, index):
        # get HR image
        # load frequence as BGR, HWDC
        img_HR = self.HR_hdf5['data'][index] if self.HR_hdf5 else None

        if 'data' in self.LR_hdf5:
            img_LR = self.LR_hdf5['data'][index]

        if self.opt['phase'] == 'train':
            # augmentation - flip, rotate
            img_LR, img_HR = util.augment(
                [img_LR, img_HR], self.opt['use_flip'], self.opt['use_rot']
            )

        # BGR to RGB,
        if self.HR_hdf5:
            if img_HR.shape[3] == 3:
                img_HR = img_HR[:, :, :, [2, 1, 0]]
            # HWDC to CHWD, numpy to tensor
            img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (3, 0, 1, 2)))).float()

        # BGR to RGB,
        if img_LR.shape[3] == 3:
            img_LR = img_LR[:, :, :, [2, 1, 0]]
        img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (3, 0, 1, 2)))).float()

        if self.HR_hdf5:
            return {'LR': img_LR, 'HR': img_HR, 'hz': self.LR_hdf5['hz'][index]}
        else:
            return {'LR': img_LR, 'hz': self.LR_hdf5['hz'][index]}
    def __getitem__(self, index):
        # get HR image
        # load frequence as BGR, HWDC
        img_HR = self.HR_hdf5['data'][index]

        # get LR image
        # load frequence as BGR, HWDC
        if 'data' in self.LRx2_hdf5:
            img_LRx2 = self.LRx2_hdf5['data'][index]
        if 'data' in self.LRx4_hdf5:
            img_LRx4 = self.LRx4_hdf5['data'][index]

        if self.opt['phase'] == 'train':
            # augmentation - flip, rotate
            img_LRx4, img_LRx2, img_HR = util.augment(
                [img_LRx4, img_LRx2, img_HR], self.opt['use_flip'], self.opt['use_rot']
            )

        # BGR to RGB,
        if img_HR.shape[3] == 3:
            img_HR = img_HR[:, :, :, [2, 1, 0]]
            img_LRx2 = img_LRx2[:, :, :, [2, 1, 0]]
            img_LRx4 = img_LRx4[:, :, :, [2, 1, 0]]
        # HWC to CHW, numpy to tensor
        img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (3, 0, 1, 2)))).float()
        img_LRx2 = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LRx2, (3, 0, 1, 2)))).float()
        img_LRx4 = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LRx4, (3, 0, 1, 2)))).float()

        return {'LRx4': img_LRx4, 'LRx2': img_LRx2, 'HR': img_HR}
    def __getitem__(self, index):
        GT_path, LQ_path = None, None
        GT_size = self.opt['GT_size']

        # get GT audio
        GT_path = self.paths_GT[index]
        audio_GT = util.read_audio(GT_path)

        # get LQ audio
        LQ_path = self.paths_LQ[index]
        resolution = None
        audio_LQ = util.read_audio(LQ_path)

        if self.opt['phase'] == 'train':

            LQ_size = GT_size

            # randomly trim
            start = random.randint(0, audio_GT.shape[1] - GT_size)
            stop = start + GT_size
            audio_GT = audio_GT[:,start:stop,:]
            audio_LQ = audio_LQ[:,start:stop,:]

            # augmentation - reverse
            audio_LQ, audio_GT = util.augment([audio_LQ, audio_GT], self.opt['use_reverse'])
        audio_GT = torch.from_numpy(np.transpose(audio_GT)).float()
        audio_LQ = torch.from_numpy(np.transpose(audio_LQ)).float()

        return {'LQ': audio_LQ, 'GT': audio_GT, 'LQ_path': LQ_path, 'GT_path': GT_path}
Exemple #5
0
    def __getitem__(self, index):
        GT_path, NOISY_path = None, None
        GT_size = self.opt['GT_size']

        GT_path = self.paths_GT[index]
        specular_ref = util_disney.loadDisneyEXR_multi_ref_shading(
            GT_path, self.opt["feature"] + ["specular"])

        NOISY_path = self.paths_NOISY[index]
        specular_in, features = util_disney.loadDisneyEXR_feature_shading(
            NOISY_path, self.opt["feature"] + ["specular"])

        if self.opt['phase'] == 'train':
            # augmentation - flip, rotate
            specular_ref, specular_in, features= util.augment([ specular_ref, specular_in, features], self.opt['use_flip'], \
                self.opt['use_rot'])

        features = torch.from_numpy(
            np.ascontiguousarray(np.transpose(features, (2, 0, 1)))).float()
        specular_in = torch.from_numpy(
            np.ascontiguousarray(np.transpose(specular_in,
                                              (2, 0, 1)))).float()
        specular_ref = torch.from_numpy(
            np.ascontiguousarray(np.transpose(specular_ref,
                                              (2, 0, 1)))).float()
        return {
            "seg": features,
            "specular_in": specular_in,
            "specular_ref": specular_ref,
            'NOISY_path': NOISY_path,
            'GT_path': GT_path,
            "x_offset": 128,
            "y_offset": 128
        }
Exemple #6
0
    def __getitem__(self, index ):
        # 从文件中读出HR图片,返回HR、LR图片对
        HR_path , LR_path = None, None
        HR_path = self.HR_paths[index]
        img_HR = util.read_img_(HR_path, self.n_channels)  # return: Numpy float32, HWC, BGR, [0,1]
        scale = self.scale
        patch_size = self.patch_size
        if self.opt.phase == 'train':
            H_s, W_s, _ = img_HR.shape
            img_HR = cv2.resize(np.copy(img_HR), (W_s, H_s), interpolation=cv2.INTER_LINEAR)
            # force to 3 channels
            if img_HR.ndim == 2:
                img_HR = cv2.cvtColor(img_HR, cv2.COLOR_GRAY2BGR)
        # using matlab imresize
        img_LR = util.imresize_np(img_HR, 1 / scale, True)
        H, W, C = img_LR.shape
        if img_LR.ndim == 2:
            img_LR = np.expand_dims(img_LR, axis=2)

        if self.opt.phase == 'train':
            patch_H, patch_L = util.paired_random_crop(img_HR, img_LR, patch_size, scale, gt_path)
            # augmentation - flip, rotate
            img_HR , img_LR = util.augment([patch_H, patch_L], self.opt.use_flip,self.opt.use_rot)


        img_HR  = util.img2tensor(img_HR)
        img_LR = util.img2tensor(img_LR)

        if LR_path is None:
            LR_path = HR_path
        return {'LR': img_LR, 'HR': img_HR, 'LR_path': LR_path, 'HR_path': HR_path}
Exemple #7
0
    def __getitem__(self, index):
        if self.HQ_envs is None:
            self._init_lmdb()

        HQ_size = self.opt["HQ_size"]
        env_idx, key = self.paths_HQ[index]
        name_a, name_b = key.split("_")
        target_frame_idx = int(name_b)

        # determine the neighbor frames
        # ensure not exceeding the borders
        neighbor_list = [target_frame_idx]
        name_b = "{:08d}".format(neighbor_list[0])

        # get the HQ image (as the center frame)
        img_HQ_l = []
        for v in neighbor_list:
            img_HQ = util.read_img(self.HQ_envs[env_idx],
                                   "{}_{:08d}".format(name_a,
                                                      v), (3, 720, 1280))
            img_HQ_l.append(img_HQ)

        # get LQ images
        img_LQ = util.read_img(self.LQ_envs[env_idx],
                               "{}_{:08d}".format(name_a, neighbor_list[-1]),
                               (3, 720, 1280))
        if self.opt["phase"] == "train":
            _, H, W = 3, 720, 1280  # LQ size
            # randomly crop
            rnd_h = random.randint(0, max(0, H - HQ_size))
            rnd_w = random.randint(0, max(0, W - HQ_size))
            img_LQ = img_LQ[rnd_h:rnd_h + HQ_size, rnd_w:rnd_w + HQ_size, :]
            img_HQ_l = [
                v[rnd_h:rnd_h + HQ_size, rnd_w:rnd_w + HQ_size, :]
                for v in img_HQ_l
            ]

            # augmentation - flip, rotate
            img_HQ_l.append(img_LQ)
            rlt = util.augment(img_HQ_l, self.opt["use_flip"],
                               self.opt["use_rot"])
            img_HQ_l = rlt[0:-1]
            img_LQ = rlt[-1]

        # stack LQ images to NHWC, N is the frame number
        img_HQs = np.stack(img_HQ_l, axis=0)
        # BGR to RGB, HWC to CHW, numpy to tensor
        img_LQ = img_LQ[:, :, [2, 1, 0]]
        img_HQs = img_HQs[:, :, :, [2, 1, 0]]
        img_LQ = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_LQ, (2, 0, 1)))).float()
        img_HQs = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_HQs, (0, 3, 1, 2)))).float()
        # print(img_LQ.shape, img_HQs.shape)

        if self.use_identical and np.random.randint(0, 10) == 0:
            img_LQ = img_HQs[-1, :, :, :]
            return {"LQ": img_LQ, "HQs": img_HQs, "identical_w": 10}

        return {"LQ": img_LQ, "HQs": img_HQs, "identical_w": 0}
Exemple #8
0
    def __getitem__(self, index):
        if self.data_type == 'lmdb' and self.GT_env is None:
            self._init_lmdb()

        image_index, scale_index = index
        scale = self.scales[scale_index]
        LQ_size = self.LQ_sizes[scale_index]
        key = self.paths_GT[image_index]
        name_a, name_b = key.split('_')
        center_frame_idx = int(name_b)

        neighbor_list, name_b = self.get_neighbor_list(center_frame_idx)
        key = name_a + '_' + name_b

        imgs = self.read_imgs(self.GT_root, name_a, neighbor_list)

        if self.opt['phase'] == 'train':
            _, H, W, _ = imgs.shape
            rnd_h = random.randint(0, max(0, H - self.GT_size))
            rnd_w = random.randint(0, max(0, W - self.GT_size))
            imgs = imgs[:, rnd_h:rnd_h + self.GT_size,
                        rnd_w:rnd_w + self.GT_size, :]
            img_GT = imgs[self.half_N_frames] / 255.
            is_PreGT = True if random.random() < self.pre_GT else False
            img_PreGT = imgs[self.half_N_frames - 1] / 255.
            img_LQs = [
                np.array(
                    Image.fromarray(img).resize(
                        (LQ_size, LQ_size), Image.BICUBIC)) / 255.
                for img in imgs
            ]
            # augmentation - flip, rotate
            img_LQs.extend([img_GT, img_PreGT])
            rlt = util.augment(img_LQs, self.opt['use_flip'],
                               self.opt['use_rot'])
            img_LQs = rlt[0:-2]
            img_GT = rlt[-2]
            img_PreGT = rlt[-1]

        # stack LQ images to NHWC, N is the frame number
        img_LQs = np.stack(img_LQs, axis=0)
        # BGR => RGB
        img_GT = img_GT[:, :, [2, 1, 0]]
        img_PreGT = img_PreGT[:, :, [2, 1, 0]]
        img_LQs = img_LQs[:, :, :, [2, 1, 0]]
        # NHWC => NCHW
        img_GT = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
        img_PreGT = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_PreGT, (2, 0, 1)))).float()
        img_LQs = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_LQs, (0, 3, 1, 2)))).float()
        return {
            'LQs': img_LQs,
            'GT': img_GT,
            'Pre': img_PreGT,
            'if_pre': is_PreGT,
            'key': key,
            'scale': scale
        }
    def __getitem__(self, index):
        # get full size image
        full_path = self.paths_hq[index % len(self.paths_hq)]
        loaded_img = util.read_img(None, full_path, None)
        img_full1 = util.channel_convert(loaded_img.shape[2], 'RGB',
                                         [loaded_img])[0]
        img_full2 = util.augment([img_full1], True, True)[0]
        img_full3 = get_square_image(img_full2)
        # This error crops up from time to time. I suspect an issue with util.read_img.
        if img_full3.shape[0] == 0 or img_full3.shape[1] == 0:
            print(
                "Error with image: %s. Loaded image shape: %s" %
                (full_path, str(loaded_img.shape)), str(img_full1.shape),
                str(img_full2.shape), str(img_full3.shape))
            # Attempt to recover by just using a fixed array of zeros, which the downstream networks should be fine training against, within reason.
            img_full3 = np.zeros((1024, 1024, 3), dtype=np.int)
        img_full = cv2.resize(img_full3, (self.hq_size_cap, self.hq_size_cap),
                              interpolation=cv2.INTER_AREA)
        patches_hq = [
            cv2.resize(img_full, (self.tile_size, self.tile_size),
                       interpolation=cv2.INTER_AREA)
        ]
        self.recursively_extract_patches(img_full, patches_hq, 1)
        # Image corruption is applied against the full size image for this dataset.
        img_corrupted = self.corruptor.corrupt_images([img_full])[0]
        patches_hq_corrupted = [
            cv2.resize(img_corrupted, (self.tile_size, self.tile_size),
                       interpolation=cv2.INTER_AREA)
        ]
        self.recursively_extract_patches(img_corrupted, patches_hq_corrupted,
                                         1)

        # BGR to RGB, HWC to CHW, numpy to tensor
        if patches_hq[0].shape[2] == 3:
            patches_hq = [
                cv2.cvtColor(p, cv2.COLOR_BGR2RGB) for p in patches_hq
            ]
            patches_hq_corrupted = [
                cv2.cvtColor(p, cv2.COLOR_BGR2RGB)
                for p in patches_hq_corrupted
            ]
        patches_hq = [
            torch.from_numpy(np.ascontiguousarray(np.transpose(
                p, (2, 0, 1)))).float() for p in patches_hq
        ]
        patches_hq = torch.stack(patches_hq, dim=0)
        patches_hq_corrupted = [
            torch.from_numpy(np.ascontiguousarray(np.transpose(
                p, (2, 0, 1)))).float() for p in patches_hq_corrupted
        ]
        patches_lq = [
            torch.nn.functional.interpolate(p.unsqueeze(0),
                                            scale_factor=1 / self.scale,
                                            mode='area').squeeze()
            for p in patches_hq_corrupted
        ]
        patches_lq = torch.stack(patches_lq, dim=0)

        d = {'lq': patches_lq, 'hq': patches_hq, 'GT_path': full_path}
        return d
Exemple #10
0
    def __getitem__(self, index):
        # 从文件中读出HR图片,返回HR、LR图片对
        HR_path, LR_path = None, None
        HR_path = self.HR_paths[index]
        scale = self.scale[self.idx_scale]
        patch_size = self.patch_size

        img_HR = util.read_img_(HR_path, self.n_colors)
        # 归一化
        img_HR = util.unit2single(img_HR)
        img_HR = util.modcrop(img_HR, self.scale)

        if self.opt['phase'] == 'train':
            l_max = 50
            theta = np.pi * np.random.rand(1)
            l1 = 0.1 + l_max * np.random.rand(1)
            l2 = 0.1 + (l1 - 0.1) * np.random.rand(1)
            kernel = util.anisotropic_Gaussian(ksize=self.ksize,
                                               theta=theta[0],
                                               l1=l1[0],
                                               l2=l2[0])
        else:
            kernel = util.anisotropic_Gaussian(ksize=self.ksize,
                                               theta=np.pi,
                                               l1=0.1,
                                               l2=0.1)
        k = np.reshape(kernel, (-1), order="F")
        k_reduced = np.dot(self.p, k)
        k_reduced = torch.from_numpy(k_reduced).float()

        H, W, _ = img_HR.shape
        img_LR = util.srmd_degradation(img_HR, kernel, self.scale)

        if not self.opt.test_only:
            patch_H, patch_L = util.paired_random_crop(img_HR, img_LR,
                                                       patch_size, scale)
            # augmentation - flip, rotate
            img_HR, img_LR = util.augment([patch_H, patch_L],
                                          self.opt.use_flip, self.opt.use_rot)
            if random.random < 0.1:
                noise_level = np.zeros(1).float()
            else:
                noise_level = np.random.uniform(0, 15)
        else:
            noise_level = np.zeros(self.sigma_test).float

        img_LR = img_LR + np.random.normal(0, noise_level / 255., img_LR.shape)

        img_HR = util.img_single2tensor(img_HR)
        img_LR = util.img_single2tensor(img_LR)

        M_vector = torch.cat((k_reduced, noise_level),
                             0).unsqueeze(1).unsqueeze(1)
        M = M_vector.repeat(1, img_LR.size()[-2], img_LR.size()[-1])

        img_LR = torch.cat((img_LR, M), 0)
        if LR_path is None:
            LR_path = HR_path
        return img_LR, img_HR, LR_path, HR_path
Exemple #11
0
    def __getitem__(self, index):
        HR_path, LR_path = None, None
        scale = self.opt['scale']

        # get HR image
        HR_path = self.paths_HR[index]
        img_HR = util.read_img(self.HR_env, HR_path)

        # get LR image
        if self.paths_LR:
            LR_path = self.paths_LR[index]
            D1_path = self.paths_D1[index]
            D2_path = self.paths_D2[index]
            D3_path = self.paths_D3[index]
            img_LR = util.read_img(self.LR_env, LR_path)
            img_D1 = util.read_img(self.D1_env, D1_path)
            img_D2 = util.read_img(self.D2_env, D2_path)
            img_D3 = util.read_img(self.D3_env, D3_path)
        else:  # down-sampling on-the-fly
            if self.opt['phase'] == 'train':
                # force to 3 channels
                if img_HR.ndim == 2:
                    img_HR = cv2.cvtColor(img_HR, cv2.COLOR_GRAY2BGR)

            H, W, _ = img_HR.shape
            # using matlab imresize
            img_LR = cv2.resize(img_HR, dsize=(int(H / scale), int(W / scale)), interpolation=cv2.INTER_CUBIC)
            img_D1 = cv2.resize(img_HR, dsize=(int(H / scale * 2), int(W / scale * 2)), interpolation=cv2.INTER_CUBIC)
            img_D2 = cv2.resize(img_HR, dsize=(int(H / scale * 4), int(W / scale * 4)), interpolation=cv2.INTER_CUBIC)
            img_D3 = cv2.resize(img_HR, dsize=(int(H / scale * 8), int(W / scale * 8)), interpolation=cv2.INTER_CUBIC)

            if img_LR.ndim == 2:
                img_LR = np.expand_dims(img_LR, axis=2)
                img_D1 = np.expand_dims(img_D1, axis=2)
                img_D2 = np.expand_dims(img_D2, axis=2)
                img_D3 = np.expand_dims(img_D3, axis=2)

        if self.opt['phase'] == 'train':
            # augmentation - flip, rotate
            img_LR, img_HR, img_D1, img_D2, img_D3 = util.augment([img_LR, img_HR, img_D1, img_D2, img_D3],
                                                                  self.opt['use_flip'], self.opt['use_rot'])
        # BGR to RGB, HWC to CHW, numpy to tensor
        if img_HR.shape[2] == 3:
            img_HR = img_HR[:, :, [2, 1, 0]]
            img_LR = img_LR[:, :, [2, 1, 0]]
            img_D1 = img_D1[:, :, [2, 1, 0]]
            img_D2 = img_D2[:, :, [2, 1, 0]]
            img_D3 = img_D3[:, :, [2, 1, 0]]
        img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
        img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
        img_D1 = torch.from_numpy(np.ascontiguousarray(np.transpose(img_D1, (2, 0, 1)))).float()
        img_D2 = torch.from_numpy(np.ascontiguousarray(np.transpose(img_D2, (2, 0, 1)))).float()
        img_D3 = torch.from_numpy(np.ascontiguousarray(np.transpose(img_D3, (2, 0, 1)))).float()

        if LR_path is None:
            LR_path = HR_path
        return {'LR': img_LR, 'HR': img_HR, 'LR_path': LR_path, 'HR_path': HR_path,
                'D1': img_D1, 'D2': img_D2, 'D3': img_D3, 'D1_path': D1_path, 'D2_path': D2_path, 'D3_path': D3_path}
    def __getitem__(self, index):
        HR_path, LR_path, MR_path = None, None, None
        scale = self.opt['scale']
        HR_size = self.opt['HR_size']
        # get HR image
        HR_path = self.paths_HR[index]
        img_HR = util.read_img(self.HR_env, HR_path)

        # modcrop in the validation / test phase
        # if self.opt['phase'] != 'train':
        #     img_HR = util.modcrop(img_HR, scale)
        # change color space if necessary
        if self.opt['color']:
            img_HR = util.channel_convert(img_HR.shape[2], self.opt['color'], [img_HR])[0]

        LR_path = self.paths_LR[index]
        img_LR = util.read_img(self.LR_env, LR_path)

        MR_path = self.paths_MR[index]
        img_MR = util.read_img(self.MR_env, MR_path)

        if self.opt['noise_gt']:
            img_MR = img_LR - img_MR

        if self.opt['phase'] == 'train':
            # if the image size is too small
            H, W, C = img_LR.shape
            LR_size = HR_size // scale

            # randomly crop
            rnd_h = random.randint(0, max(0, H - LR_size))
            rnd_w = random.randint(0, max(0, W - LR_size))
            img_MR = img_MR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
            img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
            rnd_h_HR, rnd_w_HR = int(rnd_h * scale), int(rnd_w * scale)
            img_HR = img_HR[rnd_h_HR:rnd_h_HR + HR_size, rnd_w_HR:rnd_w_HR + HR_size, :]

            # augmentation - flip, rotate
            img_MR, img_LR, img_HR = util.augment([img_MR, img_LR, img_HR], self.opt['use_flip'], \
                                          self.opt['use_rot'])

        # channel conversion
        if self.opt['color']:
            # img_HR, img_LR, img_MR = util.channel_convert(C, self.opt['color'], [img_HR, img_LR, img_MR])
            img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0]
            img_MR = util.channel_convert(C, self.opt['color'], [img_MR])[0]

        # BGR to RGB, HWC to CHW, numpy to tensor
        if img_HR.shape[2] == 3:
            img_HR = img_HR[:, :, [2, 1, 0]]
            img_LR = img_LR[:, :, [2, 1, 0]]
            img_MR = img_MR[:, :, [2, 1, 0]]
        img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
        img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
        img_MR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_MR, (2, 0, 1)))).float()

        return {'HR': img_HR, 'LR': img_LR, 'MR': img_MR, 'HR_path': HR_path, 'MR_path': MR_path, 'LR_path': LR_path}
    def __getitem__(self, index):
        if self.data_type == 'lmdb':
            if (self.GT_env is None) or (self.LQ_env is None):
                self._init_lmdb()
        GT_path, LQ_path = None, None
        scale = self.opt['scale']
        GT_size = self.opt['GT_size']

        # get GT image
        GT_path = self.paths_GT[index]
        if self.data_type == 'lmdb':
            resolution = [int(s) for s in self.sizes_GT[index].split('_')]
        else:
            resolution = None
        img_GT = util.read_img(self.GT_env, GT_path, resolution)
        # modcrop in the validation / test phase
        if self.opt['phase'] != 'train':
            img_GT = util.modcrop(img_GT, scale)
            # print('val img_GT shape: ', img_GT.shape)
        # change color space if necessary
        if self.opt['color']:
            img_GT = util.channel_convert(img_GT.shape[2], self.opt['color'],
                                          [img_GT])[0]
            # print('img_GT shape: ', img_GT.shape)
        if img_GT.ndim == 2:
            img_GT = np.expand_dims(img_GT, axis=2)
        if self.opt['phase'] == 'train':
            # if the image size is too small
            H, W, C = img_GT.shape
            if H < GT_size or W < GT_size:
                img_GT = cv2.resize(np.copy(img_GT), (GT_size, GT_size),
                                    interpolation=cv2.INTER_LINEAR)

            # randomly crop - GT img shape: [H, W, 1]
            rnd_h = random.randint(0, max(0, H - GT_size))
            rnd_w = random.randint(0, max(0, W - GT_size))
            img_GT = img_GT[rnd_h:rnd_h + GT_size, rnd_w:rnd_w + GT_size, :]

            # augmentation - flip, rotate
            img_GT = util.augment([img_GT], self.opt['use_flip'],
                                  self.opt['use_rot'])[0]

        # transform() - subsample k space - img_LF
        # input img shape [H, W, 1] or [H, W, 3]
        img_LF, img_GT = self.transform(img_GT, self.mask_func, self.seed)

        if LQ_path is None:
            LQ_path = GT_path
        return {
            'LQ': img_LF,
            'GT': img_GT,
            'LQ_path': LQ_path,
            'GT_path': GT_path
        }
Exemple #14
0
    def __getitem__(self, index):
        if self.data_type == 'lmdb' and self.GT_env is None:
            self._init_lmdb()

        LQ_size = self.opt['LQ_size']
        key = self.paths_GT[index]
        name_a, name_b = key.split('_')
        center_frame_idx = int(name_b)

        neighbor_list, name_b = self.get_neighbor_list(center_frame_idx)
        key = name_a + '_' + name_b

        extra = self.valid_nf // 2
        name_bs = neighbor_list[extra:self.nf - extra]

        img_GTs = self.read_imgs(self.GT_root, name_a, name_bs, True)
        img_LQs = self.read_imgs(self.LQ_root, name_a, neighbor_list)

        if self.opt['phase'] == 'train':
            H, W, _ = img_LQs[0].shape
            rnd_h = random.randint(0, max(0, H - LQ_size))
            rnd_w = random.randint(0, max(0, W - LQ_size))
            img_LQs = [
                v[rnd_h:rnd_h + LQ_size, rnd_w:rnd_w + LQ_size, :]
                for v in img_LQs
            ]
            rnd_h_HR, rnd_w_HR = int(rnd_h * self.scale), int(rnd_w *
                                                              self.scale)
            GT_size = int(LQ_size * self.scale)
            img_GTs = [
                v[rnd_h_HR:rnd_h_HR + GT_size, rnd_w_HR:rnd_w_HR + GT_size, :]
                for v in img_GTs
            ]
            # augmentation - flip, rotate
            num_gt = len(img_GTs)
            img_LQs += img_GTs
            img_LQs = util.augment(img_LQs, self.opt['use_flip'],
                                   self.opt['use_rot'])
            img_GTs = img_LQs[-num_gt:]
            img_LQs = img_LQs[:-num_gt]

        # stack LQ images to NHWC, N is the frame number
        img_LQs = np.stack(img_LQs, axis=0)
        img_GTs = np.stack(img_GTs, axis=0)
        # BGR => RGB
        img_GTs = img_GTs[:, :, :, [2, 1, 0]]
        img_LQs = img_LQs[:, :, :, [2, 1, 0]]
        # NHWC => NCHW
        img_GTs = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_GTs, (0, 3, 1, 2)))).float()
        img_LQs = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_LQs, (0, 3, 1, 2)))).float()
        return {'LQs': img_LQs, 'GT': img_GTs, 'key': key, 'scale': self.scale}
Exemple #15
0
    def __getitem__(self, index):
        if self.opt["data_type"] == "lmdb":
            if self.LR_env is None:
                self._init_lmdb()

        LR_path = None
        scale = self.opt["scale"]
        LR_size = self.opt["LR_size"]

        # get LR image
        LR_path = self.LR_paths[index]
        if self.opt["data_type"] == "lmdb":
            resolution = [int(s) for s in self.LR_sizes[index].split("_")]
        else:
            resolution = None
        img_LR = util.read_img(
            self.LR_env, LR_path, resolution
        )  # return: Numpy float32, HWC, BGR, [0,1]

        # modcrop in the validation / test phase
        if self.opt["phase"] != "train":
            img_LR = util.modcrop(img_LR, scale)

        if self.opt["phase"] == "train":
            H, W, C = img_LR.shape

            rnd_h = random.randint(0, max(0, H - LR_size))
            rnd_w = random.randint(0, max(0, W - LR_size))
            img_LR = img_LR[rnd_h : rnd_h + LR_size, rnd_w : rnd_w + LR_size, :]

            # augmentation - flip, rotate
            img_LR = util.augment(
                img_LR,
                self.opt["use_flip"],
                self.opt["use_rot"],
                self.opt["mode"],
            )

        # change color space if necessary
        if self.opt["color"]:
            img_LR = util.channel_convert(img_LR.shape[2], self.opt["color"], [img_LR])[
                0
            ]

        # BGR to RGB, HWC to CHW, numpy to tensor
        if img_LR.shape[2] == 3:
            img_LR = img_LR[:, :, [2, 1, 0]]
        img_LR = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))
        ).float()

        return {"LR": img_LR, "LR_path": LR_path}
Exemple #16
0
    def __getitem__(self, index):
        HR_path, LR_path = None, None
        scale = 4

        HR_path = self.paths_HR[index]
        img_HR = util.read_img(self.HR_env, HR_path)

        if self.cnf['phase'] != 'train':
            img_HR = util.modcrop(img_HR, scale)

        LR_path = self.paths_LR[index]
        img_LR = util.read_img(self.LR_env, LR_path)

        if self.cnf['phase'] == 'train':

            H, W, _ = img_HR.shape
            if H < self.hr_sz or W < self.hr_sz:
                img_HR = cv2.resize(np.copy(img_HR), (self.hr_sz, self.hr_sz),
                                    interpolation=cv2.INTER_LINEAR)

                img_LR = util.imresize_np(img_HR, 1 / scale, True)
                if img_LR.ndim == 2:
                    img_LR = np.expand_dims(img_LR, axis=2)

            H, W, C = img_LR.shape
            LR_size = self.hr_sz // scale

            rnd_h = random.randint(0, max(0, H - LR_size))
            rnd_w = random.randint(0, max(0, W - LR_size))
            img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
            rnd_h_HR, rnd_w_HR = int(rnd_h * scale), int(rnd_w * scale)
            img_HR = img_HR[rnd_h_HR:rnd_h_HR + self.hr_sz,
                            rnd_w_HR:rnd_w_HR + self.hr_sz, :]

            img_LR, img_HR = util.augment([img_LR, img_HR], True, True)

        if img_HR.shape[2] == 3:
            img_HR = img_HR[:, :, [2, 1, 0]]
            img_LR = img_LR[:, :, [2, 1, 0]]
        img_HR = from_numpy(
            np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
        img_LR = from_numpy(
            np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()

        if LR_path is None:
            LR_path = HR_path
        return {
            'LR': img_LR,
            'HR': img_HR,
            'LR_path': LR_path,
            'HR_path': HR_path
        }
Exemple #17
0
    def __getitem__(self, index):
        GT_path, NOISY_path = None, None
        GT_size = self.opt['GT_size']

        # get GT image
        GT_path = self.paths_GT[index]
        # img_GT = util.read_img(self.GT_env, GT_path)
        # for HDRs
        img_GT = util.load_reference_mat("", GT_path)

        # get NOISY image
        if self.paths_NOISY:
            NOISY_path = self.paths_NOISY[index]
            all_feature = util.load_feature_mat(
                self.FEATURE_DIR,
                NOISY_path)  #util.load_feature_mat_complete_tungsten
            img_NOISY = all_feature[:, :, :3]
            features = all_feature[:, :, 3:]
        # do the cropping
        H, W, _ = img_GT.shape
        # x = random.randint(0, np.maximum(0, W - GT_size))
        # y = random.randint(0, np.maximum(0, H - GT_size))
        x = 0
        y = 0
        img_GT = util._crop(img_GT, (y, x), GT_size)
        img_NOISY = util._crop(img_NOISY, (y, x), GT_size)
        features = util._crop(features, (y, x), GT_size)

        if self.opt['phase'] == 'train':
            # augmentation - flip, rotate
            img_NOISY, img_GT, features = util.augment([img_NOISY, img_GT, features], self.opt['use_flip'], \
                self.opt['use_rot'])

        img_GT = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
        img_NOISY = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_NOISY, (2, 0, 1)))).float()
        features = torch.from_numpy(
            np.ascontiguousarray(np.transpose(features, (2, 0, 1)))).float()

        return {
            'diffuse_in': img_NOISY,
            'diffuse_ref': img_GT,
            "seg": features,
            "category": 1,
            'NOISY_path': NOISY_path,
            'GT_path': GT_path,
            "x_offset": 512,
            "y_offset": 512
        }
Exemple #18
0
    def __getitem__(self, index):
        HR_path, LR_path = None, None
        HR_size = self.opt['HR_size']

        # get HR image
        HR_path = self.paths_HR[index]
        img_HR = util.read_img(self.HR_env, HR_path)

        # get LR image
        LR_path = self.paths_LR[index]
        img_LR = util.read_img(self.LR_env, LR_path)

        # modcrop in the validation / test phase
        if self.opt['phase'] != 'train':
            img_HR = util.modcrop(img_HR, 2)
            img_LR = util.modcrop(img_LR, 2)

        if self.opt['phase'] == 'train':

            H, W, C = img_LR.shape
            LR_size = HR_size

            # randomly crop
            rnd_h = random.randint(0, max(0, H - LR_size))
            rnd_w = random.randint(0, max(0, W - LR_size))
            img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
            img_HR = img_HR[rnd_h:rnd_h + HR_size, rnd_w:rnd_w + HR_size, :]

            # augmentation - flip, rotate
            img_LR, img_HR = util.augment([img_LR, img_HR], self.opt['use_flip'], \
                self.opt['use_rot'])

        # BGR to RGB, HWC to CHW, numpy to tensor
        if img_HR.shape[2] == 3:
            img_HR = img_HR[:, :, [2, 1, 0]]
            img_LR = img_LR[:, :, [2, 1, 0]]
        img_HR = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
        img_LR = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()

        if LR_path is None:
            LR_path = HR_path
        return {
            'LR': img_LR,
            'HR': img_HR,
            'LR_path': LR_path,
            'HR_path': HR_path
        }
Exemple #19
0
    def __getitem__(self, index):
        if self.opt["data_type"] == "lmdb":
            if self.LR_env is None:
                self._init_lmdb()

        LR_size = self.LR_size
        SR_size = self.SR_size
        scale = self.opt["scale"]

        # get real kernel map
        real_ker_map = self.real_ker_map_list[index]
        # get each kernel map
        ker_map = self.ker_map_list[index]

        # get LR image
        LR_path = self.LR_paths[index]
        if self.opt["data_type"] == "lmdb":
            resolution = [int(s) for s in self.LR_sizes[index].split("_")]
        else:
            resolution = None
        img_LR = util.read_img(self.LR_env, LR_path, resolution)
        H, W, C = img_LR.shape

        # get SR image
        img_SR = self.SR_img_list[index]

        if self.opt["phase"] == "train":
            # randomly crop
            rnd_h = random.randint(0, max(0, H - LR_size))
            rnd_w = random.randint(0, max(0, W - LR_size))
            rnd_h_SR, rnd_w_SR = int(rnd_h * scale), int(rnd_w * scale)
            img_SR = img_SR[rnd_h_SR:rnd_h_SR + SR_size,
                            rnd_w_SR:rnd_w_SR + SR_size, :]

            # augmentation - flip, rotate
            img_SR = util.augment(img_SR, self.opt["use_flip"],
                                  self.opt["use_rot"], self.opt["mode"])

        # change color space if necessary
        if self.opt["color"]:
            img_SR = util.channel_convert(C, self.opt["color"], [img_SR])[0]

        # BGR to RGB, HWC to CHW, numpy to tensor
        if img_SR.shape[2] == 3:
            img_SR = img_SR[:, :, [2, 1, 0]]
        img_SR = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_SR, (2, 0, 1)))).float()

        return {"SR": img_SR, "real_ker": real_ker_map, "ker": ker_map}
    def __getitem__(self, index):
        GT_path, NOISY_path = None, None
        GT_size = self.opt['GT_size']

        # get GT image
        GT_path = self.paths_GT[index]
        NOISY_path = self.paths_NOISY[index]
        #   print(GT_path)
        #  print(NOISY_path)
        img_GT = util.read_img_GIGAN_version(GT_path)
        # for HDRs
        #  img_GT = util.load_reference_mat_shading("", GT_path, NOISY_path)#图片那边还需要改!!!!

        # get NOISY image

        all_feature = util.load_feature_GIGAN_version(
            NOISY_path,
            self.DIFFUSE_DIR)  #util.load_feature_mat_complete_tungsten
        img_NOISY = all_feature[:, :, :3]
        features = all_feature[:, :, 3:]
        # do the cropping
        # H, W, _ = img_GT.shape
        # x = random.randint(0, np.maximum(0, W - GT_size))
        # y = random.randint(0, np.maximum(0, H - GT_size))
        # img_GT = util._crop(img_GT, (y,x), GT_size)
        # img_NOISY = util._crop(img_NOISY, (y,x), GT_size)
        # features = util._crop(features, (y,x), GT_size)

        if self.opt['phase'] == 'train':
            # augmentation - flip, rotate
            img_NOISY, img_GT, features = util.augment([img_NOISY, img_GT, features], self.opt['use_flip'], \
                self.opt['use_rot'])

        img_GT = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
        img_NOISY = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_NOISY, (2, 0, 1)))).float()
        features = torch.from_numpy(
            np.ascontiguousarray(np.transpose(features, (2, 0, 1)))).float()
        #  all_feature = torch.from_numpy(np.ascontiguousarray(np.transpose(all_feature, (2, 0, 1)))).float()

        return {
            'NOISY': img_NOISY,
            'GT': img_GT,
            "seg": features,
            "category": 1,
            'NOISY_path': NOISY_path,
            'GT_path': GT_path
        }
Exemple #21
0
 def __getitem__(self, index):
     # 从文件中读出HR图片,返回HR、LR图片对
     HR_path = self.HR_paths[index]
     LR_path = self.LR_paths[index]
     img_HR = util.read_img_(HR_path, self.n_colors)
     img_LR = util.read_img_(LR_path, self.n_colors)
     if not self.opt.test_only == 'train':
         patch_H, patch_L = util.paired_random_crop(img_HR, img_LR,
                                                    self.patch_size,
                                                    self.scale)
         img_HR, img_LR = util.augment([patch_H, patch_L],
                                       self.opt.use_flip, self.opt.use_rot)
     img_HR = util.img2tensor(img_HR)
     img_LR = util.img2tensor(img_LR)
     return {'LR': img_LR, 'HR': img_HR}
Exemple #22
0
 def __getitem__(self, index):
     lrs, gt, video_name, filename = self.read_data(index)
     if self.train:
         lrs, gt = self.random_crop(lrs, gt, self.patch_size, self.scale)
         lrs.append(gt)
         rlt = util.augment(lrs)
         # rlt = [util.rgb2ycbcr(img, only_y=True) for img in rlt]
         lrs, gt = rlt[:-1], rlt[-1]
     else:
         lrs.append(gt)
         # rlt = [util.rgb2ycbcr(img, only_y=True) for img in lrs]
         lrs, gt = rlt[:-1], rlt[-1]
     lrs, gt = self.to_tensor(lrs, gt)
     # return lrs, gt, video_name, filename
     key = '{}_{}'.format(video_name, filename)
     return {'LQs': lrs, 'GT': gt, 'key': key}
Exemple #23
0
 def __getitem__(self, index):
     # 从文件中读出HR图片,返回HR、LR图片对
     HR_path, LR_path = None, None
     HR_path = self.HR_paths[index]
     img_HR = util.read_img_(HR_path, self.n_colors)
     # 归一化
     img_HR = util.unit2single(img_HR)
     img_HR = util.modcrop(img_HR, self.scale)
     kernel = fspecial('gaussian', 15, 2.6)
     H, W, _ = img_HR.shape
     img_LR = util.srmd_degradation(img_HR, kernel, self.scale)
     if self.opt.phase == 'train':
         patch_H, patch_L = util.paired_random_crop(img_HR, img_LR,
                                                    self.patch_size,
                                                    self.scale)
         # augmentation - flip, rotate
         img_HR, img_LR = util.augment([patch_H, patch_L],
                                       self.opt.use_flip, self.opt.use_rot)
         img_HR = util.img_single2tensor(img_HR)
         img_LR = util.img_single2tensor(img_LR)
         if random.random() < 0.1:
             noise_level = torch.zeros(1).float()
         else:
             noise_level = torch.FloatTensor([
                 np.random.uniform(self.sigma_min, self.sigma_max)
             ]) / 255.0
     else:
         img_HR = util.img_single2tensor(img_HR)
         img_LR = util.img_single2tensor(img_LR)
         noise_level = torch.FloatTensor([self.sigma_test])
     noise = torch.randn(img_LR.size()).mul_(noise_level).float()
     img_LR.add_(noise)
     kernel = util.single2tensor3(np.expand_dims(np.float32(kernel),
                                                 axis=2))
     noise_level = torch.FloatTensor([noise_level]).view([1, 1, 1])
     if LR_path is None:
         LR_path = HR_path
     return {
         'LR': img_LR,
         'HR': img_HR,
         'k': kernel,
         'sigma': noise_level,
         'sf': self.scale,
         'LR_path': LR_path,
         'HR_path': HR_path
     }
Exemple #24
0
    def __getitem__(self, index):
        image_index, scale_index = index
        scale = self.scales[scale_index]
        LQ_size = self.opt['LQ_size']
        key = self.paths_GT[image_index]
        name_a, name_b = key.split('_')
        center_frame_idx = int(name_b)

        neighbor_list, name_b = self.get_neighbor_list(center_frame_idx)

        #### get the images
        img_GT = self.read_imgs(self.GT_root, name_a, name_b)[0]
        img_LQs = self.read_imgs(self.LQ_root, name_a, neighbor_list, scale)

        if self.opt['phase'] == 'train':
            H, W, _ = img_LQs[0].shape
            rnd_h = random.randint(0, max(0, H - LQ_size))
            rnd_w = random.randint(0, max(0, W - LQ_size))
            img_LQs = [
                v[rnd_h:rnd_h + LQ_size, rnd_w:rnd_w + LQ_size, :]
                for v in img_LQs
            ]
            rnd_h_HR, rnd_w_HR = int(rnd_h * scale), int(rnd_w * scale)
            GT_size = int(LQ_size * scale)
            img_GT = img_GT[rnd_h_HR:rnd_h_HR + GT_size,
                            rnd_w_HR:rnd_w_HR + GT_size, :]
            # augmentation - flip, rotate
            img_LQs.append(img_GT)
            rlt = util.augment(img_LQs, self.opt['use_flip'],
                               self.opt['use_rot'])
            img_LQs = rlt[0:-1]
            img_GT = rlt[-1]

        # stack LQ images to NHWC, N is the frame number
        img_LQs = np.stack(img_LQs, axis=0)
        # BGR => RGB
        img_GT = img_GT[:, :, [2, 1, 0]]
        img_LQs = img_LQs[:, :, :, [2, 1, 0]]
        # NHWC => NCHW
        img_GT = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
        img_LQs = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_LQs, (0, 3, 1, 2)))).float()
        return {'LQs': img_LQs, 'GT': img_GT, 'key': key, 'scale': scale}
    def __getitem__(self, index):


        # get LQ and HQ image
        LQ_path = self.paths_LQ[index]
        HQ_path = self.paths_HQ[index]

        img_LQ = util.read_img(None, LQ_path, None)
        img_HQ = util.read_img(None, HQ_path, None)

        if self.opt['crop_size']:
            # randomly crop
            LH, LW, _ = img_LQ.shape
            rnd_h = random.randint(0, max(0, LH - self.crop_size))
            rnd_w = random.randint(0, max(0, LW - self.crop_size))

            rnd_hh = rnd_h * self.opt['scale']
            rnd_wh = rnd_w * self.opt['scale']
            patch_size = self.crop_size * self.opt['scale']
            img_LQ = img_LQ[rnd_h:rnd_h + self.crop_size, rnd_w:rnd_w + self.crop_size, :]
            img_HQ = img_HQ[rnd_hh:rnd_hh + patch_size, rnd_wh:rnd_wh + patch_size, :]

        if self.opt['phase'] == 'train':
            # augmenttation cutclur
            if self.opt['use_cutblur']:
                img_LQ,img_HQ=util.argment_cutblur(img_LQ, img_HQ,self.opt['scale'])
            if self.opt['use_rgbpermute']:
                img_LQ, img_HQ=util.argment_rgb(img_LQ, img_HQ)
            # augmentation - flip, rotate
            img_LQ, img_HQ = util.augment([img_LQ, img_HQ], self.opt['use_flip'],
                                          self.opt['use_rot'])

        # BGR to RGB, HWC to CHW, numpy to tensor
        if img_LQ.shape[2] == 3:
            img_LQ = img_LQ[:, :, [2, 1, 0]]
        if img_HQ.shape[2] == 3:
            img_HQ = img_HQ[:, :, [2, 1, 0]]

        img_LQ = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LQ, (2, 0, 1)))).float()
        img_HQ = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HQ, (2, 0, 1)))).float()

        return {'LQ': img_LQ, 'LQ_path': LQ_path, 'HQ': img_HQ, 'HQ_path': HQ_path}
Exemple #26
0
    def __getitem__(self, index):
        if self.data_type == 'lmdb' and self.GT_env is None:
            self._init_lmdb()

        image_index, scale_index = index
        scale = self.scales[scale_index]
        LQ_size = self.LQ_sizes[scale_index]
        key = self.paths_GT[image_index]

        imgs = self.read_imgs(key)

        if self.opt['phase'] == 'train':
            _, H, W, _ = imgs.shape
            rnd_h = random.randint(0, max(0, H - self.GT_size))
            rnd_w = random.randint(0, max(0, W - self.GT_size))
            imgs = imgs[:, rnd_h:rnd_h + self.GT_size,
                        rnd_w:rnd_w + self.GT_size, :]
            img_GT = imgs[self.half_N_frames] / 255.
            img_LQs = [
                np.array(
                    Image.fromarray(img).resize(
                        (LQ_size, LQ_size), Image.BICUBIC)) / 255.
                for img in imgs
            ]
            img_LQs.append(img_GT)
            rlt = util.augment(img_LQs, self.opt['use_flip'],
                               self.opt['use_rot'])
            img_LQs = rlt[0:-1]
            img_GT = rlt[-1]

        # stack LQ images to NHWC, N is the frame number
        img_LQs = np.stack(img_LQs, axis=0)
        # BGR => RGB, HWC to CHW, numpy to tensor
        img_GT = img_GT[:, :, [2, 1, 0]]
        img_LQs = img_LQs[:, :, :, [2, 1, 0]]
        # NHWC => NCHW
        img_GT = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
        img_LQs = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_LQs, (0, 3, 1, 2)))).float()
        return {'LQs': img_LQs, 'GT': img_GT, 'key': key, 'scale': scale}
Exemple #27
0
    def __getitem__(self, index):
        if self.data_type == 'lmdb' and self.GT_env is None:
            self._init_lmdb()

        key = self.paths_GT[index]
        name_a, name_b = key.split('_')
        center_frame_idx = int(name_b)
        neighbors, name_b = self.get_neighbor_list(center_frame_idx)
        key = name_a + "_" + name_b
        keys = [name_a + "_{:08d}".format(nei) for nei in neighbors]

        img_GTs = self.read_imgs(keys, True)
        img_LQs = self.read_imgs(keys, False)

        if self.opt['phase'] == 'train':
            H, W, _ = img_LQs[0].shape
            rnd_h = random.randint(0, max(0, H - self.LQ_size))
            rnd_w = random.randint(0, max(0, W - self.LQ_size))
            img_LQs = [v[rnd_h : rnd_h + self.LQ_size, rnd_w : rnd_w + \
                            self.LQ_size, :] for v in img_LQs]
            rnd_h_HR, rnd_w_HR = int(rnd_h * self.scale), int(rnd_w *
                                                              self.scale)

            img_GTs = [v[rnd_h_HR : rnd_h_HR + self.GT_size, rnd_w_HR : rnd_w_HR + \
                            self.GT_size, :] for v in img_GTs]
            # augmentation - flip, rotate
            img_LQs += img_GTs
            rlts = util.augment(img_LQs, self.opt['use_flip'],
                                self.opt['use_rot'])
            img_GTs = rlts[self.nframes:]
            img_LQs = rlts[:self.nframes]

        # stack images to NHWC; BGR -> RGB; NHWC -> CNHW
        img_GTs = np.stack(img_GTs, axis=0)[:, :, :, (2, 1, 0)].transpose(
            (3, 0, 1, 2))
        img_LQs = np.stack(img_LQs, axis=0)[:, :, :, (2, 1, 0)].transpose(
            (3, 0, 1, 2))
        img_GTs = torch.from_numpy(img_GTs).float()
        img_LQs = torch.from_numpy(img_LQs).float()
        return {'LQs': img_LQs, 'GT': img_GTs, 'key': key, 'scale': self.scale}
    def __getitem__(self, index):
        if self.data_type == "mc":
            self._ensure_memcached()
        elif self.data_type == "lmdb" and (self.HQ_env is None
                                           or self.LQ_env is None):
            self._init_lmdb()

        HQ_size = self.opt["HQ_size"]
        key = self.paths_HQ[index]
        name_a, name_b = key.split("_")

        # get the HQ image
        img_HQ = util.read_img(self.HQ_env, key, (3, 720, 1280))

        # get the LQ image
        img_LQ = util.read_img(self.LQ_env, key, (3, 720, 1280))

        if self.opt["phase"] == "train":
            _, H, W = 3, 720, 1280  # LQ size
            # randomly crop
            rnd_h = random.randint(0, max(0, H - HQ_size))
            rnd_w = random.randint(0, max(0, W - HQ_size))
            img_LQ = img_LQ[rnd_h:rnd_h + HQ_size, rnd_w:rnd_w + HQ_size, :]
            img_HQ = img_HQ[rnd_h:rnd_h + HQ_size, rnd_w:rnd_w + HQ_size, :]

            # augmentation - flip, rotate
            imgs = [img_HQ, img_LQ]
            rlt = util.augment(imgs, self.opt["use_flip"], self.opt["use_rot"])
            img_HQ = rlt[0]
            img_LQ = rlt[1]

        # BGR to RGB, HWC to CHW, numpy to tensor
        img_LQ = img_LQ[:, :, [2, 1, 0]]
        img_HQ = img_HQ[:, :, [2, 1, 0]]
        img_LQ = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_LQ, (2, 0, 1)))).float()
        img_HQ = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_HQ, (2, 0, 1)))).float()

        return {"LQ": img_LQ, "HQ": img_HQ}
Exemple #29
0
    def __getitem__(self, index):
        if self.opt['data_type'] == 'lmdb':
            if self.LR_env is None:
                self._init_lmdb()

        LR_size = self.LR_size

        # get LR image, kernel map
        LR_path = self.LR_paths[index]
        ker_map = self.ker_maps[index]
        if self.opt['data_type'] == 'lmdb':
            resolution = [int(s) for s in self.LR_sizes[index].split('_')]
        else:
            resolution = None
        img_LR = util.read_img(self.LR_env, LR_path, resolution)
        H, W, C = img_LR.shape

        if self.opt['phase'] == 'train':
            #randomly crop
            rnd_h = random.randint(0, max(0, H - LR_size))
            rnd_w = random.randint(0, max(0, W - LR_size))
            img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]

            # augmentation - flip, rotate
            img_LR = util.augment(img_LR, self.opt['use_flip'],
                                  self.opt['use_rot'], self.opt['mode'])

        # change color space if necessary
        if self.opt['color']:
            img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0]

        # BGR to RGB, HWC to CHW, numpy to tensor
        if img_LR.shape[2] == 3:
            img_LR = img_LR[:, :, [2, 1, 0]]
        img_LR = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()

        return {'LQ': img_LR, 'ker': ker_map, 'LQ_path': LR_path}
Exemple #30
0
    def __getitem__(self, index):
        path, target = self.samples[index]
        '''img_dis = self.loader(path[0])
        img_ref = self.loader(path[1])'''
        img_dis = read_img(env=None, path=path[0])
        img_ref = read_img(env=None, path=path[1])
        '''if self.transform is not None:
            img_dis = self.transform(img_dis)
            img_ref = self.transform(img_ref)'''

        if self.patch_size < 288:
            H, W, _ = img_ref.shape
            crop_size = self.patch_size
            rnd_h = random.randint(0, max(0, (H - crop_size)))
            rnd_w = random.randint(0, max(0, (W - crop_size)))
            img_dis = img_dis[rnd_h:rnd_h + crop_size,
                              rnd_w:rnd_w + crop_size, :]
            img_ref = img_ref[rnd_h:rnd_h + crop_size,
                              rnd_w:rnd_w + crop_size, :]

        # augmentation - flip, rotate
        img_dis, img_ref = augment([img_dis, img_ref],
                                   self.opt['use_flip'],
                                   rot=False)

        if img_ref.shape[2] == 3:
            img_ref = img_ref[:, :, [2, 1, 0]]
            img_dis = img_dis[:, :, [2, 1, 0]]

        img_ref = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_ref, (2, 0, 1)))).float()
        img_dis = torch.from_numpy(
            np.ascontiguousarray(np.transpose(img_dis, (2, 0, 1)))).float()

        img_dis = self.transform(img_dis)
        img_ref = self.transform(img_ref)

        return {'Dis': img_dis, 'Ref': img_ref, 'Label': target}