import os import io import gradio as gr from PIL import Image import tempfile from pathlib import Path import argparse import shutil 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 os.system('wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pth -P model_zoo/swinir') os.system('wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/004_grayDN_DFWB_s128w8_SwinIR-M_noise15.pth -P model_zoo/swinir') os.system('wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/004_grayDN_DFWB_s128w8_SwinIR-M_noise25.pth -P model_zoo/swinir') os.system('wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/004_grayDN_DFWB_s128w8_SwinIR-M_noise50.pth -P model_zoo/swinir') os.system('wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/005_colorDN_DFWB_s128w8_SwinIR-M_noise15.pth -P model_zoo/swinir') os.system('wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/005_colorDN_DFWB_s128w8_SwinIR-M_noise25.pth -P model_zoo/swinir') os.system('wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/005_colorDN_DFWB_s128w8_SwinIR-M_noise50.pth -P model_zoo/swinir') os.system('wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/006_CAR_DFWB_s126w7_SwinIR-M_jpeg10.pth -P model_zoo/swinir') os.system('wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/006_CAR_DFWB_s126w7_SwinIR-M_jpeg20.pth -P model_zoo/swinir') os.system('wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/006_CAR_DFWB_s126w7_SwinIR-M_jpeg30.pth -P model_zoo/swinir') os.system('wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/006_CAR_DFWB_s126w7_SwinIR-M_jpeg40.pth -P model_zoo/swinir') def sentence_builder(image, task_type, noise, jpeg): model_dir = 'model_zoo/swinir' 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=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') args = parser.parse_args('') device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') tasks = { 'Real-World Image Super-Resolution': 'real_sr', 'Grayscale Image Denoising': 'gray_dn', 'Color Image Denoising': 'color_dn', 'JPEG Compression Artifact Reduction': 'jpeg_car' } args.task = tasks[task_type] args.noise = noise args.jpeg = jpeg # set model path if args.task == 'real_sr': args.scale = 4 args.model_path = model_zoo[args.task][4] elif args.task in ['gray_dn', 'color_dn']: args.model_path = model_zoo[args.task][noise] else: args.model_path = model_zoo[args.task][jpeg] try: # set input folder input_dir = 'input_cog_temp' os.makedirs(input_dir, exist_ok=True) input_path = os.path.join(input_dir, guess_filename(image)) #shutil.copy(str(image), input_path) image.save(input_path, "JPEG") if args.task == 'real_sr': args.folder_lq = input_dir else: args.folder_gt = input_dir model = define_model(args) model.eval() model = model.to(device) # setup folder and path folder, save_dir, border, window_size = setup(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(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(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 * args.scale, :w_old * 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) finally: clean_folder(input_dir) return out_path def guess_filename(obj: io.IOBase) -> str: """Tries to guess the filename of the given object.""" name = getattr(obj, "name", "input") return os.path.basename(name) 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)) title = "Dmonin-SwinIR" description = "Gradio for SwinIR. SwinIR achieves state-of-the-art performance on six tasks: image super-resolution (including classical, lightweight and real-world image super-resolution), image denoising (including grayscale and color image denoising) and JPEG compression artifact reduction. See the paper and project page for detailed results below. Here, we provide a demo for real-world image SR.To use it, simply upload your image, or click one of the examples to load them." article = "

SwinIR: Image Restoration Using Swin Transformer | Github Repo

" examples=[['ETH_LR.png']] gr.Interface( sentence_builder, [ gr.inputs.Image(type="pil", label="Input"), gr.Dropdown(choices=["Real-World Image Super-Resolution", "Grayscale Image Denoising", "Color Image Denoising", "JPEG Compression Artifact Reduction"], default= "Real-World Image Super-Resolution", value= "Real-World Image Super-Resolution"), gr.Dropdown(choices=["15", "25", "50"], default = "15", value="15"), gr.Dropdown(choices=["10", "20", "30", "40"], default="40", value="40") ], gr.outputs.Image(type="filepath", label="Output"), title=title, description=description, article=article, enable_queue=True, examples=examples ).launch()