|
|
''' |
|
|
Author: Naiyuan liu |
|
|
Github: https://github.com/NNNNAI |
|
|
Date: 2021-11-23 17:03:58 |
|
|
LastEditors: Naiyuan liu |
|
|
LastEditTime: 2021-11-24 16:46:04 |
|
|
Description: |
|
|
''' |
|
|
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_func.utils import face_align_ffhqandnewarc as 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: |
|
|
|
|
|
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), mode ='None'): |
|
|
self.det_thresh = det_thresh |
|
|
self.mode = mode |
|
|
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 |
|
|
|
|
|
det_score = bboxes[..., 4] |
|
|
best_index = np.argmax(det_score) |
|
|
|
|
|
kps = None |
|
|
if kpss is not None: |
|
|
kps = kpss[best_index] |
|
|
|
|
|
|
|
|
|
|
|
h, w, _ = img.shape |
|
|
face_h = bboxes[best_index][3] - bboxes[best_index][1] |
|
|
kps[:,1] += int(0.1 * face_h) |
|
|
|
|
|
|
|
|
M, _ = face_align.estimate_norm(kps, crop_size, mode=self.mode) |
|
|
align_img = cv2.warpAffine(img, M, (crop_size, crop_size), borderValue=0.0) |
|
|
|
|
|
return [align_img], [M] |
|
|
|
