Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

NhanNT: Update some code to calculate task real_sr #163

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
*.png
14 changes: 12 additions & 2 deletions main_test_swinir.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import os
import torch
import requests
import pandas as pd

from models.network_swinir import SwinIR as net
from utils import util_calculate_psnr_ssim as util
Expand Down Expand Up @@ -56,6 +57,7 @@ def main():
test_results['psnrb'] = []
test_results['psnrb_y'] = []
psnr, ssim, psnr_y, ssim_y, psnrb, psnrb_y = 0, 0, 0, 0, 0, 0
results = list()

for idx, path in enumerate(sorted(glob.glob(os.path.join(folder, '*')))):
# read image
Expand Down Expand Up @@ -105,9 +107,15 @@ def main():
print('Testing {:d} {:20s} - PSNR: {:.2f} dB; SSIM: {:.4f}; PSNRB: {:.2f} dB;'
'PSNR_Y: {:.2f} dB; SSIM_Y: {:.4f}; PSNRB_Y: {:.2f} dB.'.
format(idx, imgname, psnr, ssim, psnrb, psnr_y, ssim_y, psnrb_y))
results.append([imgname, psnr, ssim, psnr_y, ssim_y])
else:
print('Testing {:d} {:20s}'.format(idx, imgname))
columns = ['Image Name', 'PSNR', 'SSIM', 'PSNR_Y', 'SSIM_Y']
df = pd.DataFrame(results, columns=columns)

csv_path = os.path.join(save_dir, 'results.csv')
df.to_csv(csv_path, index=False)
print(f'Results saved to {csv_path}')
# summarize psnr/ssim
if img_gt is not None:
ave_psnr = sum(test_results['psnr']) / len(test_results['psnr'])
Expand Down Expand Up @@ -228,16 +236,18 @@ def setup(args):

def get_image_pair(args, path):
(imgname, imgext) = os.path.splitext(os.path.basename(path))
print(imgname)
print(imgext)

# 001 classical image sr/ 002 lightweight image sr (load lq-gt image pairs)
if args.task in ['classical_sr', 'lightweight_sr']:
img_gt = cv2.imread(path, cv2.IMREAD_COLOR).astype(np.float32) / 255.
img_lq = cv2.imread(f'{args.folder_lq}/{imgname}x{args.scale}{imgext}', cv2.IMREAD_COLOR).astype(
np.float32) / 255.

# 003 real-world image sr (load lq image only)
# 003 real-world image sr (load lq image and gt image)
elif args.task in ['real_sr']:
img_gt = None
img_gt = cv2.imread(f'{args.folder_gt}/{imgname}{imgext}', cv2.IMREAD_COLOR).astype(np.float32) / 255.
img_lq = cv2.imread(path, cv2.IMREAD_COLOR).astype(np.float32) / 255.

# 004 grayscale image denoising (load gt image and generate lq image on-the-fly)
Expand Down