| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788 | from __future__ import divisionimport collectionsimport numpy as npimport globimport osimport os.path as ospimport cv2from insightface.model_zoo import model_zoofrom 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]
 |