forked from leftthomas/SRGAN
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_stats.py
53 lines (41 loc) · 2.08 KB
/
test_stats.py
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
import os
import numpy as np
from skimage import io
from skimage.metrics import structural_similarity as ssim, peak_signal_noise_ratio as psnr
def calculate_metrics(image1_path, image2_path):
img1 = io.imread(image1_path)
img2 = io.imread(image2_path)
if img1.shape != img2.shape:
raise ValueError("Images must have the same dimensions.")
ssim_value = ssim(img1, img2, channel_axis=-1, data_range=img2.max() - img2.min())
psnr_value = psnr(img1, img2, data_range=img2.max() - img2.min())
return ssim_value, psnr_value
def process_folders(folder_path1, folder_path2):
directories1 = [d for d in os.listdir(folder_path1) if os.path.isdir(os.path.join(folder_path1, d))]
directories2 = [d for d in os.listdir(folder_path2) if os.path.isdir(os.path.join(folder_path2, d))]
common_directories = set(directories1).intersection(directories2)
all_ssim = []
all_psnr = []
for directory in common_directories:
dir1 = os.path.join(folder_path1, directory)
dir2 = os.path.join(folder_path2, directory)
images1 = {f: os.path.join(dir1, f) for f in os.listdir(dir1) if f.endswith(('.png', '.jpg'))}
images2 = {f: os.path.join(dir2, f) for f in os.listdir(dir2) if f.endswith(('.png', '.jpg'))}
common_images = set(images1.keys()).intersection(images2.keys())
for image in common_images:
ssim_value, psnr_value = calculate_metrics(images1[image], images2[image])
all_ssim.append(ssim_value)
all_psnr.append(psnr_value)
mean_ssim = np.mean(all_ssim) if all_ssim else 0
mean_psnr = np.mean(all_psnr) if all_psnr else 0
return mean_ssim, mean_psnr
# Example usage
folder_path1 = 'data/REDS/val_REDS4_sharp_bicubic/X4'
folder_path2 = 'data/REDS/base_pred_val'
# folder_path1 = 'data/VID4/GT'
# folder_path2 = 'data/VID4/BDx4/base_pred'
folder_path1 = 'data/REDS/val_REDS4_sharp'
folder_path2 = 'data/REDS/base_highres_pred_val'
mean_ssim, mean_psnr = process_folders(folder_path1, folder_path2)
print(f"Mean SSIM: {mean_ssim}")
print(f"Mean PSNR: {mean_psnr}")