Spaces:
Running
Running
Delete face_enhancer.py
Browse files- face_enhancer.py +0 -72
face_enhancer.py
DELETED
@@ -1,72 +0,0 @@
|
|
1 |
-
import os
|
2 |
-
import cv2
|
3 |
-
import torch
|
4 |
-
import gfpgan
|
5 |
-
from PIL import Image
|
6 |
-
from upscaler.RealESRGAN import RealESRGAN
|
7 |
-
from upscaler.codeformer import CodeFormerEnhancer
|
8 |
-
|
9 |
-
def gfpgan_runner(img, model):
|
10 |
-
_, imgs, _ = model.enhance(img, paste_back=True, has_aligned=True)
|
11 |
-
return imgs[0]
|
12 |
-
|
13 |
-
|
14 |
-
def realesrgan_runner(img, model):
|
15 |
-
img = model.predict(img)
|
16 |
-
return img
|
17 |
-
|
18 |
-
|
19 |
-
def codeformer_runner(img, model):
|
20 |
-
img = model.enhance(img)
|
21 |
-
return img
|
22 |
-
|
23 |
-
|
24 |
-
supported_enhancers = {
|
25 |
-
"CodeFormer": ("./assets/pretrained_models/codeformer.onnx", codeformer_runner),
|
26 |
-
"GFPGAN": ("./assets/pretrained_models/GFPGANv1.4.pth", gfpgan_runner),
|
27 |
-
"REAL-ESRGAN 2x": ("./assets/pretrained_models/RealESRGAN_x2.pth", realesrgan_runner),
|
28 |
-
"REAL-ESRGAN 4x": ("./assets/pretrained_models/RealESRGAN_x4.pth", realesrgan_runner),
|
29 |
-
"REAL-ESRGAN 8x": ("./assets/pretrained_models/RealESRGAN_x8.pth", realesrgan_runner)
|
30 |
-
}
|
31 |
-
|
32 |
-
cv2_interpolations = ["LANCZOS4", "CUBIC", "NEAREST"]
|
33 |
-
|
34 |
-
def get_available_enhancer_names():
|
35 |
-
available = []
|
36 |
-
for name, data in supported_enhancers.items():
|
37 |
-
path = os.path.join(os.path.abspath(os.path.dirname(__file__)), data[0])
|
38 |
-
if os.path.exists(path):
|
39 |
-
available.append(name)
|
40 |
-
return available
|
41 |
-
|
42 |
-
|
43 |
-
def load_face_enhancer_model(name='GFPGAN', device="cpu"):
|
44 |
-
assert name in get_available_enhancer_names() + cv2_interpolations, f"Face enhancer {name} unavailable."
|
45 |
-
if name in supported_enhancers.keys():
|
46 |
-
model_path, model_runner = supported_enhancers.get(name)
|
47 |
-
model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model_path)
|
48 |
-
if name == 'CodeFormer':
|
49 |
-
model = CodeFormerEnhancer(model_path=model_path, device=device)
|
50 |
-
elif name == 'GFPGAN':
|
51 |
-
model = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=device)
|
52 |
-
elif name == 'REAL-ESRGAN 2x':
|
53 |
-
model = RealESRGAN(device, scale=2)
|
54 |
-
model.load_weights(model_path, download=False)
|
55 |
-
elif name == 'REAL-ESRGAN 4x':
|
56 |
-
model = RealESRGAN(device, scale=4)
|
57 |
-
model.load_weights(model_path, download=False)
|
58 |
-
elif name == 'REAL-ESRGAN 8x':
|
59 |
-
model = RealESRGAN(device, scale=8)
|
60 |
-
model.load_weights(model_path, download=False)
|
61 |
-
elif name == 'LANCZOS4':
|
62 |
-
model = None
|
63 |
-
model_runner = lambda img, _: cv2.resize(img, (512,512), interpolation=cv2.INTER_LANCZOS4)
|
64 |
-
elif name == 'CUBIC':
|
65 |
-
model = None
|
66 |
-
model_runner = lambda img, _: cv2.resize(img, (512,512), interpolation=cv2.INTER_CUBIC)
|
67 |
-
elif name == 'NEAREST':
|
68 |
-
model = None
|
69 |
-
model_runner = lambda img, _: cv2.resize(img, (512,512), interpolation=cv2.INTER_NEAREST)
|
70 |
-
else:
|
71 |
-
model = None
|
72 |
-
return (model, model_runner)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|