ohayonguy
commited on
Commit
·
5320385
1
Parent(s):
5afc7ad
updated title and description
Browse files
app.py
CHANGED
|
@@ -108,13 +108,8 @@ def inference(img, aligned, scale, num_flow_steps):
|
|
| 108 |
if scale > 4:
|
| 109 |
scale = 4 # avoid too large scale value
|
| 110 |
img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
|
| 111 |
-
if len(img.shape) ==
|
| 112 |
-
img_mode = 'RGBA'
|
| 113 |
-
elif len(img.shape) == 2: # for gray inputs
|
| 114 |
-
img_mode = None
|
| 115 |
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
|
| 116 |
-
else:
|
| 117 |
-
img_mode = None
|
| 118 |
|
| 119 |
h, w = img.shape[0:2]
|
| 120 |
if h > 3500 or w > 3500:
|
|
@@ -136,34 +131,25 @@ def inference(img, aligned, scale, num_flow_steps):
|
|
| 136 |
|
| 137 |
has_aligned = True if aligned == 'Yes' else False
|
| 138 |
_, restored_aligned, restored_img = enhance_face(img, face_helper, has_aligned, only_center_face=False,
|
| 139 |
-
paste_back=True, num_flow_steps=num_flow_steps)
|
| 140 |
if has_aligned:
|
| 141 |
output = restored_aligned[0]
|
| 142 |
else:
|
| 143 |
output = restored_img
|
| 144 |
|
| 145 |
-
|
| 146 |
-
# try:
|
| 147 |
-
# if scale != 2:
|
| 148 |
-
# interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
|
| 149 |
-
# h, w = img.shape[0:2]
|
| 150 |
-
# output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
|
| 151 |
-
# except Exception as error:
|
| 152 |
-
# print('Wrong scale input.', error)
|
| 153 |
-
if img_mode == 'RGBA': # RGBA images should be saved in png format
|
| 154 |
-
extension = 'png'
|
| 155 |
-
else:
|
| 156 |
-
extension = 'jpg'
|
| 157 |
-
save_path = f'output/out.{extension}'
|
| 158 |
cv2.imwrite(save_path, output)
|
| 159 |
|
| 160 |
output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
|
| 161 |
return output, save_path
|
| 162 |
-
# except Exception as error:
|
| 163 |
-
# print('global exception', error)
|
| 164 |
-
# return None, None
|
| 165 |
|
| 166 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 167 |
css = r"""
|
| 168 |
"""
|
| 169 |
|
|
@@ -171,12 +157,14 @@ demo = gr.Interface(
|
|
| 171 |
inference, [
|
| 172 |
gr.Image(type="filepath", label="Input"),
|
| 173 |
gr.Radio(['Yes', 'No'], type="value", value='aligned', label='Is the input an aligned face image?'),
|
| 174 |
-
gr.Number(label="
|
| 175 |
-
gr.Number(label="Number of flow steps. A higher value should result in better image quality, but
|
| 176 |
], [
|
| 177 |
gr.Image(type="numpy", label="Output"),
|
| 178 |
gr.File(label="Download the output image")
|
| 179 |
],
|
|
|
|
|
|
|
| 180 |
)
|
| 181 |
|
| 182 |
|
|
|
|
| 108 |
if scale > 4:
|
| 109 |
scale = 4 # avoid too large scale value
|
| 110 |
img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
|
| 111 |
+
if len(img.shape) == 2: # for gray inputs
|
|
|
|
|
|
|
|
|
|
| 112 |
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
|
|
|
|
|
|
|
| 113 |
|
| 114 |
h, w = img.shape[0:2]
|
| 115 |
if h > 3500 or w > 3500:
|
|
|
|
| 131 |
|
| 132 |
has_aligned = True if aligned == 'Yes' else False
|
| 133 |
_, restored_aligned, restored_img = enhance_face(img, face_helper, has_aligned, only_center_face=False,
|
| 134 |
+
paste_back=True, num_flow_steps=num_flow_steps, scale=scale)
|
| 135 |
if has_aligned:
|
| 136 |
output = restored_aligned[0]
|
| 137 |
else:
|
| 138 |
output = restored_img
|
| 139 |
|
| 140 |
+
save_path = f'output/out.png'
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 141 |
cv2.imwrite(save_path, output)
|
| 142 |
|
| 143 |
output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
|
| 144 |
return output, save_path
|
|
|
|
|
|
|
|
|
|
| 145 |
|
| 146 |
|
| 147 |
+
title = "Posterior-Mean Rectified Flow: Towards Minimum MSE Photo-Realistic Image Restoration"
|
| 148 |
+
|
| 149 |
+
description = r"""
|
| 150 |
+
Gradio demo for Posterior-Mean Rectified Flow (PMRF). Please refer to our project's page: https://pmrf-ml.github.io/.
|
| 151 |
+
"""
|
| 152 |
+
|
| 153 |
css = r"""
|
| 154 |
"""
|
| 155 |
|
|
|
|
| 157 |
inference, [
|
| 158 |
gr.Image(type="filepath", label="Input"),
|
| 159 |
gr.Radio(['Yes', 'No'], type="value", value='aligned', label='Is the input an aligned face image?'),
|
| 160 |
+
gr.Number(label="Scale factor for the background upsampler. Insert a value between 1 and 4 (including). Applicable only to non-aligned face images.", value=1),
|
| 161 |
+
gr.Number(label="Number of flow steps. A higher value should result in better image quality, but will inference will take a longer time.", value=25),
|
| 162 |
], [
|
| 163 |
gr.Image(type="numpy", label="Output"),
|
| 164 |
gr.File(label="Download the output image")
|
| 165 |
],
|
| 166 |
+
title=title,
|
| 167 |
+
description=description
|
| 168 |
)
|
| 169 |
|
| 170 |
|