victorisgeek commited on
Commit
12c456b
·
verified ·
1 Parent(s): a97be07

Delete face_enhance

Browse files
face_enhance/8 DELETED
@@ -1 +0,0 @@
1
- o
 
 
face_enhance/__init__.py DELETED
@@ -1 +0,0 @@
1
- from .gfpgan import GFPGAN
 
 
face_enhance/gfpgan.py DELETED
@@ -1,148 +0,0 @@
1
- import numpy as np
2
-
3
- import cv2
4
-
5
- from insightface.utils import face_align
6
- from insightface import model_zoo
7
- from dofaker.utils import download_file, get_model_url
8
-
9
-
10
- class GFPGAN:
11
-
12
- def __init__(self, name='gfpgan', root='weights/models') -> None:
13
- _, model_file = download_file(get_model_url(name),
14
- save_dir=root,
15
- overwrite=False)
16
- providers = model_zoo.model_zoo.get_default_providers()
17
- self.session = model_zoo.model_zoo.PickableInferenceSession(
18
- model_file, providers=providers)
19
-
20
- self.input_mean = 127.5
21
- self.input_std = 127.5
22
- inputs = self.session.get_inputs()
23
- self.input_names = []
24
- for inp in inputs:
25
- self.input_names.append(inp.name)
26
- outputs = self.session.get_outputs()
27
- output_names = []
28
- for out in outputs:
29
- output_names.append(out.name)
30
- self.output_names = output_names
31
- assert len(
32
- self.output_names
33
- ) == 1, "The output number of GFPGAN model should be 1, but got {}, please check your model.".format(
34
- len(self.output_names))
35
- output_shape = outputs[0].shape
36
- input_cfg = inputs[0]
37
- input_shape = input_cfg.shape
38
- self.input_shape = input_shape
39
- print('face_enhance-shape:', self.input_shape)
40
- self.input_size = tuple(input_shape[2:4][::-1])
41
-
42
- def forward(self, image, image_format='bgr'):
43
- if isinstance(image, str):
44
- image = cv2.imread(image, 1)
45
- elif isinstance(image, np.ndarray):
46
- if image_format == 'bgr':
47
- image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
48
- elif image_format == 'rgb':
49
- pass
50
- else:
51
- raise UserWarning(
52
- "gfpgan not support image format {}".format(image_format))
53
- else:
54
- raise UserWarning(
55
- "gfpgan input must be str or np.ndarray, but got {}.".format(
56
- type(image)))
57
- img = (image - self.input_mean) / self.input_std
58
- pred = self.session.run(self.output_names,
59
- {self.input_names[0]: img})[0]
60
- return pred
61
-
62
- def _get(self, img, image_format='bgr'):
63
- if image_format.lower() == 'bgr':
64
- img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
65
- elif image_format.lower() == 'rgb':
66
- pass
67
- else:
68
- raise UserWarning(
69
- "gfpgan not support image format {}".format(image_format))
70
- h, w, c = img.shape
71
- img = cv2.resize(img, (self.input_shape[-1], self.input_shape[-2]))
72
- blob = cv2.dnn.blobFromImage(
73
- img,
74
- 1.0 / self.input_std,
75
- self.input_size,
76
- (self.input_mean, self.input_mean, self.input_mean),
77
- swapRB=False)
78
- pred = self.session.run(self.output_names,
79
- {self.input_names[0]: blob})[0]
80
- image_aug = pred.transpose((0, 2, 3, 1))[0]
81
- rgb_aug = np.clip(self.input_std * image_aug + self.input_mean, 0,
82
- 255).astype(np.uint8)
83
- rgb_aug = cv2.resize(rgb_aug, (w, h))
84
- bgr_image = rgb_aug[:, :, ::-1]
85
- return bgr_image
86
-
87
- def get(self, img, target_face, paste_back=True, image_format='bgr'):
88
- aimg, M = face_align.norm_crop2(img, target_face.kps,
89
- self.input_size[0])
90
- bgr_fake = self._get(aimg, image_format='bgr')
91
- if not paste_back:
92
- return bgr_fake, M
93
- else:
94
- target_img = img
95
- fake_diff = bgr_fake.astype(np.float32) - aimg.astype(np.float32)
96
- fake_diff = np.abs(fake_diff).mean(axis=2)
97
- fake_diff[:2, :] = 0
98
- fake_diff[-2:, :] = 0
99
- fake_diff[:, :2] = 0
100
- fake_diff[:, -2:] = 0
101
- IM = cv2.invertAffineTransform(M)
102
- img_white = np.full((aimg.shape[0], aimg.shape[1]),
103
- 255,
104
- dtype=np.float32)
105
- bgr_fake = cv2.warpAffine(
106
- bgr_fake,
107
- IM, (target_img.shape[1], target_img.shape[0]),
108
- borderValue=0.0)
109
- img_white = cv2.warpAffine(
110
- img_white,
111
- IM, (target_img.shape[1], target_img.shape[0]),
112
- borderValue=0.0)
113
- fake_diff = cv2.warpAffine(
114
- fake_diff,
115
- IM, (target_img.shape[1], target_img.shape[0]),
116
- borderValue=0.0)
117
- img_white[img_white > 20] = 255
118
- fthresh = 10
119
- fake_diff[fake_diff < fthresh] = 0
120
- fake_diff[fake_diff >= fthresh] = 255
121
- img_mask = img_white
122
- mask_h_inds, mask_w_inds = np.where(img_mask == 255)
123
- mask_h = np.max(mask_h_inds) - np.min(mask_h_inds)
124
- mask_w = np.max(mask_w_inds) - np.min(mask_w_inds)
125
- mask_size = int(np.sqrt(mask_h * mask_w))
126
- k = max(mask_size // 10, 10)
127
- #k = max(mask_size//20, 6)
128
- #k = 6
129
- kernel = np.ones((k, k), np.uint8)
130
- img_mask = cv2.erode(img_mask, kernel, iterations=1)
131
- kernel = np.ones((2, 2), np.uint8)
132
- fake_diff = cv2.dilate(fake_diff, kernel, iterations=1)
133
- k = max(mask_size // 20, 5)
134
- kernel_size = (k, k)
135
- blur_size = tuple(2 * i + 1 for i in kernel_size)
136
- img_mask = cv2.GaussianBlur(img_mask, blur_size, 0)
137
- k = 5
138
- kernel_size = (k, k)
139
- blur_size = tuple(2 * i + 1 for i in kernel_size)
140
- fake_diff = cv2.GaussianBlur(fake_diff, blur_size, 0)
141
- img_mask /= 255
142
- fake_diff /= 255
143
- img_mask = np.reshape(img_mask,
144
- [img_mask.shape[0], img_mask.shape[1], 1])
145
- fake_merged = img_mask * bgr_fake + (
146
- 1 - img_mask) * target_img.astype(np.float32)
147
- fake_merged = fake_merged.astype(np.uint8)
148
- return fake_merged