Spaces:
Running
Running
| import cv2 | |
| import numpy as np | |
| import math | |
| # from skimage.metrics import structural_similarity as ssim | |
| from skimage.measure import compare_ssim | |
| from scipy.misc import imread | |
| from glob import glob | |
| import argparse | |
| parser = argparse.ArgumentParser(description="evaluation codes") | |
| parser.add_argument("--path", type=str, help="Path to evaluate images.") | |
| args = parser.parse_args() | |
| def psnr(img1, img2): | |
| mse = np.mean((img1 / 255.0 - img2 / 255.0) ** 2) | |
| if mse < 1.0e-10: | |
| return 100 | |
| PIXEL_MAX = 1 | |
| return 20 * math.log10(PIXEL_MAX / math.sqrt(mse)) | |
| def psnr_raw(img1, img2): | |
| mse = np.mean((img1 - img2) ** 2) | |
| if mse < 1.0e-10: | |
| return 100 | |
| PIXEL_MAX = 1 | |
| return 20 * math.log10(PIXEL_MAX / math.sqrt(mse)) | |
| def my_ssim(img1, img2): | |
| return compare_ssim( | |
| img1, img2, data_range=img1.max() - img1.min(), multichannel=True | |
| ) | |
| def quan_eval(path, suffix="jpg"): | |
| # path: /disk2/yazhou/projects/IISP/exps/test_final_unet_globalEDV2/ | |
| # ours | |
| gt_imgs = sorted(glob(path + "tar*.%s" % suffix)) | |
| pred_imgs = sorted(glob(path + "pred*.%s" % suffix)) | |
| # with open(split_path + "test_gt.txt", 'r') as f_gt, open(split_path+"test_rgb.txt","r") as f_rgb: | |
| # gt_imgs = [line.rstrip() for line in f_gt.readlines()] | |
| # pred_imgs = [line.rstrip() for line in f_rgb.readlines()] | |
| assert len(gt_imgs) == len(pred_imgs) | |
| psnr_avg = 0.0 | |
| ssim_avg = 0.0 | |
| for i in range(len(gt_imgs)): | |
| gt = imread(gt_imgs[i]) | |
| pred = imread(pred_imgs[i]) | |
| psnr_temp = psnr(gt, pred) | |
| psnr_avg += psnr_temp | |
| ssim_temp = my_ssim(gt, pred) | |
| ssim_avg += ssim_temp | |
| print("psnr: ", psnr_temp) | |
| print("ssim: ", ssim_temp) | |
| psnr_avg /= float(len(gt_imgs)) | |
| ssim_avg /= float(len(gt_imgs)) | |
| print("psnr_avg: ", psnr_avg) | |
| print("ssim_avg: ", ssim_avg) | |
| return psnr_avg, ssim_avg | |
| def mse(gt, pred): | |
| return np.mean((gt - pred) ** 2) | |
| def mse_raw(path, suffix="npy"): | |
| gt_imgs = sorted(glob(path + "raw_tar*.%s" % suffix)) | |
| pred_imgs = sorted(glob(path + "raw_pred*.%s" % suffix)) | |
| # with open(split_path + "test_gt.txt", 'r') as f_gt, open(split_path+"test_rgb.txt","r") as f_rgb: | |
| # gt_imgs = [line.rstrip() for line in f_gt.readlines()] | |
| # pred_imgs = [line.rstrip() for line in f_rgb.readlines()] | |
| assert len(gt_imgs) == len(pred_imgs) | |
| mse_avg = 0.0 | |
| psnr_avg = 0.0 | |
| for i in range(len(gt_imgs)): | |
| gt = np.load(gt_imgs[i]) | |
| pred = np.load(pred_imgs[i]) | |
| mse_temp = mse(gt, pred) | |
| mse_avg += mse_temp | |
| psnr_temp = psnr_raw(gt, pred) | |
| psnr_avg += psnr_temp | |
| print("mse: ", mse_temp) | |
| print("psnr: ", psnr_temp) | |
| mse_avg /= float(len(gt_imgs)) | |
| psnr_avg /= float(len(gt_imgs)) | |
| print("mse_avg: ", mse_avg) | |
| print("psnr_avg: ", psnr_avg) | |
| return mse_avg, psnr_avg | |
| test_full = False | |
| # if test_full: | |
| # psnr_avg, ssim_avg = quan_eval(ROOT_PATH+"%s/vis_%s_full/"%(args.task, args.ckpt), "jpeg") | |
| # mse_avg, psnr_avg_raw = mse_raw(ROOT_PATH+"%s/vis_%s_full/"%(args.task, args.ckpt)) | |
| # else: | |
| psnr_avg, ssim_avg = quan_eval(args.path, "jpg") | |
| mse_avg, psnr_avg_raw = mse_raw(args.path) | |
| print( | |
| "pnsr: {}, ssim: {}, mse: {}, psnr raw: {}".format( | |
| psnr_avg, ssim_avg, mse_avg, psnr_avg_raw | |
| ) | |
| ) | |