def test(img_path, img_path2, crop_border, test_y_channel=False):
img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
img2 = cv2.imread(img_path2, cv2.IMREAD_UNCHANGED)
# --------------------- Numpy ---------------------
psnr = calculate_psnr(img,
img2,
crop_border=crop_border,
input_order='HWC',
test_y_channel=test_y_channel)
ssim = calculate_ssim(img,
img2,
crop_border=crop_border,
input_order='HWC',
test_y_channel=test_y_channel)
print(f'\tNumpy\tPSNR: {psnr:.6f} dB, \tSSIM: {ssim:.6f}')
# --------------------- PyTorch (CPU) ---------------------
img = img2tensor(img / 255., bgr2rgb=True, float32=True).unsqueeze_(0)
img2 = img2tensor(img2 / 255., bgr2rgb=True, float32=True).unsqueeze_(0)
psnr_pth = calculate_psnr_pt(img,
img2,
crop_border=crop_border,
test_y_channel=test_y_channel)
ssim_pth = calculate_ssim_pt(img,
img2,
crop_border=crop_border,
test_y_channel=test_y_channel)
print(
f'\tTensor (CPU) \tPSNR: {psnr_pth[0]:.6f} dB, \tSSIM: {ssim_pth[0]:.6f}'
)
# --------------------- PyTorch (GPU) ---------------------
img = img.cuda()
img2 = img2.cuda()
psnr_pth = calculate_psnr_pt(img,
img2,
crop_border=crop_border,
test_y_channel=test_y_channel)
ssim_pth = calculate_ssim_pt(img,
img2,
crop_border=crop_border,
test_y_channel=test_y_channel)
print(
f'\tTensor (GPU) \tPSNR: {psnr_pth[0]:.6f} dB, \tSSIM: {ssim_pth[0]:.6f}'
)
psnr_pth = calculate_psnr_pt(torch.repeat_interleave(img, 2, dim=0),
torch.repeat_interleave(img2, 2, dim=0),
crop_border=crop_border,
test_y_channel=test_y_channel)
ssim_pth = calculate_ssim_pt(torch.repeat_interleave(img, 2, dim=0),
torch.repeat_interleave(img2, 2, dim=0),
crop_border=crop_border,
test_y_channel=test_y_channel)
print(
f'\tTensor (GPU batch) \tPSNR: {psnr_pth[0]:.6f}, {psnr_pth[1]:.6f} dB,'
f'\tSSIM: {ssim_pth[0]:.6f}, {ssim_pth[1]:.6f}')
def main():
"""Calculate PSNR and SSIM for images.
Configurations:
folder_gt (str): Path to gt (Ground-Truth).
folder_restored (str): Path to restored images.
crop_border (int): Crop border for each side.
suffix (str): Suffix for restored images.
test_y_channel (bool): If True, test Y channel (In MatLab YCbCr format)
If False, test RGB channels.
"""
# Configurations
# -------------------------------------------------------------------------
folder_gt = 'datasets/val_set14/Set14'
folder_restored = 'results/exp/visualization/val_set14'
crop_border = 4
suffix = '_expname'
test_y_channel = False
# -------------------------------------------------------------------------
psnr_all = []
ssim_all = []
img_list = sorted(mmcv.scandir(folder_gt, recursive=True))
if test_y_channel:
print('Testing Y channel.')
else:
print('Testing RGB channels.')
for i, img_path in enumerate(img_list):
basename, ext = osp.splitext(osp.basename(img_path))
img_gt = mmcv.imread(
osp.join(folder_gt, img_path), flag='unchanged').astype(
np.float32) / 255.
img_restored = mmcv.imread(
osp.join(folder_restored, basename + suffix + ext),
flag='unchanged').astype(np.float32) / 255.
if test_y_channel and img_gt.ndim == 3 and img_gt.shape[2] == 3:
img_gt = mmcv.bgr2ycbcr(img_gt, y_only=True)
img_restored = mmcv.bgr2ycbcr(img_restored, y_only=True)
# calculate PSNR and SSIM
psnr = calculate_psnr(
img_gt * 255,
img_restored * 255,
crop_border=crop_border,
input_order='HWC')
ssim = calculate_ssim(
img_gt * 255,
img_restored * 255,
crop_border=crop_border,
input_order='HWC')
print(f'{i+1:3d}: {basename:25}. \tPSNR: {psnr:.6f} dB, '
f'\tSSIM: {ssim:.6f}')
psnr_all.append(psnr)
ssim_all.append(ssim)
print(f'Average: PSNR: {sum(psnr_all) / len(psnr_all):.6f} dB, '
f'SSIM: {sum(ssim_all) / len(ssim_all):.6f}')
def main(args):
"""Calculate PSNR and SSIM for images.
"""
psnr_all = []
ssim_all = []
img_list_gt = sorted(list(scandir(args.gt, recursive=True, full_path=True)))
img_list_restored = sorted(list(scandir(args.restored, recursive=True, full_path=True)))
if args.test_y_channel:
print('Testing Y channel.')
else:
print('Testing RGB channels.')
for i, img_path in enumerate(img_list_gt):
basename, ext = osp.splitext(osp.basename(img_path))
img_gt = cv2.imread(img_path, cv2.IMREAD_UNCHANGED).astype(np.float32) / 255.
if args.suffix == '':
img_path_restored = img_list_restored[i]
else:
img_path_restored = osp.join(args.restored, basename + args.suffix + ext)
img_restored = cv2.imread(img_path_restored, cv2.IMREAD_UNCHANGED).astype(np.float32) / 255.
if args.correct_mean_var:
mean_l = []
std_l = []
for j in range(3):
mean_l.append(np.mean(img_gt[:, :, j]))
std_l.append(np.std(img_gt[:, :, j]))
for j in range(3):
# correct twice
mean = np.mean(img_restored[:, :, j])
img_restored[:, :, j] = img_restored[:, :, j] - mean + mean_l[j]
std = np.std(img_restored[:, :, j])
img_restored[:, :, j] = img_restored[:, :, j] / std * std_l[j]
mean = np.mean(img_restored[:, :, j])
img_restored[:, :, j] = img_restored[:, :, j] - mean + mean_l[j]
std = np.std(img_restored[:, :, j])
img_restored[:, :, j] = img_restored[:, :, j] / std * std_l[j]
if args.test_y_channel and img_gt.ndim == 3 and img_gt.shape[2] == 3:
img_gt = bgr2ycbcr(img_gt, y_only=True)
img_restored = bgr2ycbcr(img_restored, y_only=True)
# calculate PSNR and SSIM
psnr = calculate_psnr(img_gt * 255, img_restored * 255, crop_border=args.crop_border, input_order='HWC')
ssim = calculate_ssim(img_gt * 255, img_restored * 255, crop_border=args.crop_border, input_order='HWC')
print(f'{i+1:3d}: {basename:25}. \tPSNR: {psnr:.6f} dB, \tSSIM: {ssim:.6f}')
psnr_all.append(psnr)
ssim_all.append(ssim)
print(args.gt)
print(args.restored)
print(f'Average: PSNR: {sum(psnr_all) / len(psnr_all):.6f} dB, SSIM: {sum(ssim_all) / len(ssim_all):.6f}')
def main():
"""Calculate PSNR and SSIM for images.
Configurations:
folder_gt (str): Path to gt (Ground-Truth).
folder_restored (str): Path to restored images.
crop_border (int): Crop border for each side.
suffix (str): Suffix for restored images.
test_y_channel (bool): If True, test Y channel (In MatLab YCbCr format)
If False, test RGB channels.
"""
# Configurations
# -------------------------------------------------------------------------
folder_gt = 'datasets/val_set14/Set14'
folder_restored = 'results/exp/visualization/val_set14'
crop_border = 4
suffix = '_expname'
test_y_channel = False
correct_mean_var = False
# -------------------------------------------------------------------------
psnr_all = []
ssim_all = []
img_list = sorted(scandir(folder_gt, recursive=True, full_path=True))
if test_y_channel:
print('Testing Y channel.')
else:
print('Testing RGB channels.')
for i, img_path in enumerate(img_list):
basename, ext = osp.splitext(osp.basename(img_path))
img_gt = cv2.imread(img_path, cv2.IMREAD_UNCHANGED).astype(
np.float32) / 255.
img_restored = cv2.imread(
osp.join(folder_restored, basename + suffix + ext),
cv2.IMREAD_UNCHANGED).astype(np.float32) / 255.
if correct_mean_var:
mean_l = []
std_l = []
for j in range(3):
mean_l.append(np.mean(img_gt[:, :, j]))
std_l.append(np.std(img_gt[:, :, j]))
for j in range(3):
# correct twice
mean = np.mean(img_restored[:, :, j])
img_restored[:, :,
j] = img_restored[:, :, j] - mean + mean_l[j]
std = np.std(img_restored[:, :, j])
img_restored[:, :, j] = img_restored[:, :, j] / std * std_l[j]
mean = np.mean(img_restored[:, :, j])
img_restored[:, :,
j] = img_restored[:, :, j] - mean + mean_l[j]
std = np.std(img_restored[:, :, j])
img_restored[:, :, j] = img_restored[:, :, j] / std * std_l[j]
if test_y_channel and img_gt.ndim == 3 and img_gt.shape[2] == 3:
img_gt = bgr2ycbcr(img_gt, y_only=True)
img_restored = bgr2ycbcr(img_restored, y_only=True)
# calculate PSNR and SSIM
psnr = calculate_psnr(img_gt * 255,
img_restored * 255,
crop_border=crop_border,
input_order='HWC')
ssim = calculate_ssim(img_gt * 255,
img_restored * 255,
crop_border=crop_border,
input_order='HWC')
print(f'{i+1:3d}: {basename:25}. \tPSNR: {psnr:.6f} dB, '
f'\tSSIM: {ssim:.6f}')
psnr_all.append(psnr)
ssim_all.append(ssim)
print(folder_gt)
print(folder_restored)
print(f'Average: PSNR: {sum(psnr_all) / len(psnr_all):.6f} dB, '
f'SSIM: {sum(ssim_all) / len(ssim_all):.6f}')