Spaces:
Running
on
Zero
Running
on
Zero
ohayonguy
commited on
Commit
•
ad47941
1
Parent(s):
917533b
improved description
Browse files
app.py
CHANGED
@@ -3,6 +3,7 @@ if os.getenv('SPACES_ZERO_GPU') == "true":
|
|
3 |
os.environ['SPACES_ZERO_GPU'] = "1"
|
4 |
os.environ['K_DIFFUSION_USE_COMPILE'] = "0"
|
5 |
import spaces
|
|
|
6 |
import cv2
|
7 |
import gradio as gr
|
8 |
import torch
|
@@ -61,16 +62,23 @@ def generate_reconstructions(pmrf_model, x, y, non_noisy_z0, num_flow_steps, dev
|
|
61 |
def enhance_face(img, face_helper, has_aligned, num_flow_steps, only_center_face=False, paste_back=True, scale=2):
|
62 |
face_helper.clean_all()
|
63 |
|
64 |
-
if has_aligned:
|
65 |
-
|
|
|
|
|
|
|
|
|
66 |
face_helper.cropped_faces = [img]
|
67 |
else:
|
68 |
face_helper.read_image(img)
|
69 |
-
|
70 |
-
|
71 |
-
|
|
|
|
|
72 |
# align and warp each face
|
73 |
face_helper.align_warp_face()
|
|
|
74 |
# face restoration
|
75 |
for cropped_face in face_helper.cropped_faces:
|
76 |
# prepare data
|
@@ -96,7 +104,7 @@ def enhance_face(img, face_helper, has_aligned, num_flow_steps, only_center_face
|
|
96 |
|
97 |
face_helper.get_inverse_affine(None)
|
98 |
# paste each restored face to the input image
|
99 |
-
restored_img = face_helper.paste_faces_to_input_image(upsample_img=bg_img)
|
100 |
return face_helper.cropped_faces, face_helper.restored_faces, restored_img
|
101 |
else:
|
102 |
return face_helper.cropped_faces, face_helper.restored_faces, None
|
@@ -161,6 +169,28 @@ You may use this demo to enhance the quality of any image which contains faces.
|
|
161 |
<b>NOTE</b>: Our model is designed to restore aligned face images, but here we incorporate mechanisms that allow restoring the quality of any image that contains any number of faces. Thus, the resulting quality of such general images is not guaranteed.
|
162 |
"""
|
163 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
164 |
css = r"""
|
165 |
"""
|
166 |
|
@@ -175,7 +205,8 @@ demo = gr.Interface(
|
|
175 |
gr.File(label="Download the output image")
|
176 |
],
|
177 |
title=title,
|
178 |
-
description=description
|
|
|
179 |
)
|
180 |
|
181 |
|
|
|
3 |
os.environ['SPACES_ZERO_GPU'] = "1"
|
4 |
os.environ['K_DIFFUSION_USE_COMPILE'] = "0"
|
5 |
import spaces
|
6 |
+
from facelib.utils.misc import is_gray
|
7 |
import cv2
|
8 |
import gradio as gr
|
9 |
import torch
|
|
|
62 |
def enhance_face(img, face_helper, has_aligned, num_flow_steps, only_center_face=False, paste_back=True, scale=2):
|
63 |
face_helper.clean_all()
|
64 |
|
65 |
+
if has_aligned:
|
66 |
+
# the input faces are already cropped and aligned
|
67 |
+
img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
|
68 |
+
face_helper.is_gray = is_gray(img, threshold=5)
|
69 |
+
if face_helper.is_gray:
|
70 |
+
print('\tgrayscale input: True')
|
71 |
face_helper.cropped_faces = [img]
|
72 |
else:
|
73 |
face_helper.read_image(img)
|
74 |
+
# get face landmarks for each face
|
75 |
+
num_det_faces = face_helper.get_face_landmarks_5(
|
76 |
+
only_center_face=only_center_face, resize=640, eye_dist_threshold=5
|
77 |
+
)
|
78 |
+
print(f'\tdetect {num_det_faces} faces')
|
79 |
# align and warp each face
|
80 |
face_helper.align_warp_face()
|
81 |
+
|
82 |
# face restoration
|
83 |
for cropped_face in face_helper.cropped_faces:
|
84 |
# prepare data
|
|
|
104 |
|
105 |
face_helper.get_inverse_affine(None)
|
106 |
# paste each restored face to the input image
|
107 |
+
restored_img = face_helper.paste_faces_to_input_image(upsample_img=bg_img, draw_box=False)
|
108 |
return face_helper.cropped_faces, face_helper.restored_faces, restored_img
|
109 |
else:
|
110 |
return face_helper.cropped_faces, face_helper.restored_faces, None
|
|
|
169 |
<b>NOTE</b>: Our model is designed to restore aligned face images, but here we incorporate mechanisms that allow restoring the quality of any image that contains any number of faces. Thus, the resulting quality of such general images is not guaranteed.
|
170 |
"""
|
171 |
|
172 |
+
|
173 |
+
article = r"""
|
174 |
+
If you find our work useful, please help to ⭐ our <a href='https://github.com/ohayonguy/PMRF' target='_blank'>GitHub repository</a>. Thanks!
|
175 |
+
[![GitHub Stars](https://img.shields.io/github/stars/ohayonguy/PMRF?style=social)](https://github.com/ohayonguy/PMRF)
|
176 |
+
---
|
177 |
+
📝 **Citation**
|
178 |
+
If our work is useful for your research, please consider citing:
|
179 |
+
```bibtex
|
180 |
+
@article{ohayon2024pmrf,
|
181 |
+
author = {Guy Ohayon and Tomer Michaeli and Michael Elad},
|
182 |
+
title = {Posterior-Mean Rectified Flow: Towards Minimum MSE Photo-Realistic Image Restoration},
|
183 |
+
journal = {arXiv preprint arXiv:2410.00418},
|
184 |
+
year = {2024},
|
185 |
+
url = {https://arxiv.org/abs/2410.00418}
|
186 |
+
}
|
187 |
+
```
|
188 |
+
📋 **License**
|
189 |
+
This project is released under the <a rel="license" href="https://github.com/ohayonguy/PMRF/blob/master/LICENSE">MIT license</a>.
|
190 |
+
Redistribution and use for non-commercial purposes should follow this license.
|
191 |
+
📧 **Contact**
|
192 |
+
If you have any questions, please feel free to contact me at <b>guyoep@gmail.com</b>.
|
193 |
+
"""
|
194 |
css = r"""
|
195 |
"""
|
196 |
|
|
|
205 |
gr.File(label="Download the output image")
|
206 |
],
|
207 |
title=title,
|
208 |
+
description=description,
|
209 |
+
article=article,
|
210 |
)
|
211 |
|
212 |
|