Ejemplo n.º 1
0
    def __getitem__(self, idx):
        # 将二进制文件转为imageio
        lr, hr, filename = self._load_file(idx)
        pair = self.get_patch(lr, hr)
        pair = common.set_channel(*pair, n_channels=self.args.n_colors)
        # 生成kernel
        pair_t = common.np2Tensor(*pair, rgb_range=self.args.rgb_range)
        k = utils_sisr.gen_kernel(scale_factor=np.array(
            [self.scale, self.scale]))  # Gaussian blur
        r_value = np.random.randint(0, 8)
        if r_value > 3:
            k = utils_deblur.blurkernel_synthesis(h=25)  # motion blur
        else:
            sf_k = random.choice(self.scale)
            k = utils_sisr.gen_kernel(scale_factor=np.array(
                [sf_k, sf_k]))  # Gaussian blur
            mode_k = np.random.randint(0, 8)
            k = util.augment_img(k, mode=mode_k)
            if np.random.randint(0, 8) == 1:
                noise_level = 0 / 255.0
            else:
                noise_level = np.random.randint(0, self.sigma_max) / 255.0

            # ---------------------------
            # Low-quality image
            # ---------------------------
            img_L = ndimage.filters.convolve(patch_H,
                                             np.expand_dims(k, axis=2),
                                             mode='wrap')
            img_L = img_L[0::self.sf, 0::self.sf, ...]
            # add Gaussian noise
            img_L = util.uint2single(img_L) + np.random.normal(
                0, noise_level, img_L.shape)
            img_H = patch_H

        return pair_t[0], pair_t[1], filename
Ejemplo n.º 2
0
    def __getitem__(self, index):

        # -------------------
        # get H image
        # -------------------
        H_path = self.paths_H[index]
        img_H = util.imread_uint(H_path, self.n_channels)
        L_path = H_path

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

            # ---------------------------
            # 1) scale factor, ensure each batch only involves one scale factor
            # ---------------------------
            if self.count % self.opt['dataloader_batch_size'] == 0:
                # sf = random.choice([1,2,3,4])
                self.sf = random.choice(self.scales)
                # self.count = 0  # optional
            self.count += 1
            H, W, _ = img_H.shape

            # ----------------------------
            # randomly crop the patch
            # ----------------------------
            rnd_h = random.randint(0, max(0, H - self.patch_size))
            rnd_w = random.randint(0, max(0, W - self.patch_size))
            patch_H = img_H[rnd_h:rnd_h + self.patch_size,
                            rnd_w:rnd_w + self.patch_size, :]

            # ---------------------------
            # augmentation - flip, rotate
            # ---------------------------
            mode = np.random.randint(0, 8)
            patch_H = util.augment_img(patch_H, mode=mode)

            # ---------------------------
            # 2) kernel
            # ---------------------------
            r_value = random.randint(0, 7)
            if r_value > 3:
                k = utils_deblur.blurkernel_synthesis(h=25)  # motion blur
            else:
                sf_k = random.choice(self.scales)
                k = utils_sisr.gen_kernel(scale_factor=np.array(
                    [sf_k, sf_k]))  # Gaussian blur
                mode_k = random.randint(0, 7)
                k = util.augment_img(k, mode=mode_k)

            # ---------------------------
            # 3) noise level
            # ---------------------------
            if random.randint(0, 8) == 1:
                noise_level = 0 / 255.0
            else:
                noise_level = np.random.randint(0, self.sigma_max) / 255.0

            # ---------------------------
            # Low-quality image
            # ---------------------------
            img_L = ndimage.filters.convolve(patch_H,
                                             np.expand_dims(k, axis=2),
                                             mode='wrap')
            img_L = img_L[0::self.sf, 0::self.sf, ...]
            # add Gaussian noise
            img_L = util.uint2single(img_L) + np.random.normal(
                0, noise_level, img_L.shape)
            img_H = patch_H

        else:

            k = self.kernels[0, 0].astype(np.float64)  # validation kernel
            k /= np.sum(k)
            noise_level = 0. / 255.0  # validation noise level
            img_L = ndimage.filters.convolve(img_H,
                                             np.expand_dims(k, axis=2),
                                             mode='wrap')  # blur
            img_L = img_L[0::self.sf_validation, 0::self.sf_validation,
                          ...]  # downsampling
            img_L = util.uint2single(img_L) + np.random.normal(
                0, noise_level, img_L.shape)

        k = util.single2tensor3(np.expand_dims(np.float32(k), axis=2))
        img_H, img_L = util.uint2tensor3(img_H), util.single2tensor3(img_L)
        noise_level = torch.FloatTensor([noise_level]).view([1, 1, 1])

        return {
            'L': img_L,
            'H': img_H,
            'k': k,
            'sigma': noise_level,
            'sf': self.sf,
            'L_path': L_path,
            'H_path': H_path
        }