KenjieDec commited on
Commit
ef2a093
·
verified ·
1 Parent(s): da7e1d8

Fixed small typo

Browse files
Files changed (1) hide show
  1. app.py +205 -205
app.py CHANGED
@@ -1,206 +1,206 @@
1
- import gradio as gr
2
- import os.path
3
- import numpy as np
4
- from collections import OrderedDict
5
- import torch
6
- import cv2
7
- from PIL import Image, ImageOps
8
- import utils_image as util
9
- from network_fbcnn import FBCNN as net
10
- import requests
11
- import datetime
12
- from gradio_imageslider import ImageSlider
13
-
14
- current_output = None
15
- for model_path in ['fbcnn_gray.pth','fbcnn_color.pth']:
16
- if os.path.exists(model_path):
17
- print(f'{model_path} exists.')
18
- else:
19
- print("downloading model")
20
- url = 'https://github.com/jiaxi-jiang/FBCNN/releases/download/v1.0/{}'.format(os.path.basename(model_path))
21
- r = requests.get(url, allow_redirects=True)
22
- open(model_path, 'wb').write(r.content)
23
-
24
- def inference(input_img, is_gray, res_percentage, input_quality, zoom, x_shift, y_shift):
25
-
26
- print("datetime:", datetime.datetime.utcnow())
27
- input_img_width, input_img_height = Image.fromarray(input_img).size
28
- print("img size:", (input_img_width, input_img_height))
29
-
30
- resized_input = Image.fromarray(input_img).resize(
31
- (
32
- int(input_img_width * (res_percentage/100)),
33
- int(input_img_height * (res_percentage/100))
34
- ), resample = Image.BICUBIC)
35
- input_img = np.array(resized_input)
36
- print("input image resized to:", resized_input.size)
37
-
38
- if is_gray:
39
- n_channels = 1
40
- model_name = 'fbcnn_gray.pth'
41
- else:
42
- n_channels = 3
43
- model_name = 'fbcnn_color.pth'
44
- nc = [64,128,256,512]
45
- nb = 4
46
-
47
- input_quality = 100 - input_quality
48
-
49
- model_path = model_name
50
-
51
- if os.path.exists(model_path):
52
- print(f'{model_path} already exists.')
53
- else:
54
- print("downloading model")
55
- os.makedirs(os.path.dirname(model_path), exist_ok=True)
56
- url = 'https://github.com/jiaxi-jiang/FBCNN/releases/download/v1.0/{}'.format(os.path.basename(model_path))
57
- r = requests.get(url, allow_redirects=True)
58
- open(model_path, 'wb').write(r.content)
59
-
60
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
61
- print("device:", device)
62
-
63
- print(f'loading model from {model_path}')
64
-
65
- model = net(in_nc=n_channels, out_nc=n_channels, nc=nc, nb=nb, act_mode='R')
66
- print("#model.load_state_dict(torch.load(model_path), strict=True)")
67
- model.load_state_dict(torch.load(model_path), strict=True)
68
- print("#model.eval()")
69
- model.eval()
70
- print("#for k, v in model.named_parameters()")
71
- for k, v in model.named_parameters():
72
- v.requires_grad = False
73
- print("#model.to(device)")
74
- model = model.to(device)
75
- print("Model loaded.")
76
-
77
- test_results = OrderedDict()
78
- test_results['psnr'] = []
79
- test_results['ssim'] = []
80
- test_results['psnrb'] = []
81
-
82
- print("#if n_channels")
83
- if n_channels == 1:
84
- open_cv_image = Image.fromarray(input_img)
85
- open_cv_image = ImageOps.grayscale(open_cv_image)
86
- open_cv_image = np.array(open_cv_image)
87
- img = np.expand_dims(open_cv_image, axis=2)
88
- elif n_channels == 3:
89
- open_cv_image = np.array(input_img)
90
- if open_cv_image.ndim == 2:
91
- open_cv_image = cv2.cvtColor(open_cv_image, cv2.COLOR_GRAY2RGB)
92
- else:
93
- open_cv_image = cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2RGB)
94
-
95
- print("#util.uint2tensor4(open_cv_image)")
96
- img_L = util.uint2tensor4(open_cv_image)
97
-
98
- print("#img_L.to(device)")
99
- img_L = img_L.to(device)
100
-
101
- print("#model(img_L)")
102
- img_E, QF = model(img_L)
103
- print("#util.tensor2single(img_E)")
104
- img_E = util.tensor2single(img_E)
105
- print("#util.single2uint(img_E)")
106
- img_E = util.single2uint(img_E)
107
-
108
- print("#torch.tensor([[1-input_quality/100]]).cuda() || torch.tensor([[1-input_quality/100]])")
109
- qf_input = torch.tensor([[1-input_quality/100]]).cuda() if device == torch.device('cuda') else torch.tensor([[1-input_quality/100]])
110
- print("#util.single2uint(img_E)")
111
- img_E, QF = model(img_L, qf_input)
112
-
113
- print("#util.tensor2single(img_E)")
114
- img_E = util.tensor2single(img_E)
115
- print("#util.single2uint(img_E)")
116
- img_E = util.single2uint(img_E)
117
-
118
- if img_E.ndim == 3:
119
- img_E = img_E[:, :, [2, 1, 0]]
120
-
121
- global current_output
122
- current_output = img_E.copy()
123
- print("--inference finished")
124
-
125
- (in_img, out_img) = zoom_image(zoom, x_shift, y_shift, input_img, img_E)
126
- print("--generating preview finished")
127
-
128
- return img_E, (in_img, out_img)
129
-
130
- def zoom_image(zoom, x_shift, y_shift, input_img, output_img = None):
131
- global current_output
132
- if output_img is None:
133
- if current_output is None:
134
- return None
135
- output_img = current_output
136
-
137
- img = Image.fromarray(input_img)
138
- out_img = Image.fromarray(output_img)
139
-
140
- img_w, img_h = img.size
141
- zoom_factor = (100 - zoom) / 100
142
- x_shift /= 100
143
- y_shift /= 100
144
-
145
- zoom_w, zoom_h = int(img_w * zoom_factor), int(img_h * zoom_factor)
146
- x_offset = int((img_w - zoom_w) * x_shift)
147
- y_offset = int((img_h - zoom_h) * y_shift)
148
-
149
- crop_box = (x_offset, y_offset, x_offset + zoom_w, y_offset + zoom_h)
150
- img = img.resize((img_w, img_h), Image.BILINEAR).crop(crop_box)
151
- out_img = out_img.resize((img_w, img_h), Image.BILINEAR).crop(crop_box)
152
-
153
- return (img, out_img)
154
-
155
- with gr.Blocks() as demo:
156
- gr.Markdown("# JPEG Artifacts Removal [FBCNN]")
157
-
158
- with gr.Row():
159
- input_img = gr.Image(label="Input Image")
160
- output_img = gr.Image(label="Result")
161
-
162
- is_gray = gr.Checkbox(label="Grayscale (Check this if your image is grayscale)")
163
- max_res = gr.Slider(1, 100, step=0.5, label="Output image resolution Percentage (Higher% = longer processing time)")
164
- input_quality = gr.Slider(1, 100, step=1, label="Intensity (Higher = stronger JPEG artifact removal)")
165
- zoom = gr.Slider(10, 100, step=1, value=50, label="Zoom Percentage (0 = original size)")
166
- x_shift = gr.Slider(0, 100, step=1, label="Horizontal shift Percentage (Before/After)")
167
- y_shift = gr.Slider(0, 100, step=1, label="Vertical shift Percentage (Before/After)")
168
-
169
- run = gr.Button("Run")
170
-
171
- with gr.Row():
172
- before_after = ImageSlider(label="Before/After", type="pil", value=None)
173
-
174
- run.click(
175
- inference,
176
- inputs=[input_img, is_gray, max_res, input_quality, zoom, x_shift, y_shift],
177
- outputs=[output_img, before_after]
178
- )
179
-
180
- gr.Examples([
181
- ["doraemon.jpg", False, 100, 60, 58, 50, 50],
182
- ["tomandjerry.jpg", False, 100, 60, 60, 57, 44],
183
- ["somepanda.jpg", True, 100, 100, 70, 8, 24],
184
- ["cemetry.jpg", False, 100, 70, 80, 76, 62],
185
- ["michelangelo_david.jpg", True, 100, 30, 88, 53, 27],
186
- ["elon_musk.jpg", False, 100, 45, 75, 33, 30],
187
- ["text.jpg", True, 100, 70, 50, 11, 29]
188
- ], inputs=[input_img, is_gray, max_res, input_quality, zoom, x_shift, y_shift])
189
-
190
- zoom.release(zoom_image, inputs=[zoom, x_shift, y_shift, input_img], outputs=[before_after])
191
- x_shift.release(zoom_image, inputs=[zoom, x_shift, y_shift, input_img], outputs=[before_after])
192
- y_shift.release(zoom_image, inputs=[zoom, x_shift, y_shift, input_img], outputs=[before_after])
193
-
194
- gr.Markdown("""
195
- JPEG Artifacts are noticeable distortions of images caused by JPEG lossy compression.
196
- Note that this is not an AI Upscaler, but just a JPEG Compression Artifact Remover.
197
-
198
- [Original Demo](https://huggingface.co/spaces/danielsapit/JPEG_Artifacts_Removal)
199
- [FBCNN GitHub Repo](https://github.com/jiaxi-jiang/FBCNN)
200
- [Towards Flexible Blind JPEG Artifacts Removal (FBCNN, ICCV 2021)](https://arxiv.org/abs/2109.14573)
201
- [Jiaxi Jiang](https://jiaxi-jiang.github.io/),
202
- [Kai Zhang](https://cszn.github.io/),
203
- [Radu Timofte](http://people.ee.ethz.ch/~timofter/)
204
- """)
205
-
206
  demo.launch()
 
1
+ import gradio as gr
2
+ import os.path
3
+ import numpy as np
4
+ from collections import OrderedDict
5
+ import torch
6
+ import cv2
7
+ from PIL import Image, ImageOps
8
+ import utils_image as util
9
+ from network_fbcnn import FBCNN as net
10
+ import requests
11
+ import datetime
12
+ from gradio_imageslider import ImageSlider
13
+
14
+ current_output = None
15
+ for model_path in ['fbcnn_gray.pth','fbcnn_color.pth']:
16
+ if os.path.exists(model_path):
17
+ print(f'{model_path} exists.')
18
+ else:
19
+ print("downloading model")
20
+ url = 'https://github.com/jiaxi-jiang/FBCNN/releases/download/v1.0/{}'.format(os.path.basename(model_path))
21
+ r = requests.get(url, allow_redirects=True)
22
+ open(model_path, 'wb').write(r.content)
23
+
24
+ def inference(input_img, is_gray, res_percentage, input_quality, zoom, x_shift, y_shift):
25
+
26
+ print("datetime:", datetime.datetime.utcnow())
27
+ input_img_width, input_img_height = Image.fromarray(input_img).size
28
+ print("img size:", (input_img_width, input_img_height))
29
+
30
+ resized_input = Image.fromarray(input_img).resize(
31
+ (
32
+ int(input_img_width * (res_percentage/100)),
33
+ int(input_img_height * (res_percentage/100))
34
+ ), resample = Image.BICUBIC)
35
+ input_img = np.array(resized_input)
36
+ print("input image resized to:", resized_input.size)
37
+
38
+ if is_gray:
39
+ n_channels = 1
40
+ model_name = 'fbcnn_gray.pth'
41
+ else:
42
+ n_channels = 3
43
+ model_name = 'fbcnn_color.pth'
44
+ nc = [64,128,256,512]
45
+ nb = 4
46
+
47
+ input_quality = 100 - input_quality
48
+
49
+ model_path = model_name
50
+
51
+ if os.path.exists(model_path):
52
+ print(f'{model_path} already exists.')
53
+ else:
54
+ print("downloading model")
55
+ os.makedirs(os.path.dirname(model_path), exist_ok=True)
56
+ url = 'https://github.com/jiaxi-jiang/FBCNN/releases/download/v1.0/{}'.format(os.path.basename(model_path))
57
+ r = requests.get(url, allow_redirects=True)
58
+ open(model_path, 'wb').write(r.content)
59
+
60
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
61
+ print("device:", device)
62
+
63
+ print(f'loading model from {model_path}')
64
+
65
+ model = net(in_nc=n_channels, out_nc=n_channels, nc=nc, nb=nb, act_mode='R')
66
+ print("#model.load_state_dict(torch.load(model_path), strict=True)")
67
+ model.load_state_dict(torch.load(model_path), strict=True)
68
+ print("#model.eval()")
69
+ model.eval()
70
+ print("#for k, v in model.named_parameters()")
71
+ for k, v in model.named_parameters():
72
+ v.requires_grad = False
73
+ print("#model.to(device)")
74
+ model = model.to(device)
75
+ print("Model loaded.")
76
+
77
+ test_results = OrderedDict()
78
+ test_results['psnr'] = []
79
+ test_results['ssim'] = []
80
+ test_results['psnrb'] = []
81
+
82
+ print("#if n_channels")
83
+ if n_channels == 1:
84
+ open_cv_image = Image.fromarray(input_img)
85
+ open_cv_image = ImageOps.grayscale(open_cv_image)
86
+ open_cv_image = np.array(open_cv_image)
87
+ img = np.expand_dims(open_cv_image, axis=2)
88
+ elif n_channels == 3:
89
+ open_cv_image = np.array(input_img)
90
+ if open_cv_image.ndim == 2:
91
+ open_cv_image = cv2.cvtColor(open_cv_image, cv2.COLOR_GRAY2RGB)
92
+ else:
93
+ open_cv_image = cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2RGB)
94
+
95
+ print("#util.uint2tensor4(open_cv_image)")
96
+ img_L = util.uint2tensor4(open_cv_image)
97
+
98
+ print("#img_L.to(device)")
99
+ img_L = img_L.to(device)
100
+
101
+ print("#model(img_L)")
102
+ img_E, QF = model(img_L)
103
+ print("#util.tensor2single(img_E)")
104
+ img_E = util.tensor2single(img_E)
105
+ print("#util.single2uint(img_E)")
106
+ img_E = util.single2uint(img_E)
107
+
108
+ print("#torch.tensor([[1-input_quality/100]]).cuda() || torch.tensor([[1-input_quality/100]])")
109
+ qf_input = torch.tensor([[1-input_quality/100]]).cuda() if device == torch.device('cuda') else torch.tensor([[1-input_quality/100]])
110
+ print("#util.single2uint(img_E)")
111
+ img_E, QF = model(img_L, qf_input)
112
+
113
+ print("#util.tensor2single(img_E)")
114
+ img_E = util.tensor2single(img_E)
115
+ print("#util.single2uint(img_E)")
116
+ img_E = util.single2uint(img_E)
117
+
118
+ if img_E.ndim == 3:
119
+ img_E = img_E[:, :, [2, 1, 0]]
120
+
121
+ global current_output
122
+ current_output = img_E.copy()
123
+ print("--inference finished")
124
+
125
+ (in_img, out_img) = zoom_image(zoom, x_shift, y_shift, input_img, img_E)
126
+ print("--generating preview finished")
127
+
128
+ return img_E, (in_img, out_img)
129
+
130
+ def zoom_image(zoom, x_shift, y_shift, input_img, output_img = None):
131
+ global current_output
132
+ if output_img is None:
133
+ if current_output is None:
134
+ return None
135
+ output_img = current_output
136
+
137
+ img = Image.fromarray(input_img)
138
+ out_img = Image.fromarray(output_img)
139
+
140
+ img_w, img_h = img.size
141
+ zoom_factor = (100 - zoom) / 100
142
+ x_shift /= 100
143
+ y_shift /= 100
144
+
145
+ zoom_w, zoom_h = int(img_w * zoom_factor), int(img_h * zoom_factor)
146
+ x_offset = int((img_w - zoom_w) * x_shift)
147
+ y_offset = int((img_h - zoom_h) * y_shift)
148
+
149
+ crop_box = (x_offset, y_offset, x_offset + zoom_w, y_offset + zoom_h)
150
+ img = img.resize((img_w, img_h), Image.BILINEAR).crop(crop_box)
151
+ out_img = out_img.resize((img_w, img_h), Image.BILINEAR).crop(crop_box)
152
+
153
+ return (img, out_img)
154
+
155
+ with gr.Blocks() as demo:
156
+ gr.Markdown("# JPEG Artifacts Removal [FBCNN]")
157
+
158
+ with gr.Row():
159
+ input_img = gr.Image(label="Input Image")
160
+ output_img = gr.Image(label="Result")
161
+
162
+ is_gray = gr.Checkbox(label="Grayscale (Check this if your image is grayscale)")
163
+ max_res = gr.Slider(1, 100, step=0.5, label="Output image resolution Percentage (Higher% = longer processing time)")
164
+ input_quality = gr.Slider(1, 100, step=1, label="Intensity (Higher = stronger JPEG artifact removal)")
165
+ zoom = gr.Slider(0, 100, step=1, value=50, label="Zoom Percentage (0 = original size)")
166
+ x_shift = gr.Slider(0, 100, step=1, label="Horizontal shift Percentage (Before/After)")
167
+ y_shift = gr.Slider(0, 100, step=1, label="Vertical shift Percentage (Before/After)")
168
+
169
+ run = gr.Button("Run")
170
+
171
+ with gr.Row():
172
+ before_after = ImageSlider(label="Before/After", type="pil", value=None)
173
+
174
+ run.click(
175
+ inference,
176
+ inputs=[input_img, is_gray, max_res, input_quality, zoom, x_shift, y_shift],
177
+ outputs=[output_img, before_after]
178
+ )
179
+
180
+ gr.Examples([
181
+ ["doraemon.jpg", False, 100, 60, 58, 50, 50],
182
+ ["tomandjerry.jpg", False, 100, 60, 60, 57, 44],
183
+ ["somepanda.jpg", True, 100, 100, 70, 8, 24],
184
+ ["cemetry.jpg", False, 100, 70, 80, 76, 62],
185
+ ["michelangelo_david.jpg", True, 100, 30, 88, 53, 27],
186
+ ["elon_musk.jpg", False, 100, 45, 75, 33, 30],
187
+ ["text.jpg", True, 100, 70, 50, 11, 29]
188
+ ], inputs=[input_img, is_gray, max_res, input_quality, zoom, x_shift, y_shift])
189
+
190
+ zoom.release(zoom_image, inputs=[zoom, x_shift, y_shift, input_img], outputs=[before_after])
191
+ x_shift.release(zoom_image, inputs=[zoom, x_shift, y_shift, input_img], outputs=[before_after])
192
+ y_shift.release(zoom_image, inputs=[zoom, x_shift, y_shift, input_img], outputs=[before_after])
193
+
194
+ gr.Markdown("""
195
+ JPEG Artifacts are noticeable distortions of images caused by JPEG lossy compression.
196
+ Note that this is not an AI Upscaler, but just a JPEG Compression Artifact Remover.
197
+
198
+ [Original Demo](https://huggingface.co/spaces/danielsapit/JPEG_Artifacts_Removal)
199
+ [FBCNN GitHub Repo](https://github.com/jiaxi-jiang/FBCNN)
200
+ [Towards Flexible Blind JPEG Artifacts Removal (FBCNN, ICCV 2021)](https://arxiv.org/abs/2109.14573)
201
+ [Jiaxi Jiang](https://jiaxi-jiang.github.io/),
202
+ [Kai Zhang](https://cszn.github.io/),
203
+ [Radu Timofte](http://people.ee.ethz.ch/~timofter/)
204
+ """)
205
+
206
  demo.launch()