Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.

1
from __future__ import division
2
import collections
3
import numpy as np
4
import glob
5
import os
6
import os.path as osp
7
import cv2
8
from insightface.model_zoo import model_zoo
9
from insightface.utils import face_align
10

11
__all__ = ['Face_detect_crop', 'Face']
12

13
Face = collections.namedtuple('Face', [
14
    'bbox', 'kps', 'det_score', 'embedding', 'gender', 'age',
15
    'embedding_norm', 'normed_embedding',
16
    'landmark'
17
])
18

19
Face.__new__.__defaults__ = (None, ) * len(Face._fields)
20

21

22
class Face_detect_crop:
23
    def __init__(self, name, root='~/.insightface_func/models'):
24
        self.models = {}
25
        root = os.path.expanduser(root)
26
        onnx_files = glob.glob(osp.join(root, name, '*.onnx'))
27
        onnx_files = sorted(onnx_files)
28
        for onnx_file in onnx_files:
29
            if onnx_file.find('_selfgen_')>0:
30
                #print('ignore:', onnx_file)
31
                continue
32
            model = model_zoo.get_model(onnx_file)
33
            if model.taskname not in self.models:
34
                print('find model:', onnx_file, model.taskname)
35
                self.models[model.taskname] = model
36
            else:
37
                print('duplicated model task type, ignore:', onnx_file, model.taskname)
38
                del model
39
        assert 'detection' in self.models
40
        self.det_model = self.models['detection']
41

42

43
    def prepare(self, ctx_id, det_thresh=0.5, det_size=(640, 640)):
44
        self.det_thresh = det_thresh
45
        assert det_size is not None
46
        print('set det-size:', det_size)
47
        self.det_size = det_size
48
        for taskname, model in self.models.items():
49
            if taskname=='detection':
50
                model.prepare(ctx_id, input_size=det_size)
51
            else:
52
                model.prepare(ctx_id)
53

54
    def get(self, img, crop_size, max_num=0):
55
        bboxes, kpss = self.det_model.detect(img,
56
                                             threshold=self.det_thresh,
57
                                             max_num=max_num,
58
                                             metric='default')
59
        if bboxes.shape[0] == 0:
60
            return None
61
        ret = []
62
        # for i in range(bboxes.shape[0]):
63
        #     bbox = bboxes[i, 0:4]
64
        #     det_score = bboxes[i, 4]
65
        #     kps = None
66
        #     if kpss is not None:
67
        #         kps = kpss[i]
68
        #     M, _ = face_align.estimate_norm(kps, crop_size, mode ='None') 
69
        #     align_img = cv2.warpAffine(img, M, (crop_size, crop_size), borderValue=0.0)
70
        # align_img_list = []
71
        # M_list = []
72
        kps_list = []
73
        for i in range(bboxes.shape[0]):
74
            kps = None
75
            if kpss is not None:
76
                kps = kpss[i]
77
            # M, _ = face_align.estimate_norm(kps, crop_size, mode ='None') 
78
            # align_img = cv2.warpAffine(img, M, (crop_size, crop_size), borderValue=0.0)
79
            # align_img_list.append(align_img)
80
            # M_list.append(M)
81
            kps_list.append(kps)
82

83
        # det_score = bboxes[..., 4]
84

85
        # best_index = np.argmax(det_score)
86

87
        # kps = None
88
        # if kpss is not None:
89
        #     kps = kpss[best_index]
90
        # M, _ = face_align.estimate_norm(kps, crop_size, mode ='None') 
91
        # align_img = cv2.warpAffine(img, M, (crop_size, crop_size), borderValue=0.0)
92
        
93
        return kps_list
94

95