| import gradio as gr | |
| import os.path | |
| import numpy as np | |
| from collections import OrderedDict | |
| import torch | |
| import cv2 | |
| from PIL import Image, ImageOps | |
| import utils_image as util | |
| from network_fbcnn import FBCNN as net | |
| import requests | |
| import datetime | |
| print(gr.__version__) | |
| for model_path in ['fbcnn_gray.pth','fbcnn_color.pth']: | |
| if os.path.exists(model_path): | |
| print(f'{model_path} exists.') | |
| else: | |
| print("Downloading model: ", f'{model_path}') | |
| url = 'https://github.com/jiaxi-jiang/FBCNN/releases/download/v1.0/{}'.format(os.path.basename(model_path)) | |
| r = requests.get(url, allow_redirects=True) | |
| open(model_path, 'wb').write(r.content) | |
| def inference(filepaths, is_gray, res_percentage, input_quality): | |
| outputs = [] | |
| before_afters = [] | |
| if filepaths is None: | |
| return [], None | |
| for filepath, *_ in filepaths: | |
| filename = os.path.basename(filepath) | |
| print("Processing: ", filename) | |
| input_img = np.array(Image.open(filepath).convert("RGB")) | |
| print("Datetime: ", datetime.datetime.utcnow()) | |
| input_img_width, input_img_height = Image.fromarray(input_img).size | |
| print("Img size: ", (input_img_width, input_img_height)) | |
| resized_input = Image.fromarray(input_img).resize( | |
| ( | |
| int(input_img_width * (res_percentage/100)), | |
| int(input_img_height * (res_percentage/100)) | |
| ), resample = Image.BICUBIC) | |
| input_img = np.array(resized_input) | |
| print("Input image resized to: ", resized_input.size) | |
| if is_gray: | |
| n_channels = 1 | |
| model_name = 'fbcnn_gray.pth' | |
| else: | |
| n_channels = 3 | |
| model_name = 'fbcnn_color.pth' | |
| nc = [64,128,256,512] | |
| nb = 4 | |
| input_quality = 100 - input_quality | |
| model_path = model_name | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| print("Device: ", device) | |
| print(f'Loading model from {model_path}') | |
| model = net(in_nc=n_channels, out_nc=n_channels, nc=nc, nb=nb, act_mode='R') | |
| print("#model.load_state_dict(torch.load(model_path), strict=True)") | |
| model.load_state_dict(torch.load(model_path), strict=True) | |
| print("#model.eval()") | |
| model.eval() | |
| print("#for k, v in model.named_parameters()") | |
| for k, v in model.named_parameters(): | |
| v.requires_grad = False | |
| print("#model.to(device)") | |
| model = model.to(device) | |
| print("Model loaded.") | |
| test_results = OrderedDict() | |
| test_results['psnr'] = [] | |
| test_results['ssim'] = [] | |
| test_results['psnrb'] = [] | |
| print("#if n_channels") | |
| if n_channels == 1: | |
| open_cv_image = Image.fromarray(input_img) | |
| open_cv_image = ImageOps.grayscale(open_cv_image) | |
| open_cv_image = np.array(open_cv_image) | |
| img = np.expand_dims(open_cv_image, axis=2) | |
| elif n_channels == 3: | |
| open_cv_image = np.array(input_img) | |
| if open_cv_image.ndim == 2: | |
| open_cv_image = cv2.cvtColor(open_cv_image, cv2.COLOR_GRAY2RGB) | |
| else: | |
| open_cv_image = cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2RGB) | |
| print("#util.uint2tensor4(open_cv_image)") | |
| img_L = util.uint2tensor4(open_cv_image) | |
| print("#img_L.to(device)") | |
| img_L = img_L.to(device) | |
| print("#model(img_L)") | |
| img_E, QF = model(img_L) | |
| print("#util.tensor2single(img_E)") | |
| img_E = util.tensor2single(img_E) | |
| print("#util.single2uint(img_E)") | |
| img_E = util.single2uint(img_E) | |
| print("#torch.tensor([[1-input_quality/100]]).cuda() || torch.tensor([[1-input_quality/100]])") | |
| qf_input = torch.tensor([[1-input_quality/100]]).cuda() if device == torch.device('cuda') else torch.tensor([[1-input_quality/100]]) | |
| print("#util.single2uint(img_E)") | |
| img_E, QF = model(img_L, qf_input) | |
| print("#util.tensor2single(img_E)") | |
| img_E = util.tensor2single(img_E) | |
| print("#util.single2uint(img_E)") | |
| img_E = util.single2uint(img_E) | |
| if img_E.ndim == 3: | |
| img_E = img_E[:, :, [2, 1, 0]] | |
| print("--inference finished") | |
| image_path = check_file_exist("output_images", filename) | |
| outputs.append((img_E, f'{filename}')) | |
| before_afters.append((input_img, img_E)) | |
| return outputs, before_afters | |
| def select_image(event: gr.SelectData, before_afters): | |
| index = event.index | |
| if index is None or index >= len(before_afters): | |
| return None | |
| return before_afters[index], index | |
| def select_changed_image(index, before_afters): | |
| if index is None or index >= len(before_afters): | |
| return None | |
| return before_afters[index] | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# JPEG Artifacts Removal [FBCNN]") | |
| with gr.Row(): | |
| input_img = gr.Gallery( | |
| label="Input Image(s)", | |
| file_types=['image'], | |
| type="filepath", | |
| height="auto" | |
| ) | |
| output_img = gr.Gallery( | |
| label="Results", | |
| height="auto", | |
| interactive=False | |
| ) | |
| is_gray = gr.Checkbox(label="Grayscale (Check this if your image is grayscale)") | |
| max_res = gr.Slider(1, 100, step=0.5, value=100, label="Output image resolution Percentage (Higher% = longer processing time)") | |
| input_quality = gr.Slider(1, 100, step=1, value=40, label="Intensity (Higher = stronger JPEG artifact removal)") | |
| run = gr.Button("Run") | |
| with gr.Row(): | |
| before_afters = gr.State([]) | |
| current_index = gr.State(0) | |
| before_after = gr.ImageSlider(label="Before/After (Select One)", value=None, height="auto") | |
| run.click( | |
| inference, | |
| inputs=[input_img, is_gray, max_res, input_quality], | |
| outputs=[output_img, before_afters] | |
| ) | |
| output_img.select( | |
| select_image, | |
| inputs=[before_afters], | |
| outputs=[before_after, current_index] | |
| ) | |
| output_img.change( | |
| select_changed_image, | |
| inputs=[current_index, before_afters], | |
| outputs=[before_after] | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| [[("doraemon.jpg", "doraemon.jpg")], False, 100, 60], | |
| [[("tomandjerry.jpg", "tomandjerry.jpg")], False, 100, 60], | |
| [[("somepanda.jpg", "somepanda.jpg")], True, 100, 100], | |
| [[("cemetry.jpg", "cemetry.jpg")], False, 100, 70], | |
| [[("michelangelo_david.jpg", "michelangelo_david.jpg")], True, 100, 30], | |
| [[("elon_musk.jpg", "elon_musk.jpg")], False, 100, 45], | |
| [[("text.jpg", "text.jpg")], True, 100, 70] | |
| ], | |
| inputs=[input_img, is_gray, max_res, input_quality], | |
| outputs=[output_img, before_afters] | |
| ) | |
| gr.Markdown(""" | |
| JPEG Artifacts are noticeable distortions of images caused by JPEG lossy compression. | |
| Note that this is not an AI Upscaler, but just a JPEG Compression Artifact Remover. | |
| [Original Demo](https://huggingface.co/spaces/danielsapit/JPEG_Artifacts_Removal) | |
| [FBCNN GitHub Repo](https://github.com/jiaxi-jiang/FBCNN) | |
| [Towards Flexible Blind JPEG Artifacts Removal (FBCNN, ICCV 2021)](https://arxiv.org/abs/2109.14573) | |
| [Jiaxi Jiang](https://jiaxi-jiang.github.io/), | |
| [Kai Zhang](https://cszn.github.io/), | |
| [Radu Timofte](http://people.ee.ethz.ch/~timofter/) | |
| """) | |
| demo.launch() |