예제 #1
0
파일: trainer.py 프로젝트: alexquach/TTSR
    def evaluate(self, current_epoch=0):
        self.logger.info('Epoch ' + str(current_epoch) +
                         ' evaluation process...')

        if (self.args.dataset == 'HMDB_FRAMES'):
            self.model.eval()
            with torch.no_grad():
                psnr, ssim, cnt = 0., 0., 0
                for i_batch, sample_batched in enumerate(
                        self.dataloader['test']['1']):
                    cnt += 1
                    sample_batched = self.prepare(sample_batched)
                    lr = sample_batched['LR']
                    lr_sr = sample_batched['LR_sr']
                    hr = sample_batched['HR']
                    ref = sample_batched['Ref']
                    ref_sr = sample_batched['Ref_sr']
                    hr = hr.repeat(1, 3, 1, 1)

                    # sr, _, _, _, _ = naive_averaging(
                    # self.model, lr, lr_sr, hr, ref, ref_sr) TODO uncomment for modified fusion model
                    sr, _, _, _, _ = self.model(lr=lr,
                                                lrsr=lr_sr,
                                                ref=ref,
                                                refsr=ref_sr)

                    if (self.args.eval_save_results):
                        sr_save = (sr + 1.) * 127.5
                        sr_save = np.transpose(
                            sr_save.squeeze().round().cpu().numpy(),
                            (1, 2, 0)).astype(np.uint8)
                        imsave(
                            os.path.join(self.args.save_dir, 'save_results',
                                         str(i_batch).zfill(5) + '.png'),
                            sr_save)

                    # calculate psnr and ssim
                    # sr = sr.squeeze(0)
                    # hr = hr.squeeze(0) TODO uncomment for modified fusion model
                    _psnr, _ssim = calc_psnr_and_ssim(sr.detach(), hr.detach())

                    psnr += _psnr
                    ssim += _ssim

                psnr_ave = psnr / cnt
                ssim_ave = ssim / cnt
                self.logger.info('Ref  PSNR (now): %.3f \t SSIM (now): %.4f' %
                                 (psnr_ave, ssim_ave))
                if (psnr_ave > self.max_psnr):
                    self.max_psnr = psnr_ave
                    self.max_psnr_epoch = current_epoch
                if (ssim_ave > self.max_ssim):
                    self.max_ssim = ssim_ave
                    self.max_ssim_epoch = current_epoch
                self.logger.info(
                    'Ref  PSNR (max): %.3f (%d) \t SSIM (max): %.4f (%d)' %
                    (self.max_psnr, self.max_psnr_epoch, self.max_ssim,
                     self.max_ssim_epoch))

        self.logger.info('Evaluation over.')
예제 #2
0
    def test_FixedRef(self, lr_path, ref_path, save_path):

        ### LR and LR_sr
        HR = imread(lr_path)
        h1, w1 = HR.shape[:2]
        h1, w1 = h1 // 4 * 4, w1 // 4 * 4
        HR = HR[:h1, :w1, :]
        LR = np.array(
            Image.fromarray(HR).resize((w1 // 4, h1 // 4), Image.BICUBIC))
        LR_sr = np.array(Image.fromarray(LR).resize((w1, h1), Image.BICUBIC))

        ### Ref and Ref_sr
        Ref = imread(ref_path)
        h2, w2 = Ref.shape[:2]
        h2, w2 = h2 // 4 * 4, w2 // 4 * 4
        Ref = Ref[:h2, :w2, :]
        Ref_sr = np.array(
            Image.fromarray(Ref).resize((w2 // 4, h2 // 4), Image.BICUBIC))
        Ref_sr = np.array(
            Image.fromarray(Ref_sr).resize((w2, h2), Image.BICUBIC))

        ### change type
        HR = HR.astype(np.float32)
        LR = LR.astype(np.float32)
        LR_sr = LR_sr.astype(np.float32)
        Ref = Ref.astype(np.float32)
        Ref_sr = Ref_sr.astype(np.float32)

        ### rgb range to [-1, 1]
        HR = HR / 127.5 - 1.
        LR = LR / 127.5 - 1.
        LR_sr = LR_sr / 127.5 - 1.
        Ref = Ref / 127.5 - 1.
        Ref_sr = Ref_sr / 127.5 - 1.

        ### to tensor
        HR_t = torch.from_numpy(HR.transpose(
            (2, 0, 1))).unsqueeze(0).float().to(self.device)
        LR_t = torch.from_numpy(LR.transpose(
            (2, 0, 1))).unsqueeze(0).float().to(self.device)
        LR_sr_t = torch.from_numpy(LR_sr.transpose(
            (2, 0, 1))).unsqueeze(0).float().to(self.device)
        Ref_t = torch.from_numpy(Ref.transpose(
            (2, 0, 1))).unsqueeze(0).float().to(self.device)
        Ref_sr_t = torch.from_numpy(Ref_sr.transpose(
            (2, 0, 1))).unsqueeze(0).float().to(self.device)

        self.model.eval()
        with torch.no_grad():
            sr, _, _, _, _ = self.model(lr=LR_t,
                                        lrsr=LR_sr_t,
                                        ref=Ref_t,
                                        refsr=Ref_sr_t)
            sr_save = (sr + 1.) * 127.5
            sr_save = np.transpose(sr_save.squeeze().round().cpu().numpy(),
                                   (1, 2, 0)).astype(np.uint8)
            imsave(save_path, sr_save)
        psnr, ssim = calc_psnr_and_ssim(sr.detach(), HR_t.detach())
        return psnr, ssim
예제 #3
0
    def evaluate(self, current_epoch=0):
        self.logger.info('Epoch ' + str(current_epoch) +
                         ' evaluation process...')

        if (
                self.args.dataset == 'IMM'
        ):  #Dataset name has to be manually altered here.. better Solution!
            self.model.eval()
            with torch.no_grad():
                psnr, ssim, cnt = 0., 0., 0
                for i_batch, sample_batched in enumerate(
                        self.dataloader['test']['1']):
                    cnt += 1
                    sample_batched = self.prepare(sample_batched)
                    lr = sample_batched['LR']
                    lr_sr = sample_batched['LR_sr']
                    hr = sample_batched['HR']
                    ref = sample_batched['Ref']
                    ref_sr = sample_batched['Ref_sr']

                    sr, _, _, _, _ = self.model(lr=lr,
                                                lrsr=lr_sr,
                                                ref=ref,
                                                refsr=ref_sr)

                    if (self.args.eval_save_results):
                        sr_save = (sr + 1.) * 127.5
                        sr_save = np.transpose(
                            sr_save.squeeze().round().cpu().numpy(), (1, 2, 0)
                        ).astype(
                            np.uint8
                        )  # squeeze delets all 1 values --- round() rounds to closest integer -- .cpu() moves object to cpu ---.numpy transforms object to numpy array--- transform to np.unit8 type
                        imsave(
                            os.path.join(self.args.save_dir, 'save_results',
                                         str(i_batch).zfill(5) + '.tif'),
                            sr_save)

                    ### calculate psnr and ssim
                    _psnr, _ssim = calc_psnr_and_ssim(sr.detach(), hr.detach())

                    psnr += _psnr
                    ssim += _ssim

                psnr_ave = psnr / cnt
                ssim_ave = ssim / cnt
                self.logger.info('Ref  PSNR (now): %.3f \t SSIM (now): %.4f' %
                                 (psnr_ave, ssim_ave))
                if (psnr_ave > self.max_psnr):
                    self.max_psnr = psnr_ave
                    self.max_psnr_epoch = current_epoch
                if (ssim_ave > self.max_ssim):
                    self.max_ssim = ssim_ave
                    self.max_ssim_epoch = current_epoch
                self.logger.info(
                    'Ref  PSNR (max): %.3f (%d) \t SSIM (max): %.4f (%d)' %
                    (self.max_psnr, self.max_psnr_epoch, self.max_ssim,
                     self.max_ssim_epoch))

        self.logger.info('Evaluation over.')
예제 #4
0
    def evaluate(self, filepath):
        input_image = align_image(load_image(filepath), self.scale)
        input_y_image = resize_image(convert_rgb_to_y(input_image),
                                     1 / self.scale)
        input_scaled_y_image = resize_image(input_y_image, self.scale)

        output_y_image = self.run(input_y_image, input_scaled_y_image)
        ground_truth_y_image = convert_rgb_to_y(input_image)

        return calc_psnr_and_ssim(ground_truth_y_image,
                                  output_y_image,
                                  border=self.scale)
예제 #5
0
파일: trainer.py 프로젝트: alexquach/TTSR
    def evaluate_flownet(self, args, current_epoch=0):
        self.logger.info('Epoch ' + str(current_epoch) +
                         ' evaluation process...')

        if (self.args.dataset == 'HMDB_FLOWNET'):
            self.model.eval()
            with torch.no_grad():
                psnr, ssim, cnt = 0., 0., 0
                for i_batch, sample_batched in enumerate(
                        self.dataloader['test']['1']):
                    cnt += 1
                    sample_batched = self.prepare(sample_batched)
                    lr = sample_batched['LR'].float(
                    ) / 127.5 - 1  # [5, 3, 40, 40]
                    lr_sr = sample_batched['LR_sr'].float(
                    ) / 127.5 - 1  # [5, 3, 160, 160]
                    hr = sample_batched['HR'].float(
                    ) / 127.5 - 1  # [1, 3, 160, 160]
                    ref = sample_batched['Ref'].float(
                    ) / 127.5 - 1  # [1, 3, 160, 160]
                    ref_sr = sample_batched['Ref_sr'].float(
                    ) / 127.5 - 1  # [1, 3, 160, 160]

                    if args.train_style == "normal":
                        sr, S, T_lv3, T_lv2, T_lv1 = self.model(
                            lr=lr[:, 2, :, :, :],
                            lrsr=lr_sr[:, 2, :, :, :],
                            ref=ref[:, :, :, :],
                            refsr=ref_sr[:, :, :, :])
                    elif args.train_style == "average":
                        sr, S, T_lv3, T_lv2, T_lv1 = flownet_naive_averaging(
                            self.model, lr, lr_sr, hr, ref, ref_sr)
                    else:
                        sr, S, T_lv3, T_lv2, T_lv1 = flownet_conv3d_1x1(
                            self.model, lr, lr_sr, hr, ref, ref_sr)

                    if (self.args.eval_save_results):
                        sr_save = (sr + 1.) * 127.5
                        sr_save = np.transpose(
                            sr_save.squeeze().round().cpu().numpy(),
                            (1, 2, 0)).astype(np.uint8)
                        imsave(
                            os.path.join(self.args.save_dir, 'save_results',
                                         str(i_batch).zfill(5) + '.png'),
                            sr_save)

                    # calculate psnr and ssim
                    # sr = sr.squeeze(0)
                    # hr = hr.squeeze(0) TODO uncomment for modified fusion model
                    _psnr, _ssim = calc_psnr_and_ssim(sr.detach(), hr.detach())

                    psnr += _psnr
                    ssim += _ssim

                psnr_ave = psnr / cnt
                ssim_ave = ssim / cnt
                self.logger.info('Ref  PSNR (now): %.3f \t SSIM (now): %.4f' %
                                 (psnr_ave, ssim_ave))
                if (psnr_ave > self.max_psnr):
                    self.max_psnr = psnr_ave
                    self.max_psnr_epoch = current_epoch
                if (ssim_ave > self.max_ssim):
                    self.max_ssim = ssim_ave
                    self.max_ssim_epoch = current_epoch
                self.logger.info(
                    'Ref  PSNR (max): %.3f (%d) \t SSIM (max): %.4f (%d)' %
                    (self.max_psnr, self.max_psnr_epoch, self.max_ssim,
                     self.max_ssim_epoch))

        self.logger.info('Evaluation over.')