Пример #1
0
Файл: test.py Проект: yhu9/RCAN
    def evaluate_ideal(self, lr, hr):
        if self.model == 'ESRGAN': lr = lr * 1.0 / 255.0
        h, w, d = lr.shape
        hrh, hrw, _ = hr.shape
        canvas = np.zeros((h * 4, w * 4, 3))
        size = self.PATCH_SIZE
        stride = self.PATCH_SIZE // 2
        pad = self.PATCH_SIZE // 8
        info = {'assignment': np.zeros((hrh, hrw))}

        for i in range(0, h - 1, stride):
            for j in range(0, w - 1, stride):
                lr_img = lr[i:i + size] if i + size < h else lr[i:]
                lr_img = lr_img[:,
                                j:j + size, :] if j + size < w else lr_img[:,
                                                                           j:]
                #lr_img = torch.from_numpy(np.transpose(lr_img[:,:,[2,1,0]],(2,0,1))).float()
                lr_img = torch.FloatTensor(lr_img).to(self.device)
                lr_img = lr_img.permute((2, 0, 1)).unsqueeze(0)
                #lr_img = lr_img.to(self.device)
                hr_img = hr[i * self.upsize:(i + size) * self.upsize,
                            j * self.upsize:(j + size) * self.upsize, :]

                psnrscores = []
                ssimscores = []
                sr_predictions = []
                for sisr in self.SRmodels:
                    hr_hat = sisr(lr_img)
                    hr_hat = hr_hat.squeeze(0).permute(1, 2,
                                                       0).data.cpu().numpy()

                    if self.model == 'ESRGAN': hr_hat = hr_hat * 255.0
                    hr_hat = np.clip(hr_hat, 0, 255)
                    sr_predictions.append(hr_hat)
                    psnr, ssim = util.calc_metrics(hr_img,
                                                   hr_hat,
                                                   crop_border=self.upsize)
                    psnrscores.append(psnr)
                    ssimscores.append(ssim)

                top, top_ = (0, 0) if i == 0 else ((i + pad) * self.upsize,
                                                   pad * self.upsize)
                bot, bot_ = (hrh, size * self.upsize) if i + size >= h else (
                    (i + size - pad) * self.upsize, -pad * self.upsize)
                lef, lef_ = (0, 0) if j == 0 else ((j + pad) * self.upsize,
                                                   pad * self.upsize)
                rig, rig_ = (hrw, size * self.upsize) if j + size >= w else (
                    (j + size - pad) * self.upsize, -pad * self.upsize)

                idx = psnrscores.index(max(psnrscores))
                canvas[top:bot, lef:rig] = sr_predictions[idx][top_:bot_,
                                                               lef_:rig_]
                info['assignment'][top:bot, lef:rig] = idx

        psnr, ssim = util.calc_metrics(hr, canvas, crop_border=self.upsize)
        info['psnr'] = psnr
        info['ssim'] = ssim
        info['SRimg'] = canvas.astype(np.uint8)
        return psnr, ssim, info
Пример #2
0
    def evaluate_ideal(self,lr,hr):
        lr = lr * 1.0 / 255
        h,w,d = lr.shape
        hrh,hrw,_ = hr.shape
        canvas = np.zeros((h*4,w*4,3))
        size = self.PATCH_SIZE
        stride = self.PATCH_SIZE // 2
        pad = self.PATCH_SIZE // 8
        info = []

        for i in range(0,h-1,stride):
            for j in range(0,w-1,stride):

                lr_img = lr[i:i+size] if i+size < h else lr[i:]
                lr_img = lr_img[:,j:j+size,:] if j+size < w else lr_img[:,j:]
                lr_img = torch.from_numpy(np.transpose(lr_img[:,:,[2,1,0]],(2,0,1))).float()
                lr_img = lr_img.unsqueeze(0)
                lr_img = lr_img.to(self.device)
                hr_img = hr[i*self.upsize:(i+size)*self.upsize,j*self.upsize:(j+size)*self.upsize,:]

                psnrscores = []
                ssimscores = []
                sr_predictions = []
                for sisr in self.SRmodels:
                    hr_hat = sisr(lr_img).data.squeeze().float().cpu().clamp_(0,1).numpy()
                    hr_hat = np.transpose(hr_hat[[2,1,0],:,:],(1,2,0))
                    hr_hat = (hr_hat * 255.0).round()
                    sr_predictions.append(hr_hat)

                    psnr,ssim = util.calc_metrics(hr_img,hr_hat,crop_border=self.upsize)
                    psnrscores.append(psnr)
                    ssimscores.append(ssim)

                top,top_= (0,0) if i == 0 else ((i+pad)*self.upsize,pad*self.upsize)
                bot,bot_ = (hrh,size*self.upsize) if i+size >= h else ((i+size-pad)*self.upsize,-pad*self.upsize)
                lef,lef_ = (0,0) if j == 0 else ((j+pad)*self.upsize,pad*self.upsize)
                rig,rig_ = (hrw,size*self.upsize) if j+size >= w else ((j+size-pad)*self.upsize,-pad*self.upsize)

                info.append(max(psnrscores))
                idx = psnrscores.index(max(psnrscores))
                canvas[top:bot,lef:rig] = sr_predictions[idx][top_:bot_,lef_:rig_]

                #cv2.imshow('srimg',canvas.astype(np.uint8))
                #cv2.imshow('gtimg',hr.astype(np.uint8))
                #cv2.waitKey(1)

        psnr,ssim = util.calc_metrics(hr,canvas,crop_border=self.upsize)
        print(psnr,ssim)
        return psnr,ssim,np.array(info)
Пример #3
0
Файл: test.py Проект: yhu9/RCAN
    def evaluate_baseline(self, lr, hr):
        if self.model == 'ESRGAN': lr = lr * 1.0 / 255.0
        img = torch.FloatTensor(lr).to(self.device)
        lr_img = img.permute((2, 0, 1)).unsqueeze(0)

        self.SRmodels[0].eval()
        SR = self.SRmodels[0](lr_img)
        SR = SR.squeeze(0).permute(1, 2, 0).data.cpu().numpy()
        if self.model == 'ESRGAN':
            SR = np.clip(SR * 255.0, 0, 255)
            psnr, ssim = util.calc_metrics(hr, SR, 4)
        elif self.model == 'RCAN':
            SR = np.clip(SR, 0, 255)
            psnr, ssim = util.calc_metrics(hr, SR, 4)
        return psnr, ssim, SR
Пример #4
0
def Test(solver, dataloader, solver_log, current_epoch):
    psnr_list = []
    ssim_list = []
    val_loss = []

    for iter, batch in enumerate(dataloader):
        solver.feed_data(batch)
        iter_loss = solver.test()
        val_loss.append(iter_loss)

        ##### calculate psnr/ssim metrics #####
        visuals = solver.get_current_visual()
        psnr, ssim = util.calc_metrics(visuals['SR'],
                                       visuals['HR'],
                                       crop_border=scale)
        psnr_list.append(psnr)
        ssim_list.append(ssim)

    print("PSNR: {:.2f} SSIM: {:.4f}  Loss: {:.3f}".format(
        sum(psnr_list) / len(psnr_list),
        sum(ssim_list) / len(ssim_list),
        sum(val_loss) / len(val_loss)))

    return sum(val_loss)/len(val_loss), sum(psnr_list)/len(psnr_list),\
            sum(ssim_list)/len(ssim_list)
Пример #5
0
Файл: test.py Проект: yhu9/RCAN
    def evaluate_baseline(self,lr,hr):
        #lr = lr * 1.0 / 255
        img = torch.FloatTensor(lr).to(self.device)
        lr_img = img.permute((2,0,1)).unsqueeze(0)

        self.SRmodels[0].eval()
        SR = self.SRmodels[0](lr_img)
        SR = SR.squeeze(0).permute(1,2,0).data.cpu().numpy()

        psnr,ssim = util.calc_metrics(hr,SR,4)
        SR = (SR).round().astype(np.uint8)
        return psnr,ssim,SR
Пример #6
0
Файл: test.py Проект: yhu9/RCAN
    def evaluate_baseline(self, lr, hr):
        lr = lr * 1.0 / 255
        lr_img = torch.from_numpy(lr).to(self.device).permute(
            2, 0, 1).unsqueeze(0).float()

        SR = self.SRmodels[0](lr_img).data.squeeze().permute(1, 2,
                                                             0).cpu().clamp_(
                                                                 0, 1).numpy()
        SR = (SR * 255).round().astype(np.uint8)

        psnr, ssim = util.calc_metrics(hr, SR, crop_border=self.upsize)
        return psnr, ssim, SR
Пример #7
0
    def applySISR(self,lr,action,hr):

        self.SRoptimizers[action].zero_grad()
        hr_hat = self.SRmodels[action](lr)
        loss = F.l1_loss(hr_hat,hr)
        loss.backward()
        self.SRoptimizers[action].step()

        hr_hat = hr_hat.squeeze(0).permute(1,2,0); hr = hr.squeeze(0).permute(1,2,0)
        hr_hat = hr_hat.detach().cpu().numpy()
        hr = hr.detach().cpu().numpy()
        psnr,ssim = util.calc_metrics(hr_hat,hr,crop_border=self.UPSIZE)

        return hr_hat, psnr, ssim, loss.item()
Пример #8
0
def Test(global_solver, val_loader, solver_log, current_r):
    psnr_list = []
    ssim_list = []
    val_loss_list = []

    for iter, batch in enumerate(val_loader):
        global_solver.feed_data(batch)
        iter_loss = global_solver.test()
        val_loss_list.append(iter_loss)

        ##### Calculate psnr/ssim metrics #####
        visuals = global_solver.get_current_visual()
        psnr, ssim = util.calc_metrics(visuals['SR'], visuals['HR'], crop_border=scale)
        psnr_list.append(psnr)
        ssim_list.append(ssim)
 
    # ##### record loss/psnr/ssim #####
    solver_log['records']['val_loss'].append(' ')
    solver_log['records']['val_loss'].append(sum(val_loss_list)/len(val_loss_list))

    solver_log['records']['psnr'].append(' ')
    solver_log['records']['psnr'].append(sum(psnr_list)/len(psnr_list))

    solver_log['records']['ssim'].append(' ')
    solver_log['records']['ssim'].append(sum(ssim_list)/len(ssim_list))

    ##### record the best epoch #####
    round_is_best = False
    if solver_log['best_pred'] < (sum(psnr_list)/len(psnr_list)):
        solver_log['best_pred'] = (sum(psnr_list)/len(psnr_list))
        round_is_best = True
        solver_log['best_round'] = current_r


    print("PSNR: %.2f  SSIM: %.4f  Loss: %.6f  Best PSNR: %.2f in Round: [%d]" 
    %(sum(psnr_list)/len(psnr_list), sum(ssim_list)/len(ssim_list), sum(val_loss_list)/len(val_loss_list),
      solver_log['best_pred'], solver_log['best_round']))

    global_solver.set_current_log(solver_log)
    global_solver.save_checkpoint(current_r, round_is_best)
    global_solver.save_current_log()

    return sum(val_loss_list)/len(val_loss_list), sum(psnr_list)/len(psnr_list),\
           sum(ssim_list)/len(ssim_list)
Пример #9
0
def Validate(client_solver, val_loader):
    psnr_list = []
    ssim_list = []
    val_loss_list = []

    for iter, batch in enumerate(val_loader):
        client_solver.feed_data(batch)
        iter_loss = client_solver.test()
        val_loss_list.append(iter_loss)

        ##### calculate psnr/ssim metrics #####
        visuals = client_solver.get_current_visual()
        psnr, ssim = util.calc_metrics(visuals['SR'], visuals['HR'], crop_border=scale)
        psnr_list.append(psnr)
        ssim_list.append(ssim)

    print("PSNR: %.2f  SSIM: %.4f" 
    %(sum(psnr_list)/len(psnr_list), sum(ssim_list)/len(ssim_list)))

    return sum(psnr_list)/len(psnr_list), sum(ssim_list)/len(ssim_list)
Пример #10
0
    def validate(self):
        scores = {}
        self.model.eval()
        for vset in self.validationsets:
            scores[vset] = []
            HR_dir = os.path.join(self.hr_rootdir, vset)
            LR_dir = os.path.join(os.path.join(self.lr_rootdir, vset),
                                  self.resfolder)

            #APPLY MODEL ON LR IMAGES
            HR_files = [os.path.join(HR_dir, f) for f in os.listdir(HR_dir)]
            LR_files = [os.path.join(LR_dir, f) for f in os.listdir(LR_dir)]
            HR_files.sort()
            LR_files.sort()

            #PSNR/SSIM SCORE FOR CURRENT VALIDATION SET
            for hr_file, lr_file in zip(HR_files, LR_files):
                hr = cv2.imread(hr_file, cv2.IMREAD_COLOR)
                lr = cv2.imread(lr_file, cv2.IMREAD_COLOR) * 1.0 / 255
                lr_img = torch.from_numpy(
                    np.transpose(lr[:, :, [2, 1, 0]], (2, 0, 1))).float()
                lr_img = lr_img.unsqueeze(0)
                lr_img = lr_img.to(self.device)

                out = self.model(lr_img).data.squeeze().float().cpu().clamp_(
                    0, 1).numpy()
                out = np.transpose(out[[2, 1, 0], :, :], (1, 2, 0))
                out = (out * 255.0).round()

                psnr, ssim = util.calc_metrics(hr,
                                               out,
                                               crop_border=self.upsize)
                print(hr_file, psnr)
                scores[vset].append([psnr, ssim])

            mu_psnr = np.mean(np.array(scores[vset])[:, 0])
            mu_ssim = np.mean(np.array(scores[vset])[:, 1])
            print(vset + ' scores', mu_psnr, mu_ssim)
def main():
    parser = argparse.ArgumentParser(
        description='Test Super Resolution Models')
    parser.add_argument('-opt',
                        type=str,
                        required=True,
                        help='Path to options JSON file.')
    opt = option.parse(parser.parse_args().opt)
    opt = option.dict_to_nonedict(opt)

    # initial configure
    scale = opt['scale']
    degrad = opt['degradation']
    network_opt = opt['networks']
    model_name = network_opt['which_model'].upper()

    # create folders
    util.mkdir_and_rename(opt['path']['res_root'])
    option.save(opt)

    # create test dataloader
    bm_names = []
    test_loaders = []
    for ds_name, dataset_opt in sorted(opt['datasets'].items()):
        test_set = create_dataset(dataset_opt)
        test_loader = create_dataloader(test_set, dataset_opt)
        test_loaders.append(test_loader)
        print('===> Test Dataset: [%s]   Number of images: [%d]' %
              (dataset_opt['name'], len(test_set)))
        bm_names.append(dataset_opt['name'])

    # create solver (and load model)
    solver = create_solver(opt)
    # Test phase
    print('===> Start Test')
    print("==================================================")
    print("Method: %s || Scale: %d || Degradation: %s" %
          (model_name, scale, degrad))

    for bm, test_loader in zip(bm_names, test_loaders):
        print("Test set : [%s]" % bm)

        sr_list = []
        path_list = []

        total_psnr = []
        total_ssim = []
        total_time = []
        res_dict = OrderedDict()

        need_HR = False if test_loader.dataset.__class__.__name__.find(
            'HR') < 0 else True

        for iter, batch in tqdm(enumerate(test_loader),
                                total=len(test_loader)):
            solver.feed_data(batch, need_HR=need_HR, need_landmark=False)

            # calculate forward time
            t0 = time.time()
            solver.test()
            t1 = time.time()
            total_time.append((t1 - t0))

            visuals = solver.get_current_visual(need_HR=need_HR)
            sr_list.append(visuals['SR'][-1])

            # calculate PSNR/SSIM metrics on Python
            if need_HR:
                psnr, ssim = util.calc_metrics(visuals['SR'][-1],
                                               visuals['HR'],
                                               crop_border=scale)
                total_psnr.append(psnr)
                total_ssim.append(ssim)
                path_list.append(
                    os.path.basename(batch['HR_path'][0]).replace(
                        'HR', model_name))
                # print(
                #     "[%d/%d] %s || PSNR(dB)/SSIM: %.2f/%.4f || Timer: %.4f sec ."
                #     % (iter + 1, len(test_loader),
                #        os.path.basename(batch['HR_path'][0]), psnr, ssim,
                #        (t1 - t0)))
                res_dict[path_list[-1]] = {
                    'psnr': psnr,
                    'ssim': ssim,
                    'time': t1 - t0
                }

            else:
                path_list.append(os.path.basename(batch['LR_path'][0]))
                # print("[%d/%d] %s || Timer: %.4f sec ." %
                #       (iter + 1, len(test_loader),
                #        os.path.basename(batch['LR_path'][0]), (t1 - t0)))

        if need_HR:
            print("---- Average PSNR(dB) /SSIM /Speed(s) for [%s] ----" % bm)
            average_res_str = "PSNR: %.2f      SSIM: %.4f      Speed: %.4f" % \
                  (sum(total_psnr) / len(total_psnr), sum(total_ssim) /
                   len(total_ssim), sum(total_time) / len(total_time))
            print(average_res_str)
        else:
            print("---- Average Speed(s) for [%s] is %.4f sec ----" %
                  (bm, sum(total_time) / len(total_time)))

        # save SR results for further evaluation on MATLAB
        save_img_path = os.path.join(opt['path']['res_root'], bm)

        print("===> Saving SR images of [%s]... Save Path: [%s]\n" %
              (bm, save_img_path))

        if not os.path.exists(save_img_path): os.makedirs(save_img_path)
        for img, name in zip(sr_list, path_list):
            imageio.imwrite(os.path.join(save_img_path, name), img)
        if need_HR:
            with open(os.path.join(save_img_path, 'result.json'), 'w') as f:
                json.dump(res_dict, f, indent=2)
            with open(os.path.join(save_img_path, 'average_result.txt'),
                      'w') as f:
                f.write(average_res_str + '\n')

    print("==================================================")
    print("===> Finished !")
Пример #12
0
Файл: test.py Проект: yhu9/RCAN
 def getstats(self, sr, hr):
     psnr, ssim = util.calc_metrics(hr, sr, crop_border=self.upsize)
     return psnr, ssim
Пример #13
0
def main():
    parser = argparse.ArgumentParser(
        description='Test Super Resolution Models')
    parser.add_argument('-opt',
                        type=str,
                        required=True,
                        help='Path to options JSON file.')
    opt = option.parse(parser.parse_args().opt)
    opt = option.dict_to_nonedict(opt)

    # initial configure
    scale = opt['scale']
    degrad = opt['degradation']
    network_opt = opt['networks']
    model_name = network_opt['which_model'].upper()
    if opt['self_ensemble']: model_name += 'plus'

    # create test dataloader
    bm_names = []
    test_loaders = []
    for _, dataset_opt in sorted(opt['datasets'].items()):
        test_set = create_dataset(dataset_opt)
        test_loader = create_dataloader(test_set, dataset_opt)
        test_loaders.append(test_loader)
        print('===> Test Dataset: [%s]   Number of images: [%d]' %
              (test_set.name(), len(test_set)))
        bm_names.append(test_set.name())

    # create solver (and load model)
    solver = create_solver(opt)
    # Test phase
    print('===> Start Test')
    print("==================================================")
    print("Method: %s || Scale: %d || Degradation: %s" %
          (model_name, scale, degrad))

    for bm, test_loader in zip(bm_names, test_loaders):
        print("Test set : [%s]" % bm)

        sr_list = []
        path_list = []

        total_psnr = []
        total_ssim = []
        total_time = []

        need_HR = False if test_loader.dataset.__class__.__name__.find(
            'LRHR') < 0 else True

        for iter, batch in enumerate(test_loader):

            solver.feed_data(batch, need_HR=need_HR)

            # calculate forward time
            t0 = time.time()
            solver.test()
            t1 = time.time()
            total_time.append((t1 - t0))

            visuals = solver.get_current_visual(need_HR=need_HR)
            sr_list.append(visuals['SR'])

            # calculate PSNR/SSIM metrics on Python
            if need_HR:
                psnr, ssim = util.calc_metrics(visuals['SR'],
                                               visuals['HR'],
                                               crop_border=scale)
                total_psnr.append(psnr)
                total_ssim.append(ssim)
                path_list.append(
                    os.path.basename(batch['HR_path'][0]).replace(
                        'HR', model_name))
                print(
                    "[%d/%d] %s || PSNR(dB)/SSIM: %.2f/%.4f || Timer: %.4f sec ."
                    % (iter + 1, len(test_loader),
                       os.path.basename(batch['LR_path'][0]), psnr, ssim,
                       (t1 - t0)))
            else:
                path_list.append(os.path.basename(batch['LR_path'][0]))
                print("[%d/%d] %s || Timer: %.4f sec ." %
                      (iter + 1, len(test_loader),
                       os.path.basename(batch['LR_path'][0]), (t1 - t0)))

        if need_HR:
            print("---- Average PSNR(dB) /SSIM /Speed(s) for [%s] ----" % bm)
            print("PSNR: %.2f      SSIM: %.4f      Speed: %.4f" %
                  (sum(total_psnr) / len(total_psnr), sum(total_ssim) /
                   len(total_ssim), sum(total_time) / len(total_time)))
        else:
            print("---- Average Speed(s) for [%s] is %.4f sec ----" %
                  (bm, sum(total_time) / len(total_time)))

        # save SR results for further evaluation on MATLAB
        if need_HR:
            save_img_path = os.path.join('./results/SR/' + degrad, model_name,
                                         bm, "x%d" % scale)
        else:
            save_img_path = os.path.join('./results/SR/' + bm, model_name,
                                         "x%d" % scale)

        print("===> Saving SR images of [%s]... Save Path: [%s]\n" %
              (bm, save_img_path))

        if not os.path.exists(save_img_path): os.makedirs(save_img_path)
        for img, name in zip(sr_list, path_list):
            imageio.imwrite(os.path.join(save_img_path, name), img)

    print("==================================================")
    print("===> Finished !")
Пример #14
0
def main():
    parser = argparse.ArgumentParser(description='Test Super Resolution Models')
    parser.add_argument('-opt', type=str, required=True, help='Path to options JSON file.')
    opt = option.parse(parser.parse_args().opt)
    opt = option.dict_to_nonedict(opt)

    # make sure the CUDA_VISIBLE_DEVICES is set before import torch.
    from utils import util
    from solvers import create_solver
    from datasets import create_dataloader
    from datasets import create_dataset

    # initial configure
    scale = opt['scale']
    degrad = opt['degradation']
    network_opt = opt['networks']
    model_name = network_opt['which_model'].upper()
    if opt['self_ensemble']: model_name += 'plus'

    # create test dataloader
    bm_names =[]
    test_loaders = []
    percent10 = True
    for _, dataset_opt in sorted(opt['datasets'].items()):
        test_set = create_dataset(dataset_opt)
        test_loader = create_dataloader(test_set, dataset_opt)
        test_loaders.append(test_loader)
        print('===> Test Dataset: [%s]   Number of images: [%d]' % (test_set.name(), len(test_set)))
        bm_names.append(test_set.name())

    # create solver (and load model)
    solver = create_solver(opt)
    # Test phase
    print('===> Start Test')
    print("==================================================")
    print("Method: %s || Scale: %d || Degradation: %s"%(model_name, scale, degrad))

    for bm, test_loader in zip(bm_names, test_loaders):
        print("Test set : [%s]"%bm)

        sr_list = []
        path_list = []

        total_psnr = []
        total_ssim = []
        total_time = []

        need_HR = False if test_loader.dataset.__class__.__name__.find('LRHR') < 0 else True

        for iter, batch in enumerate(test_loader):
            solver.feed_data(batch, need_HR=need_HR)

            # calculate forward time
            t0 = time.time()
            solver.test()
            t1 = time.time()
            total_time.append((t1 - t0))

            visuals = solver.get_current_visual(need_HR=need_HR)
            sr_list.append(visuals['SR'])

            # calculate PSNR/SSIM metrics on Python
            # 这里仅支持batch size = 1的情况!!!
            if need_HR:
                psnr, ssim = util.calc_metrics(visuals['SR'], visuals['HR'], crop_border=scale)
                total_psnr.append(psnr)
                total_ssim.append(ssim)
                path_list.append(os.path.basename(batch['HR_path'][0]).replace('HR', model_name))
                print("[%d/%d] %s || PSNR(dB)/SSIM: %.2f/%.4f || Timer: %.4f sec ." % (iter+1, len(test_loader),
                                                                                       os.path.basename(batch['LR_path'][0]),
                                                                                       psnr, ssim,
                                                                                       (t1 - t0)))
            else:
                file_dir = batch['LR_path'][0].split('/')[-2]
                path_list.append(os.path.join(file_dir, os.path.basename(batch['LR_path'][0])))
                print("[%d/%d] %s || Timer: %.4f sec ." % (iter + 1, len(test_loader),
                                                           os.path.join(file_dir, os.path.basename(batch['LR_path'][0])),
                                                           (t1 - t0)))

        if need_HR:
            print("---- Average PSNR(dB) /SSIM /Speed(s) for [%s] ----" % bm)
            print("PSNR: %.2f      SSIM: %.4f      Speed: %.4f" % (sum(total_psnr)/len(total_psnr),
                                                                  sum(total_ssim)/len(total_ssim),
                                                                  sum(total_time)/len(total_time)))
        else:
            print("---- Average Speed(s) for [%s] is %.4f sec ----" % (bm,
                                                                      sum(total_time)/len(total_time)))

        if need_HR:
            save_img_path = os.path.join('../submit/SR/'+degrad, model_name, bm, "x%d"%scale)
        else:
            save_img_path = os.path.join('../submit/')

        print("===> Saving SR images of [%s]... Save Path: [%s]\n" % (bm, save_img_path))

        middle_name = 'h_Res' if percent10 else 'h_Sub25_Res'
        filter_idx = -1 if percent10 else -7

        if not os.path.exists(save_img_path): os.makedirs(save_img_path)
        for img, name in zip(sr_list, path_list):
            store_path = os.path.join(save_img_path, name)
            base_dir = os.path.dirname(store_path)[:filter_idx] + middle_name
            if not os.path.exists(base_dir): os.makedirs(base_dir)
            store_path = os.path.join(base_dir, os.path.basename(name))
            print('write into {}.'.format(store_path))
            imageio.imwrite(store_path, img)

        percent10 = False

    print("==================================================")
    print("===> Finished !")
Пример #15
0
Файл: train.py Проект: ashj9/AFN
def main():
    parser = argparse.ArgumentParser(
        description='Train Super Resolution Models')
    parser.add_argument('-opt',
                        type=str,
                        required=True,
                        help='Path to options JSON file.')
    opt = option.parse(parser.parse_args().opt)
    print('Inside Train.py')
    print(opt['datasets']['train']['data_path'])

    # random seed
    seed = opt['solver']['manual_seed']
    if seed is None: seed = random.randint(1, 10000)
    print("===> Random Seed: [%d]" % seed)
    random.seed(seed)
    torch.manual_seed(seed)

    # create train and val dataloader
    for phase, dataset_opt in sorted(opt['datasets'].items()):
        if phase == 'train':
            train_set = create_dataset(dataset_opt)
            train_loader = create_dataloader(train_set, dataset_opt)
            print('===> Train Dataset: %s   Number of images: [%d]' %
                  (train_set.name(), len(train_set)))
            if train_loader is None:
                raise ValueError("[Error] The training data does not exist")

        elif phase == 'val':
            val_set = create_dataset(dataset_opt)
            val_loader = create_dataloader(val_set, dataset_opt)
            print('===> Val Dataset: %s   Number of images: [%d]' %
                  (val_set.name(), len(val_set)))

        else:
            raise NotImplementedError(
                "[Error] Dataset phase [%s] in *.json is not recognized." %
                phase)

    solver = create_solver(opt)

    scale = opt['scale']
    model_name = opt['networks']['which_model'].upper()

    print('===> Start Train')
    print("==================================================")

    solver_log = solver.get_current_log()

    NUM_EPOCH = int(opt['solver']['num_epochs'])
    start_epoch = solver_log['epoch']

    print("Method: %s || Scale: %d || Epoch Range: (%d ~ %d)" %
          (model_name, scale, start_epoch, NUM_EPOCH))

    for epoch in range(start_epoch, NUM_EPOCH + 1):
        print('\n===> Training Epoch: [%d/%d]...  Learning Rate: %f' %
              (epoch, NUM_EPOCH, solver.get_current_learning_rate()))

        # Initialization
        solver_log['epoch'] = epoch

        # Train model
        train_loss_list = []
        with tqdm(total=len(train_loader),
                  desc='Epoch: [%d/%d]' % (epoch, NUM_EPOCH),
                  miniters=1) as t:
            for iter, batch in enumerate(train_loader):
                solver.feed_data(batch)
                iter_loss = solver.train_step()
                batch_size = batch['LR'].size(0)
                train_loss_list.append(iter_loss * batch_size)
                t.set_postfix_str("Batch Loss: %.4f" % iter_loss)
                t.update()

        solver_log['records']['train_loss'].append(
            sum(train_loss_list) / len(train_set))
        solver_log['records']['lr'].append(solver.get_current_learning_rate())

        print('\nEpoch: [%d/%d]   Avg Train Loss: %.6f' %
              (epoch, NUM_EPOCH, sum(train_loss_list) / len(train_set)))

        print('===> Validating...')

        psnr_list = []
        ssim_list = []
        val_loss_list = []

        for iter, batch in enumerate(val_loader):
            solver.feed_data(batch)
            iter_loss = solver.test()
            val_loss_list.append(iter_loss)

            # calculate evaluation metrics
            visuals = solver.get_current_visual()
            # print(visuals['SR'].shape)
            psnr, ssim = util.calc_metrics(visuals['SR'],
                                           visuals['HR'],
                                           crop_border=scale)
            psnr_list.append(psnr)
            ssim_list.append(ssim)
            # if opt["save_image"]:
            #     solver.save_current_visual(epoch, iter)
        solver_log['records']['val_loss'].append(
            sum(val_loss_list) / len(val_loss_list))
        solver_log['records']['psnr'].append(sum(psnr_list) / len(psnr_list))
        solver_log['records']['ssim'].append(sum(ssim_list) / len(ssim_list))

        # record the best epoch
        epoch_is_best = False
        if solver_log['best_pred'] < (sum(psnr_list) / len(psnr_list)):
            solver_log['best_pred'] = (sum(psnr_list) / len(psnr_list))
            epoch_is_best = True
            solver_log['best_epoch'] = epoch

        print(
            "[%s] PSNR: %.2f   SSIM: %.4f   Loss: %.6f   Best PSNR: %.2f in Epoch: [%d]"
            %
            (val_set.name(), sum(psnr_list) / len(psnr_list), sum(ssim_list) /
             len(ssim_list), sum(val_loss_list) / len(val_loss_list),
             solver_log['best_pred'], solver_log['best_epoch']))

        solver.set_current_log(solver_log)
        solver.save_checkpoint(epoch, epoch_is_best)
        solver.save_current_log()

        # update lr
        solver.update_learning_rate(epoch)

    print('===> Finished !')
Пример #16
0
Файл: test.py Проект: yhu9/RCAN
 def getstats(self, sr, hr):
     sr = (sr * 255).clip(0, 255)
     hr = hr * 255
     psnr, ssim = util.calc_metrics(hr, sr, crop_border=self.upsize)
     return psnr, ssim
Пример #17
0
def main():
    torch.backends.cudnn.benchmark = True

    args = option.add_args()
    opt = option.parse(args.opt,
                       nblocks=args.nblocks,
                       nlayers=args.nlayers,
                       iterations=args.iterations,
                       trained_model=args.trained_path,
                       lr_path=args.lr_path)

    # fix random seed
    # seed_torch(opt['solver']['manual_seed'])

    # create train and val dataloader
    for phase, dataset_opt in sorted(opt['datasets'].items()):
        if phase == 'train':
            train_set = create_dataset(dataset_opt)
            train_loader = create_dataloader(train_set, dataset_opt)
            print('===> Train Dataset: %s   Number of images: [%d]' %
                  (train_set.name(), len(train_set)))
            if train_loader is None:
                raise ValueError("[Error] The training data does not exist")

        elif phase == 'val':
            val_set = create_dataset(dataset_opt)
            val_loader = create_dataloader(val_set, dataset_opt)
            print('===> Val Dataset: %s   Number of images: [%d]' %
                  (val_set.name(), len(val_set)))

        else:
            raise NotImplementedError(
                "[Error] Dataset phase [%s] in *.json is not recognized." %
                phase)

    solver = create_solver(opt)

    scale = opt['scale']
    model_name = opt['networks']['which_model'].upper()

    print('===> Start Train')
    print("==================================================")

    solver_log = solver.get_current_log()

    NUM_EPOCH = int(opt['solver']['num_epochs'])
    start_epoch = solver_log['epoch']

    print("Method: %s || Scale: %d || Epoch Range: (%d ~ %d)" %
          (model_name, scale, start_epoch, NUM_EPOCH))

    for epoch in range(start_epoch, NUM_EPOCH + 1):
        print('\n===> Training Epoch: [%d/%d]...  Learning Rate: %f' %
              (epoch, NUM_EPOCH, solver.get_current_learning_rate()))

        # Initialization
        solver_log['epoch'] = epoch

        # Train model
        train_loss_list = []
        with tqdm(total=len(train_loader),
                  desc='Epoch: [%d/%d]' % (epoch, NUM_EPOCH),
                  miniters=1) as t:
            for iter, batch in enumerate(train_loader):
                solver.feed_data(batch)
                iter_loss = solver.train_step()
                batch_size = batch['LR'].size(0)
                train_loss_list.append(iter_loss * batch_size)
                t.set_postfix_str("Batch Loss: %.4f" % iter_loss)
                t.update()

        solver_log['records']['train_loss'].append(
            sum(train_loss_list) / len(train_set))
        solver_log['records']['lr'].append(solver.get_current_learning_rate())

        print('\nEpoch: [%d/%d]   Avg Train Loss: %.6f' %
              (epoch, NUM_EPOCH, sum(train_loss_list) / len(train_set)))

        print('===> Validating...', )

        psnr_list = []
        ssim_list = []
        val_loss_list = []

        for iter, batch in enumerate(val_loader):
            solver.feed_data(batch)
            iter_loss = solver.test()
            val_loss_list.append(iter_loss)

            # calculate evaluation metrics
            visuals = solver.get_current_visual()
            psnr, ssim = util.calc_metrics(visuals['SR'],
                                           visuals['HR'],
                                           crop_border=scale)
            psnr_list.append(psnr)
            ssim_list.append(ssim)

            if opt["save_image"]:
                solver.save_current_visual(epoch, iter)

        solver_log['records']['val_loss'].append(
            sum(val_loss_list) / len(val_loss_list))
        solver_log['records']['psnr'].append(sum(psnr_list) / len(psnr_list))
        solver_log['records']['ssim'].append(sum(ssim_list) / len(ssim_list))

        # record the best epoch
        epoch_is_best = False
        if solver_log['best_pred'] < (sum(psnr_list) / len(psnr_list)):
            solver_log['best_pred'] = (sum(psnr_list) / len(psnr_list))
            epoch_is_best = True
            solver_log['best_epoch'] = epoch

        print(
            "[%s] PSNR: %.2f   SSIM: %.4f   Loss: %.6f   Best PSNR: %.2f in Epoch: [%d]"
            %
            (val_set.name(), sum(psnr_list) / len(psnr_list), sum(ssim_list) /
             len(ssim_list), sum(val_loss_list) / len(val_loss_list),
             solver_log['best_pred'], solver_log['best_epoch']))

        solver.set_current_log(solver_log)
        solver.save_checkpoint(epoch, epoch_is_best)
        solver.save_current_log()

        # update lr
        # solver.update_learning_rate()
        solver.scheduler.step()

    print('===> Finished !')
Пример #18
0
def SR(solver, opt, model_name):

    # dataset가져오기-많이 걸리면 0.002
    bm_names = []
    test_loaders = []
    for _, dataset_opt in sorted(opt['datasets'].items()):
        start = time.time()
        test_set = create_dataset(dataset_opt)
        test_loader = create_dataloader(test_set, dataset_opt)
        test_loaders.append(test_loader)
        print(
            '===> Test Dataset: [%s]   Number of images: [%d]  elapsed time: %.4f sec'
            % (test_set.name(), len(test_set), time.time() - start))
        bm_names.append(test_set.name())

    #Testset개수만큼 SR
    for bm, test_loader in zip(bm_names, test_loaders):
        print("Test set : [%s]" % bm)

        sr_list = []
        path_list = []

        total_psnr = []
        total_ssim = []
        total_time = []
        scale = 4

        need_HR = False if test_loader.dataset.__class__.__name__.find(
            'LRHR') < 0 else True

        for iter, batch in enumerate(test_loader):
            solver.feed_data(batch, need_HR=need_HR)

            # 시간측정
            t0 = time.time()
            solver.test()  #SR
            t1 = time.time()
            total_time.append((t1 - t0))

            visuals = solver.get_current_visual(need_HR=need_HR)
            sr_list.append(visuals['SR'])

            # calculate PSNR/SSIM metrics on Python
            if need_HR:
                psnr, ssim = util.calc_metrics(visuals['SR'],
                                               visuals['HR'],
                                               crop_border=scale)
                total_psnr.append(psnr)
                total_ssim.append(ssim)
                path_list.append(
                    os.path.basename(batch['HR_path'][0]).replace(
                        'HR', model_name))
                print(
                    "[%d/%d] %s || PSNR(dB)/SSIM: %.2f/%.4f || Timer: %.4f sec ."
                    % (iter + 1, len(test_loader),
                       os.path.basename(batch['LR_path'][0]), psnr, ssim,
                       (t1 - t0)))
            else:
                path_list.append(os.path.basename(batch['LR_path'][0]))
                print("[%d/%d] %s || Timer: %.4f sec ." %
                      (iter + 1, len(test_loader),
                       os.path.basename(batch['LR_path'][0]), (t1 - t0)))

        if need_HR:
            print("---- Average PSNR(dB) /SSIM /Speed(s) for [%s] ----" % bm)
            print("PSNR: %.2f      SSIM: %.4f      Speed: %.4f" %
                  (sum(total_psnr) / len(total_psnr), sum(total_ssim) /
                   len(total_ssim), sum(total_time) / len(total_time)))
        else:
            print("---- Average Speed(s) for [%s] is %.4f sec ----" %
                  (bm, sum(total_time) / len(total_time)))

        # save SR results for further evaluation on MATLAB
        if need_HR:
            save_img_path = os.path.join('./results/SR/' + degrad, model_name,
                                         bm, "x%d" % scale)
        else:
            save_img_path = os.path.join('./results/SR/' + bm, model_name,
                                         "x%d" % scale)

        if not os.path.exists(save_img_path): os.makedirs(save_img_path)
        for img, name in zip(sr_list, path_list):
            s = time.time()
            #matplotlib.image.save(os.path.join(save_img_path, name),img)
            cv2.imwrite(
                os.path.join(save_img_path, name),
                cv2.cvtColor(img,
                             cv2.COLOR_RGB2BGR))  # 0.609sec 평균 이미지 하나당 0.07sec
            print("NAME: %s DOWNLOAD TIME:%s\n" % (name, time.time() - s))
            #save(os.path.join(save_img_path, name),img) 5.3sec
            #Image.fromarray(img).save(os.path.join(save_img_path, name)) 2.4sec
            #imageio.imwrite(os.path.join(save_img_path, name), img) 9.8sec

        print(
            "===> Total Saving SR images of [%s]... Save Path: [%s] Time: %s\n"
            % (bm, save_img_path, time.time() - s))

    print("==================================================")
    print("===> Finished !!")
Пример #19
0
def main():
    args = option.add_args()
    opt = option.parse(args.opt,
                       nblocks=args.nblocks,
                       nlayers=args.nlayers,
                       iterations=args.iterations,
                       trained_model=args.trained_model,
                       lr_path=args.lr_path
                       )
    opt = option.dict_to_nonedict(opt)

    # initial configure
    scale = opt['scale']
    degrad = opt['degradation']
    network_opt = opt['networks']
    model_name = network_opt['which_model'].upper()
    if opt['self_ensemble']:
        model_name += 'plus'

    # create test dataloader
    bm_names =[]
    test_loaders = []
    for _, dataset_opt in sorted(opt['datasets'].items()):
        test_set = create_dataset(dataset_opt)
        test_loader = create_dataloader(test_set, dataset_opt)
        test_loaders.append(test_loader)
        print('===> Test Dataset: [%s]   Number of images: [%d]' % (test_set.name(), len(test_set)))
        bm_names.append(test_set.name())

    # create solver (and load model)
    solver = create_solver(opt)
    # Test phase
    print('===> Start Test')
    print("==================================================")
    print("Method: %s || Scale: %d || Degradation: %s"%(model_name, scale, degrad))

    # whether save the SR image?
    if opt['save_image']:
        para_save = Paralle_save_img()
        para_save.begin_background()
    # with para_save.begin_background() as para_save_imag

    for bm, test_loader in zip(bm_names, test_loaders):
        print("Test set : [%s]" % bm)

        total_psnr = []
        total_ssim = []
        total_time = []

        need_HR = False if test_loader.dataset.__class__.__name__.find('LRHR') < 0 else True

        if need_HR:
            save_img_path = os.path.join('./results/SR/' + degrad, model_name, bm, "x%d" % scale)
        else:
            save_img_path = os.path.join('./results/SR/' + bm, model_name, "x%d" % scale)

        if not os.path.exists(save_img_path):
            os.makedirs(save_img_path)

        for iter, batch in enumerate(test_loader):
            solver.feed_data(batch, need_HR=need_HR)

            # calculate forward time
            t0 = time.time()
            solver.test()
            t1 = time.time()
            total_time.append((t1 - t0))

            visuals = solver.get_current_visual(need_HR=need_HR)

            # calculate PSNR/SSIM metrics on Python
            if need_HR:
                psnr, ssim = util.calc_metrics(visuals['SR'], visuals['HR'], crop_border=scale)
                total_psnr.append(psnr)
                total_ssim.append(ssim)
                name = os.path.basename(batch['HR_path'][0]).replace('.', ('_x{}_' + model_name + '.').format(scale))

                print("[%d/%d] %s || PSNR(dB)/SSIM: %.2f/%.4f || Timer: %.4f sec ." % (iter + 1, len(test_loader),
                                                                                       os.path.basename(
                                                                                           batch['LR_path'][0]),
                                                                                       psnr, ssim,
                                                                                       (t1 - t0)))
            else:
                print("[%d/%d] %s || Timer: %.4f sec ." % (iter + 1, len(test_loader),
                                                           os.path.basename(batch['LR_path'][0]),
                                                           (t1 - t0)))
            if opt['save_image']:
                name = os.path.basename(batch['LR_path'][0]).replace('.', ('_x{}_' + model_name + '.').format(scale))
                para_save.put_image_path(filename=os.path.join(save_img_path, name), img=visuals['SR'])

        total_psnr, total_ssim = np.array(total_psnr), np.array(total_ssim)
        if need_HR:
            print("---- Average PSNR(dB) /SSIM /Speed(s) for [%s] ----" % bm)
            print("PSNR: %.2f(+/-%.2f)      SSIM: %.4f      Speed: %.4f" % (total_psnr.mean(), total_psnr.std(),
                                                                   total_ssim.mean(),
                                                                   sum(total_time) / len(total_time)))
        else:
            print("---- Average Speed(s) for [%s] is %.4f sec ----" % (bm,

                                                                       sum(total_time) / len(total_time)))
    if opt['save_image']:
        para_save.end_background()

    print("==================================================")
    print("===> Finished !")
Пример #20
0
    lrdata.sort()
    hrdata.sort()
    psnr_scores = []
    ssim_scores = []

    #RUN ON SET 5 DATASET
    for lrname,hrname in zip(lrdata,hrdata):

        img = imageio.imread(os.path.join(LRDIR,lrname))
        hr = imageio.imread(os.path.join(HRDIR,hrname))

        img = torch.FloatTensor(img).to(device)
        img = img.permute((2,0,1)).unsqueeze(0)

        #MAKE THE INFERENCE
        model.eval()
        with torch.no_grad():
            sr = model(img)
            sr = sr.squeeze(0).permute(1,2,0).data.cpu().numpy()

        # hr = hr
        psnr,ssim = util.calc_metrics(hr,sr,4)
        psnr_scores.append(psnr)
        ssim_scores.append(ssim)
        print('psnr score: {:.4f}   | {}'.format(psnr,lrdata))

    print('mean and ssim: ',np.mean(psnr_scores),np.mean(ssim_scores))
    checkpoint.done()


Пример #21
0
def main():
    # Test phase
    print('===> Start Test')

    for bm, test_loader in zip(bm_names, test_loaders):
        print("Test set : [%s]"%bm)

        sr_list = []
        path_list = []

        total_psnr = []
        total_ssim = []
        total_time = []

        need_HR = False if test_loader.dataset.__class__.__name__.find('LRHR') < 0 else True

        for iter, batch in enumerate(test_loader):
            solver.feed_data(batch, need_HR=need_HR)

            # calculate forward time
            t0 = time.time()
            solver.test()
            t1 = time.time()
            total_time.append((t1 - t0))

            visuals = solver.get_current_visual(need_HR=need_HR)
            sr_list.append(visuals['SR'])

            # calculate PSNR/SSIM metrics on Python
            if need_HR:
                psnr, ssim = util.calc_metrics(visuals['SR'], visuals['HR'], crop_border=scale)
                total_psnr.append(psnr)
                total_ssim.append(ssim)
                path_list.append(os.path.basename(batch['HR_path'][0]).replace('HR', model_name))
                print("[%d/%d] %s || PSNR(dB)/SSIM: %.2f/%.4f || Timer: %.4f sec ." % (iter+1, len(test_loader),
                                                                                       os.path.basename(batch['LR_path'][0]),
                                                                                       psnr, ssim,
                                                                                       (t1 - t0)))
            else:
                path_list.append(os.path.basename(batch['LR_path'][0]))
                print("[%d/%d] %s || Timer: %.4f sec ." % (iter + 1, len(test_loader),
                                                           os.path.basename(batch['LR_path'][0]),
                                                           (t1 - t0)))

        if need_HR:
            print("---- Average PSNR(dB) /SSIM /Speed(s) for [%s] ----" % bm)
            print("PSNR: %.2f      SSIM: %.4f      Speed: %.4f" % (sum(total_psnr)/len(total_psnr),
                                                                  sum(total_ssim)/len(total_ssim),
                                                                  sum(total_time)/len(total_time)))
        else:
            print("---- Average Speed(s) for [%s] is %.4f sec ----" % (bm,
                                                                      sum(total_time)/len(total_time)))

        # save SR results for further evaluation on MATLAB
        if need_HR:
            save_img_path = os.path.join('./results/SR/'+degrad, model_name, bm, "x%d"%scale)
        else:
            save_img_path = os.path.join('./results/SR/'+bm, model_name, "x%d"%scale)

        print("===> Saving SR images of [%s]... Save Path: [%s]\n" % (bm, save_img_path))

        if not os.path.exists(save_img_path): os.makedirs(save_img_path)
        for img, name in zip(sr_list, path_list):
            imageio.imwrite(os.path.join(save_img_path, name), img)

    print("==================================================")
    print("===> Finished !")