Spaces:
Sleeping
Sleeping
| from argparse import ArgumentParser | |
| import os | |
| import json | |
| import sys | |
| from tqdm import tqdm | |
| import numpy as np | |
| import torch | |
| from torch.utils.data import DataLoader | |
| import torchvision.transforms as transforms | |
| sys.path.append(".") | |
| sys.path.append("..") | |
| from criteria.lpips.lpips import LPIPS | |
| from datasets.gt_res_dataset import GTResDataset | |
| def parse_args(): | |
| parser = ArgumentParser(add_help=False) | |
| parser.add_argument('--mode', type=str, default='lpips', choices=['lpips', 'l2']) | |
| parser.add_argument('--data_path', type=str, default='results') | |
| parser.add_argument('--gt_path', type=str, default='gt_images') | |
| parser.add_argument('--workers', type=int, default=4) | |
| parser.add_argument('--batch_size', type=int, default=4) | |
| args = parser.parse_args() | |
| return args | |
| def run(args): | |
| transform = transforms.Compose([transforms.Resize((256, 256)), | |
| transforms.ToTensor(), | |
| transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]) | |
| print('Loading dataset') | |
| dataset = GTResDataset(root_path=args.data_path, | |
| gt_dir=args.gt_path, | |
| transform=transform) | |
| dataloader = DataLoader(dataset, | |
| batch_size=args.batch_size, | |
| shuffle=False, | |
| num_workers=int(args.workers), | |
| drop_last=True) | |
| if args.mode == 'lpips': | |
| loss_func = LPIPS(net_type='alex') | |
| elif args.mode == 'l2': | |
| loss_func = torch.nn.MSELoss() | |
| else: | |
| raise Exception('Not a valid mode!') | |
| loss_func.cuda() | |
| global_i = 0 | |
| scores_dict = {} | |
| all_scores = [] | |
| for result_batch, gt_batch in tqdm(dataloader): | |
| for i in range(args.batch_size): | |
| loss = float(loss_func(result_batch[i:i+1].cuda(), gt_batch[i:i+1].cuda())) | |
| all_scores.append(loss) | |
| im_path = dataset.pairs[global_i][0] | |
| scores_dict[os.path.basename(im_path)] = loss | |
| global_i += 1 | |
| all_scores = list(scores_dict.values()) | |
| mean = np.mean(all_scores) | |
| std = np.std(all_scores) | |
| result_str = 'Average loss is {:.2f}+-{:.2f}'.format(mean, std) | |
| print('Finished with ', args.data_path) | |
| print(result_str) | |
| out_path = os.path.join(os.path.dirname(args.data_path), 'inference_metrics') | |
| if not os.path.exists(out_path): | |
| os.makedirs(out_path) | |
| with open(os.path.join(out_path, 'stat_{}.txt'.format(args.mode)), 'w') as f: | |
| f.write(result_str) | |
| with open(os.path.join(out_path, 'scores_{}.json'.format(args.mode)), 'w') as f: | |
| json.dump(scores_dict, f) | |
| if __name__ == '__main__': | |
| args = parse_args() | |
| run(args) | |