Upload 2 files

dofaker/face_det/__init__.py

@@ -0,0 +1 @@
+from .face_analysis import FaceAnalysis

dofaker/face_det/face_analysis.py

@@ -0,0 +1,107 @@
+'''
+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