Spaces:
Sleeping
Sleeping
File size: 2,904 Bytes
f4c3c2b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
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)
parser.add_argument('--is_cars', action='store_true')
args = parser.parse_args()
return args
def run(args):
resize_dims = (256, 256)
if args.is_cars:
resize_dims = (192, 256)
transform = transforms.Compose([transforms.Resize(resize_dims),
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)
|