Spaces:
Running
Running
File size: 4,109 Bytes
6c49922 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
'''
The following code references:: https://github.com/deepinsight/insightface
'''
import glob
import os.path as osp
import onnxruntime
from insightface import model_zoo
from insightface.utils import ensure_available
from insightface.app.common import Face
from dofaker.utils import download_file, get_model_url
__all__ = ['FaceAnalysis']
class FaceAnalysis:
def __init__(self,
name='buffalo_l',
root='weights',
allowed_modules=None,
**kwargs):
self.model_dir, _ = download_file(get_model_url(name),
save_dir=root,
overwrite=False)
onnxruntime.set_default_logger_severity(3)
self.models = {}
print(self.model_dir)
onnx_files = glob.glob(osp.join(self.model_dir, '*.onnx'))
onnx_files = sorted(onnx_files)
for onnx_file in onnx_files:
model = model_zoo.get_model(onnx_file, **kwargs)
if model is None:
print('model not recognized:', onnx_file)
elif allowed_modules is not None and model.taskname not in allowed_modules:
print('model ignore:', onnx_file, model.taskname)
del model
elif model.taskname not in self.models and (allowed_modules is None
or model.taskname
in allowed_modules):
print('find model:', onnx_file, model.taskname,
model.input_shape, model.input_mean, model.input_std)
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, "det_size can't be None."
self.det_size = det_size
for taskname, model in self.models.items():
if taskname == 'detection':
model.prepare(ctx_id,
input_size=det_size,
det_thresh=det_thresh)
else:
model.prepare(ctx_id)
def get(self, img, max_num=0):
bboxes, kpss = self.det_model.detect(img,
max_num=max_num,
metric='default')
if bboxes.shape[0] == 0:
return []
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]
face = Face(bbox=bbox, kps=kps, det_score=det_score)
for taskname, model in self.models.items():
if taskname == 'detection':
continue
model.get(img, face)
ret.append(face)
return ret
def draw_on(self, img, faces):
import cv2
dimg = img.copy()
for i in range(len(faces)):
face = faces[i]
box = face.bbox.astype('int')
color = (0, 0, 255)
cv2.rectangle(dimg, (box[0], box[1]), (box[2], box[3]), color, 2)
if face.kps is not None:
kps = face.kps.astype('int')
for l in range(kps.shape[0]):
color = (0, 0, 255)
if l == 0 or l == 3:
color = (0, 255, 0)
cv2.circle(dimg, (kps[l][0], kps[l][1]), 1, color, 2)
if face.gender is not None and face.age is not None:
cv2.putText(dimg, '%s,%d' % (face.sex, face.age),
(box[0] - 1, box[1] - 4), cv2.FONT_HERSHEY_COMPLEX,
0.7, (0, 255, 0), 1)
return dimg
|