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 from gradio_imageslider import ImageSlider current_output = None 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") 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(input_img, is_gray, input_quality, zoom, x_shift, y_shift): 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)) if (input_img_width > 1080) or (input_img_height > 1080): resize_ratio = min(1080/input_img_width, 1080/input_img_height) resized_input = Image.fromarray(input_img).resize( (int(input_img_width*resize_ratio) + (input_img_width*resize_ratio < 1), int(input_img_height*resize_ratio) + (input_img_height*resize_ratio < 1)), 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 if os.path.exists(model_path): print(f'{model_path} already exists.') else: print("downloading model") os.makedirs(os.path.dirname(model_path), exist_ok=True) 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) 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]] global current_output current_output = img_E.copy() print("--inference finished") (in_img, out_img) = zoom_image(zoom, x_shift, y_shift, input_img, img_E) print("--generating preview finished") return img_E, (in_img, out_img) def zoom_image(zoom, x_shift, y_shift, input_img, output_img = None): global current_output if output_img is None: if current_output is None: return None output_img = current_output img = Image.fromarray(input_img) out_img = Image.fromarray(output_img) img_w, img_h = img.size zoom_factor = (100 - zoom) / 100 x_shift /= 100 y_shift /= 100 zoom_w, zoom_h = int(img_w * zoom_factor), int(img_h * zoom_factor) x_offset = int((img_w - zoom_w) * x_shift) y_offset = int((img_h - zoom_h) * y_shift) crop_box = (x_offset, y_offset, x_offset + zoom_w, y_offset + zoom_h) img = img.crop(crop_box).resize((img_w, img_h), Image.BILINEAR) out_img = out_img.crop(crop_box).resize((img_w, img_h), Image.BILINEAR) return (img, out_img) with gr.Blocks() as demo: gr.Markdown("# JPEG Artifacts Removal [FBCNN]") with gr.Row(): input_img = gr.Image(label="Input Image") output_img = gr.Image(label="Result") is_gray = gr.Checkbox(label="Grayscale (Check this if your image is grayscale)") input_quality = gr.Slider(1, 100, step=1, label="Intensity (Higher = stronger JPEG artifact removal)") zoom = gr.Slider(10, 100, step=1, value=50, label="Zoom Percentage (0 = original size)") x_shift = gr.Slider(0, 100, step=1, label="Horizontal shift Percentage (Before/After)") y_shift = gr.Slider(0, 100, step=1, label="Vertical shift Percentage (Before/After)") run = gr.Button("Run") with gr.Row(): before_after = ImageSlider(label="Before/After", type="pil", value=None) run.click( inference, inputs=[input_img, is_gray, input_quality, zoom, x_shift, y_shift], outputs=[output_img, before_after] ) gr.Examples([ ["doraemon.jpg", False, 60, 58, 50, 50], ["tomandjerry.jpg", False, 60, 60, 57, 44], ["somepanda.jpg", True, 100, 70, 8, 24], ["cemetry.jpg", False, 70, 80, 76, 62], ["michelangelo_david.jpg", True, 30, 88, 53, 27], ["elon_musk.jpg", False, 45, 75, 33, 30], ["text.jpg", True, 70, 50, 11, 29] ], inputs=[input_img, is_gray, input_quality, zoom, x_shift, y_shift]) zoom.release(zoom_image, inputs=[zoom, x_shift, y_shift, input_img], outputs=[before_after]) x_shift.release(zoom_image, inputs=[zoom, x_shift, y_shift, input_img], outputs=[before_after]) y_shift.release(zoom_image, inputs=[zoom, x_shift, y_shift, input_img], outputs=[before_after]) 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()