CalvinBackup/neuralchen-SimSwap
1
0
Fork 0
Code Issues 288 Pull Requests 16 Actions Packages Projects Releases 2 Wiki Activity

Insightface: used to detect and align faces

Browse Source
This commit is contained in:
NNNNAI
2021-06-20 11:58:03 +08:00
parent 29f9bb6d0e
commit 6e5ee6a3f6
3 changed files with 179 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
insightface_func/__init__.py
View File
insightface_func/face_detect_crop_mutil.py
+91
View File
@@ -0,0 +1,91 @@
from __future__ import division
import collections
import numpy as np
import glob
import os
import os.path as osp
import cv2
from insightface.model_zoo import model_zoo
from insightface.utils import face_align
__all__ = ['Face_detect_crop', 'Face']
Face = collections.namedtuple('Face', [
'bbox', 'kps', 'det_score', 'embedding', 'gender', 'age',
'embedding_norm', 'normed_embedding',
'landmark'
])
Face.__new__.__defaults__ = (None, ) * len(Face._fields)
class Face_detect_crop:
def __init__(self, name, root='~/.insightface_func/models'):
self.models = {}
root = os.path.expanduser(root)
onnx_files = glob.glob(osp.join(root, name, '*.onnx'))
onnx_files = sorted(onnx_files)
for onnx_file in onnx_files:
if onnx_file.find('_selfgen_')>0:
#print('ignore:', onnx_file)
continue
model = model_zoo.get_model(onnx_file)
if model.taskname not in self.models:
print('find model:', onnx_file, model.taskname)
self.models[model.taskname] = model
else:
print('duplicated model task type, ignore:', onnx_file, model.taskname)
del model
assert 'detection' in self.models
self.det_model = self.models['detection']
def prepare(self, ctx_id, det_thresh=0.5, det_size=(640, 640)):
self.det_thresh = det_thresh
assert det_size is not None
print('set det-size:', det_size)
self.det_size = det_size
for taskname, model in self.models.items():
if taskname=='detection':
model.prepare(ctx_id, input_size=det_size)
else:
model.prepare(ctx_id)
def get(self, img, crop_size, max_num=0):
bboxes, kpss = self.det_model.detect(img,
threshold=self.det_thresh,
max_num=max_num,
metric='default')
if bboxes.shape[0] == 0:
return None
ret = []
# for i in range(bboxes.shape[0]):
# bbox = bboxes[i, 0:4]
# det_score = bboxes[i, 4]
# kps = None
# if kpss is not None:
# kps = kpss[i]
# M, _ = face_align.estimate_norm(kps, crop_size, mode ='None')
# align_img = cv2.warpAffine(img, M, (crop_size, crop_size), borderValue=0.0)
align_img_list = []
M_list = []
for i in range(bboxes.shape[0]):
kps = None
if kpss is not None:
kps = kpss[i]
M, _ = face_align.estimate_norm(kps, crop_size, mode ='None')
align_img = cv2.warpAffine(img, M, (crop_size, crop_size), borderValue=0.0)
align_img_list.append(align_img)
M_list.append(M)
# det_score = bboxes[..., 4]
# best_index = np.argmax(det_score)
# kps = None
# if kpss is not None:
# kps = kpss[best_index]
# M, _ = face_align.estimate_norm(kps, crop_size, mode ='None')
# align_img = cv2.warpAffine(img, M, (crop_size, crop_size), borderValue=0.0)
return align_img_list, M_list
insightface_func/face_detect_crop_single.py
+88
View File
@@ -0,0 +1,88 @@
from __future__ import division
import collections
import numpy as np
import glob
import os
import os.path as osp
import cv2
from insightface.model_zoo import model_zoo
from insightface.utils import face_align
__all__ = ['Face_detect_crop', 'Face']
Face = collections.namedtuple('Face', [
'bbox', 'kps', 'det_score', 'embedding', 'gender', 'age',
'embedding_norm', 'normed_embedding',
'landmark'
])
Face.__new__.__defaults__ = (None, ) * len(Face._fields)
class Face_detect_crop:
def __init__(self, name, root='~/.insightface_func/models'):
self.models = {}
root = os.path.expanduser(root)
onnx_files = glob.glob(osp.join(root, name, '*.onnx'))
onnx_files = sorted(onnx_files)
for onnx_file in onnx_files:
if onnx_file.find('_selfgen_')>0:
#print('ignore:', onnx_file)
continue
model = model_zoo.get_model(onnx_file)
if model.taskname not in self.models:
print('find model:', onnx_file, model.taskname)
self.models[model.taskname] = model
else:
print('duplicated model task type, ignore:', onnx_file, model.taskname)
del model
assert 'detection' in self.models
self.det_model = self.models['detection']
def prepare(self, ctx_id, det_thresh=0.5, det_size=(640, 640)):
self.det_thresh = det_thresh
assert det_size is not None
print('set det-size:', det_size)
self.det_size = det_size
for taskname, model in self.models.items():
if taskname=='detection':
model.prepare(ctx_id, input_size=det_size)
else:
model.prepare(ctx_id)
def get(self, img, crop_size, max_num=0):
bboxes, kpss = self.det_model.detect(img,
threshold=self.det_thresh,
max_num=max_num,
metric='default')
if bboxes.shape[0] == 0:
return None
# ret = []
# for i in range(bboxes.shape[0]):
# bbox = bboxes[i, 0:4]
# det_score = bboxes[i, 4]
# kps = None
# if kpss is not None:
# kps = kpss[i]
# M, _ = face_align.estimate_norm(kps, crop_size, mode ='None')
# align_img = cv2.warpAffine(img, M, (crop_size, crop_size), borderValue=0.0)
# for i in range(bboxes.shape[0]):
# kps = None
# if kpss is not None:
# kps = kpss[i]
# M, _ = face_align.estimate_norm(kps, crop_size, mode ='None')
# align_img = cv2.warpAffine(img, M, (crop_size, crop_size), borderValue=0.0)
det_score = bboxes[..., 4]
# select the face with the hightest detection score
best_index = np.argmax(det_score)
kps = None
if kpss is not None:
kps = kpss[best_index]
M, _ = face_align.estimate_norm(kps, crop_size, mode ='None')
align_img = cv2.warpAffine(img, M, (crop_size, crop_size), borderValue=0.0)
return [align_img], [M]