File size: 3,146 Bytes
36d9761 |
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 |
import pyiqa
import os
import argparse
from pathlib import Path
import torch
from utils import util_image
import tqdm
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
print(pyiqa.list_models())
def evaluate(in_path, ref_path, ntest):
metric_dict = {}
metric_dict["clipiqa"] = pyiqa.create_metric('clipiqa').to(device)
metric_dict["musiq"] = pyiqa.create_metric('musiq').to(device)
metric_dict["niqe"] = pyiqa.create_metric('niqe').to(device)
metric_dict["maniqa"] = pyiqa.create_metric('maniqa').to(device)
metric_paired_dict = {}
in_path = Path(in_path) if not isinstance(in_path, Path) else in_path
assert in_path.is_dir()
ref_path_list = None
if ref_path is not None:
ref_path = Path(ref_path) if not isinstance(ref_path, Path) else ref_path
ref_path_list = sorted([x for x in ref_path.glob("*.[jpJP][pnPN]*[gG]")])
if ntest is not None: ref_path_list = ref_path_list[:ntest]
metric_paired_dict["psnr"]=pyiqa.create_metric('psnr', test_y_channel=True, color_space='ycbcr').to(device)
metric_paired_dict["lpips"]=pyiqa.create_metric('lpips').to(device)
metric_paired_dict["dists"]=pyiqa.create_metric('dists').to(device)
metric_paired_dict["ssim"]=pyiqa.create_metric('ssim', test_y_channel=True, color_space='ycbcr' ).to(device)
lr_path_list = sorted([x for x in in_path.glob("*.[jpJP][pnPN]*[gG]")])
if ntest is not None: lr_path_list = lr_path_list[:ntest]
print(f'Find {len(lr_path_list)} images in {in_path}')
result = {}
for i in tqdm.tqdm(range(len(lr_path_list))):
_in_path = lr_path_list[i]
_ref_path = ref_path_list[i] if ref_path_list is not None else None
im_in = util_image.imread(_in_path, chn='rgb', dtype='float32') # h x w x c
im_in_tensor = util_image.img2tensor(im_in).cuda() # 1 x c x h x w
for key, metric in metric_dict.items():
with torch.cuda.amp.autocast():
result[key] = result.get(key, 0) + metric(im_in_tensor).item()
if ref_path is not None:
im_ref = util_image.imread(_ref_path, chn='rgb', dtype='float32') # h x w x c
im_ref_tensor = util_image.img2tensor(im_ref).cuda()
for key, metric in metric_paired_dict.items():
result[key] = result.get(key, 0) + metric(im_in_tensor, im_ref_tensor).item()
if ref_path is not None:
fid_metric = pyiqa.create_metric('fid')
result['fid'] = fid_metric(in_path, ref_path)
for key, res in result.items():
if key == 'fid':
print(f"{key}: {res:.2f}")
else:
print(f"{key}: {res/len(lr_path_list):.5f}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('-i',"--in_path", type=str, required=True)
parser.add_argument("-r", "--ref_path", type=str, default=None)
parser.add_argument("--ntest", type=int, default=None)
args = parser.parse_args()
evaluate(args.in_path, args.ref_path, args.ntest)
|