victorisgeek commited on
Commit
45fce2a
1 Parent(s): 2262586

Delete face_enhancer.py

Browse files
Files changed (1) hide show
  1. 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)