File size: 8,701 Bytes
3b909cb
 
 
 
 
 
 
fe525fc
 
3b909cb
fb4f2b9
3b909cb
7b0d035
7fd1595
494344f
7b0d035
5cb71e1
7b0d035
 
 
3b909cb
10907b9
fb4f2b9
 
 
 
 
 
 
397cc99
 
 
fb4f2b9
 
3ae9a61
3b909cb
 
 
 
 
 
 
 
 
 
 
 
f7e9c78
 
 
3ae9a61
f7e9c78
5cb71e1
f7e9c78
 
 
 
3b909cb
 
3ae9a61
3b909cb
 
 
 
3ae9a61
 
 
10907b9
f3e03bc
10907b9
f3e03bc
10907b9
f3e03bc
10907b9
 
f3e03bc
10907b9
f3e03bc
 
10907b9
 
 
 
 
 
 
 
 
3ae9a61
10907b9
 
 
 
 
 
 
 
 
 
 
 
3ae9a61
10907b9
3ae9a61
 
10907b9
 
 
 
 
3ae9a61
 
10907b9
3ae9a61
10907b9
3ae9a61
10907b9
3ae9a61
 
10907b9
3ae9a61
10907b9
b002866
3ae9a61
10907b9
3ae9a61
10907b9
 
 
 
 
 
 
3b909cb
 
 
 
 
 
 
 
10907b9
3b909cb
 
 
 
3ae9a61
 
 
10907b9
3b909cb
494344f
3b909cb
8cdc393
3b909cb
535092b
 
03b810c
eb456f3
03b810c
eb456f3
03b810c
10907b9
3b909cb
 
 
10907b9
 
 
 
 
 
 
 
535092b
 
b002866
81e4d97
4165a15
03b810c
 
 
 
5cb71e1
 
 
 
5701c9d
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
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

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 # set 1 for grayscale image, set 3 for color image
        model_name = 'fbcnn_gray.pth'
    else:
        n_channels = 3 # set 1 for grayscale image, set 3 for color image
        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)

    # ----------------------------------------
    # load model
    # ----------------------------------------
    
    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'] = []

    # ------------------------------------
    # (1) img_L
    # ------------------------------------

    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) # PIL to open cv image
        img = np.expand_dims(open_cv_image, axis=2)  # HxWx1
    elif n_channels == 3:
        open_cv_image = np.array(input_img) # PIL to open cv image
        if open_cv_image.ndim == 2:
            open_cv_image = cv2.cvtColor(open_cv_image, cv2.COLOR_GRAY2RGB)  # GGG
        else:
            open_cv_image = cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2RGB)  # RGB

    print("#util.uint2tensor4(open_cv_image)")
    img_L = util.uint2tensor4(open_cv_image)
    
    print("#img_L.to(device)")
    img_L = img_L.to(device)

    # ------------------------------------
    # (2) img_E
    # ------------------------------------
    
    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")

    out_img = Image.fromarray(img_E)
    out_img_w, out_img_h = out_img.size # output image size
    zoom = zoom/100
    x_shift = x_shift/100
    y_shift = y_shift/100
    zoom_w, zoom_h = out_img_w*zoom, out_img_h*zoom
    zoom_left, zoom_right = int((out_img_w - zoom_w)*x_shift), int(zoom_w + (out_img_w - zoom_w)*x_shift)
    zoom_top, zoom_bottom = int((out_img_h - zoom_h)*y_shift), int(zoom_h + (out_img_h - zoom_h)*y_shift)
    in_img = Image.fromarray(input_img)
    in_img = in_img.crop((zoom_left, zoom_top, zoom_right, zoom_bottom))
    in_img = in_img.resize((int(zoom_w/zoom), int(zoom_h/zoom)), Image.NEAREST)
    out_img = out_img.crop((zoom_left, zoom_top, zoom_right, zoom_bottom))
    out_img = out_img.resize((int(zoom_w/zoom), int(zoom_h/zoom)), Image.NEAREST)
    
    print("--generating preview finished")
    
    return img_E, in_img, out_img
    
gr.Interface(
    fn = inference,
    inputs = [gr.inputs.Image(label="Input Image"),
              gr.inputs.Checkbox(label="Grayscale (Check this if your image is grayscale)"),
              gr.inputs.Slider(minimum=1, maximum=100, step=1, label="Intensity (Higher = stronger JPEG artifact removal)"),
              gr.inputs.Slider(minimum=10, maximum=100, step=1, default=50, label="Zoom Image "
                                                                                  "(Use this to see a copy of the output image up close. "
                                                                                   "100 = original size)"),
              gr.inputs.Slider(minimum=0, maximum=100, step=1, label="Zoom horizontal shift "
                                                                     "(Increase to shift to the right)"),
              gr.inputs.Slider(minimum=0, maximum=100, step=1, label="Zoom vertical shift "
                                                                     "(Increase to shift downwards)")
              ],
    outputs = [gr.outputs.Image(label="Result"),
               gr.outputs.Image(label="Before:"),
               gr.outputs.Image(label="After:")],
    examples = [["doraemon.jpg",False,60,42,50,50],
               ["tomandjerry.jpg",False,60,40,57,44],
               ["somepanda.jpg",True,100,30,8,24],
               ["cemetry.jpg",False,70,20,76,62],
               ["michelangelo_david.jpg",True,30,12,53,27],
               ["elon_musk.jpg",False,45,15,33,30],
               ["text.jpg",True,70,50,11,29]],
    title = "JPEG Artifacts Removal [FBCNN]",
    description = "Gradio Demo for JPEG Artifacts Removal. To use it, simply upload your image, "
                  "or click one of the examples to load them. Check out the paper and the original GitHub repo at the links below. "
                  "JPEG artifacts are noticeable distortions of images caused by JPEG lossy compression. "
                  "This is not a super-resolution AI but a JPEG compression artifact remover. "
                  "Written below the examples are the limitations of the input image. ",
    article = "<p style='text-align: left;'>Uploaded images with a length longer than 1080 pixels will be downscaled to a smaller size "
              "with a length of 1080 pixels. Uploaded images with transparency will be incorrectly reconstructed at the output.</p>"
              "<p style='text-align: center;'><a href='https://github.com/jiaxi-jiang/FBCNN'>FBCNN GitHub Repo</a><br>"
              "<a href='https://arxiv.org/abs/2109.14573'>Towards Flexible Blind JPEG Artifacts Removal (FBCNN, ICCV 2021)</a><br>"
              "<a href='https://jiaxi-jiang.github.io/'>Jiaxi Jiang, </a>"
              "<a href='https://cszn.github.io/'>Kai Zhang, </a>"
              "<a href='http://people.ee.ethz.ch/~timofter/'>Radu Timofte</a></p>",
    allow_flagging="never"
).launch(enable_queue=True)