import cog
import tempfile
from pathlib import Path
import argparse
import shutil
import os
import cv2
import glob
import torch
from collections import OrderedDict
import numpy as np
from main_test_swinir import define_model, setup, get_image_pair
class Predictor(cog.Predictor):
def setup(self):
model_dir = 'experiments/pretrained_models'
self.model_zoo = {
'real_sr': {
4: os.path.join(model_dir, '003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pth')
},
'gray_dn': {
15: os.path.join(model_dir, '004_grayDN_DFWB_s128w8_SwinIR-M_noise15.pth'),
25: os.path.join(model_dir, '004_grayDN_DFWB_s128w8_SwinIR-M_noise25.pth'),
50: os.path.join(model_dir, '004_grayDN_DFWB_s128w8_SwinIR-M_noise50.pth')
},
'color_dn': {
15: os.path.join(model_dir, '005_colorDN_DFWB_s128w8_SwinIR-M_noise15.pth'),
25: os.path.join(model_dir, '005_colorDN_DFWB_s128w8_SwinIR-M_noise25.pth'),
50: os.path.join(model_dir, '005_colorDN_DFWB_s128w8_SwinIR-M_noise50.pth')
},
'jpeg_car': {
10: os.path.join(model_dir, '006_CAR_DFWB_s126w7_SwinIR-M_jpeg10.pth'),
20: os.path.join(model_dir, '006_CAR_DFWB_s126w7_SwinIR-M_jpeg20.pth'),
30: os.path.join(model_dir, '006_CAR_DFWB_s126w7_SwinIR-M_jpeg30.pth'),
40: os.path.join(model_dir, '006_CAR_DFWB_s126w7_SwinIR-M_jpeg40.pth')
}
}
parser = argparse.ArgumentParser()
parser.add_argument('--task', type=str, default='real_sr', help='classical_sr, lightweight_sr, real_sr, '
'gray_dn, color_dn, jpeg_car')
parser.add_argument('--scale', type=int, default=1, help='scale factor: 1, 2, 3, 4, 8') # 1 for dn and jpeg car
parser.add_argument('--noise', type=int, default=15, help='noise level: 15, 25, 50')
parser.add_argument('--jpeg', type=int, default=40, help='scale factor: 10, 20, 30, 40')
parser.add_argument('--training_patch_size', type=int, default=128, help='patch size used in training SwinIR. '
'Just used to differentiate two different settings in Table 2 of the paper. '
'Images are NOT tested patch by patch.')
parser.add_argument('--large_model', action='store_true',
help='use large model, only provided for real image sr')
parser.add_argument('--model_path', type=str,
default=self.model_zoo['real_sr'][4])
parser.add_argument('--folder_lq', type=str, default=None, help='input low-quality test image folder')
parser.add_argument('--folder_gt', type=str, default=None, help='input ground-truth test image folder')
self.args = parser.parse_args('')
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.tasks = {
'Real-World Image Super-Resolution': 'real_sr',
'Grayscale Image Denoising': 'gray_dn',
'Color Image Denoising': 'color_dn',
'JPEG Compression Artifact Reduction': 'jpeg_car'
}
@cog.input("image", type=Path, help="input image")
@cog.input("task_type", type=str, default='Real-World Image Super-Resolution',
options=['Real-World Image Super-Resolution', 'Grayscale Image Denoising', 'Color Image Denoising',
'JPEG Compression Artifact Reduction'],
help="image restoration task type")
@cog.input("noise", type=int, default=15, options=[15, 25, 50],
help='noise level, activated for Grayscale Image Denoising and Color Image Denoising. '
'Leave it as default or arbitrary if other tasks are selected')
@cog.input("jpeg", type=int, default=40, options=[10, 20, 30, 40],
help='scale factor, activated for JPEG Compression Artifact Reduction. '
'Leave it as default or arbitrary if other tasks are selected')
def predict(self, image, task_type='Real-World Image Super-Resolution', jpeg=40, noise=15):
self.args.task = self.tasks[task_type]
self.args.noise = noise
self.args.jpeg = jpeg
# set model path
if self.args.task == 'real_sr':
self.args.scale = 4
self.args.model_path = self.model_zoo[self.args.task][4]
elif self.args.task in ['gray_dn', 'color_dn']:
self.args.model_path = self.model_zoo[self.args.task][noise]
else:
self.args.model_path = self.model_zoo[self.args.task][jpeg]
# set input folder
input_dir = 'input_cog_temp'
os.makedirs(input_dir, exist_ok=True)
input_path = os.path.join(input_dir, os.path.basename(image))
shutil.copy(str(image), input_path)
if self.args.task == 'real_sr':
self.args.folder_lq = input_dir
else:
self.args.folder_gt = input_dir
model = define_model(self.args)
model.eval()
model = model.to(self.device)
# setup folder and path
folder, save_dir, border, window_size = setup(self.args)
os.makedirs(save_dir, exist_ok=True)
test_results = OrderedDict()
test_results['psnr'] = []
test_results['ssim'] = []
test_results['psnr_y'] = []
test_results['ssim_y'] = []
test_results['psnr_b'] = []
# psnr, ssim, psnr_y, ssim_y, psnr_b = 0, 0, 0, 0, 0
out_path = Path(tempfile.mkdtemp()) / "out.png"
for idx, path in enumerate(sorted(glob.glob(os.path.join(folder, '*')))):
# read image
imgname, img_lq, img_gt = get_image_pair(self.args, path) # image to HWC-BGR, float32
img_lq = np.transpose(img_lq if img_lq.shape[2] == 1 else img_lq[:, :, [2, 1, 0]],
(2, 0, 1)) # HCW-BGR to CHW-RGB
img_lq = torch.from_numpy(img_lq).float().unsqueeze(0).to(self.device) # CHW-RGB to NCHW-RGB
# inference
with torch.no_grad():
# pad input image to be a multiple of window_size
_, _, h_old, w_old = img_lq.size()
h_pad = (h_old // window_size + 1) * window_size - h_old
w_pad = (w_old // window_size + 1) * window_size - w_old
img_lq = torch.cat([img_lq, torch.flip(img_lq, [2])], 2)[:, :, :h_old + h_pad, :]
img_lq = torch.cat([img_lq, torch.flip(img_lq, [3])], 3)[:, :, :, :w_old + w_pad]
output = model(img_lq)
output = output[..., :h_old * self.args.scale, :w_old * self.args.scale]
# save image
output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
if output.ndim == 3:
output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0)) # CHW-RGB to HCW-BGR
output = (output * 255.0).round().astype(np.uint8) # float32 to uint8
cv2.imwrite(str(out_path), output)
clean_folder(input_dir)
return out_path
def clean_folder(folder):
for filename in os.listdir(folder):
file_path = os.path.join(folder, filename)
try:
if os.path.isfile(file_path) or os.path.islink(file_path):
os.unlink(file_path)
elif os.path.isdir(file_path):
shutil.rmtree(file_path)
except Exception as e:
print('Failed to delete %s. Reason: %s' % (file_path, e))