| import sys | |
| import tqdm | |
| import torch | |
| import argparse | |
| import numpy as np | |
| from omegaconf import OmegaConf | |
| sys.path.append('.') | |
| from utils.build_utils import build_from_cfg | |
| from datasets.gopro_datasets import GoPro_Test_Dataset | |
| from metrics.psnr_ssim import calculate_psnr, calculate_ssim | |
| parser = argparse.ArgumentParser( | |
| prog = 'AMT', | |
| description = 'GOPRO evaluation', | |
| ) | |
| parser.add_argument('-c', '--config', default='cfgs/AMT-S_gopro.yaml') | |
| parser.add_argument('-p', '--ckpt', default='pretrained/gopro_amt-s.pth',) | |
| parser.add_argument('-r', '--root', default='data/GOPRO',) | |
| args = parser.parse_args() | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| cfg_path = args.config | |
| ckpt_path = args.ckpt | |
| root = args.root | |
| network_cfg = OmegaConf.load(cfg_path).network | |
| network_name = network_cfg.name | |
| model = build_from_cfg(network_cfg) | |
| ckpt = torch.load(ckpt_path) | |
| model.load_state_dict(ckpt['state_dict']) | |
| model = model.to(device) | |
| model.eval() | |
| dataset = GoPro_Test_Dataset(dataset_dir=root) | |
| psnr_list = [] | |
| ssim_list = [] | |
| pbar = tqdm.tqdm(dataset, total=len(dataset)) | |
| for data in pbar: | |
| input_dict = {} | |
| for k, v in data.items(): | |
| input_dict[k] = v.to(device).unsqueeze(0) | |
| with torch.no_grad(): | |
| imgt_pred = model(**input_dict)['imgt_pred'] | |
| psnr = calculate_psnr(imgt_pred, input_dict['imgt']) | |
| ssim = calculate_ssim(imgt_pred, input_dict['imgt']) | |
| psnr_list.append(psnr) | |
| ssim_list.append(ssim) | |
| avg_psnr = np.mean(psnr_list) | |
| avg_ssim = np.mean(ssim_list) | |
| desc_str = f'[{network_name}/GOPRO] psnr: {avg_psnr:.02f}, ssim: {avg_ssim:.04f}' | |
| pbar.set_description_str(desc_str) | |