Upload folder using huggingface_hub

.gitattributes

@@ -34,3 +34,12 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 *models/buffalo_l/det_10g.axmodel filter=lfs diff=lfs merge=lfs -text
+insightface/data/images/t1.jpg filter=lfs diff=lfs merge=lfs -text
+models/buffalo_l/1k3d68.axmodel filter=lfs diff=lfs merge=lfs -text
+models/buffalo_l/2d106det.axmodel filter=lfs diff=lfs merge=lfs -text
+models/buffalo_l/genderage.axmodel filter=lfs diff=lfs merge=lfs -text
+models/buffalo_l/w600k_r50.axmodel filter=lfs diff=lfs merge=lfs -text
+output/kobe0.jpg filter=lfs diff=lfs merge=lfs -text
+output/ssd_horse.jpg filter=lfs diff=lfs merge=lfs -text
+query/ssd_horse.jpg filter=lfs diff=lfs merge=lfs -text
+query/wade2.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,110 @@
----
-license: bsd-3-clause
----
+# Insightface
+基于Insightface的人脸识别Pipeline，使用其buffalo_l系列模型中的人脸检测+定点模型，人脸属性模型，人脸特征提取模型实现人脸识别，支持650N系列平台的1v1人脸比对和1vN人脸识别。
+
+支持芯片:
+- AX650N
+
+支持硬件
+
+  - [M4N-Dock(爱芯派Pro)](https://wiki.sipeed.com/hardware/zh/maixIV/m4ndock/m4ndock.html)
+  - [M.2 Accelerator card](https://axcl-docs.readthedocs.io/zh-cn/latest/doc_guide_hardware.html)
+
+原始模型请参考
+- [Insightface Github](https://github.com/deepinsight/insightface)
+
+## 性能对比
+
+| Models        |input size        | Latency (ms) | CMM Usage (MB) |
+| --------------|------- | ---------------------- | -------------- |
+| det_10g       |640x640 |6.947                    | 14.0          |
+| genderage     |96x96 |0.295                    | 0.62          |
+| w600k_r50     |112x112|3.993                    | 44.0          |
+
+## 模型转换
+- 模型转换工具链[Pulsar2](https://huggingface.co/AXERA-TECH/Pulsar2)
+- 转换文档[TODO]
+
+## 环境准备
+- NPU Python API: [pyaxengine](https://github.com/AXERA-TECH/pyaxengine)
+
+安装需要的python库
+```pip install -r requirements.txt```
+
+## 运行（算力卡版本）
+
+```bash
+(insightface) axera@dell:~/insightface/ax$ python insightface_pipeline.py -h
+[INFO] Available providers:  ['AXCLRTExecutionProvider']
+usage: insightface_pipeline.py [-h] [--model_path MODEL_PATH] [--type TYPE] [--gallery_path GALLERY_PATH] [--query_path QUERY_PATH] [--draw]
+
+Face Recognition Pipeline Example
+
+options:
+  -h, --help            show this help message and exit
+  --model_path, -m MODEL_PATH
+                        Path to the model directory
+  --type, -t TYPE       Type of operation: 1: 1v1 compare, 2: 1vN recognize
+  --gallery_path, -g GALLERY_PATH
+                        Path to the gallery image for image file
+  --query_path, -q QUERY_PATH
+                        Path to the query image
+  --draw, -d            Whether to draw results on the image
+
+
+(insightface) axera@dell:~/insightface/ax$ python insightface_pipeline.py -g ./gallery_imgs/ -d -t 1
+[INFO] Available providers:  ['AXCLRTExecutionProvider']
+ModelRouter ./models/buffalo_l/1k3d68.axmodel
+[INFO] Using provider: AXCLRTExecutionProvider
+[INFO] SOC Name: AX650N
+[INFO] VNPU type: VNPUType.DISABLED
+[INFO] Compiler version: 5.0-patch1-dirty 5cf5ab99-dirty
+Applied providers: AXCLRTExecutionProvider, with options: None
+model not recognized: ./models/buffalo_l/1k3d68.axmodel
+ModelRouter ./models/buffalo_l/2d106det.axmodel
+[INFO] Using provider: AXCLRTExecutionProvider
+[INFO] SOC Name: AX650N
+[INFO] VNPU type: VNPUType.DISABLED
+[INFO] Compiler version: 5.0-patch1-dirty 5cf5ab99-dirty
+Applied providers: AXCLRTExecutionProvider, with options: None
+model not recognized: ./models/buffalo_l/2d106det.axmodel
+ModelRouter ./models/buffalo_l/det_10g.axmodel
+[INFO] Using provider: AXCLRTExecutionProvider
+[INFO] SOC Name: AX650N
+[INFO] VNPU type: VNPUType.DISABLED
+[INFO] Compiler version: 5.0-patch1-dirty 5cf5ab99-dirty
+Applied providers: AXCLRTExecutionProvider, with options: None
+find model: ./models/buffalo_l/det_10g.axmodel detection [1, 640, 640, 3] 0.0 1.0
+ModelRouter ./models/buffalo_l/genderage.axmodel
+[INFO] Using provider: AXCLRTExecutionProvider
+[INFO] SOC Name: AX650N
+[INFO] VNPU type: VNPUType.DISABLED
+[INFO] Compiler version: 5.0-patch1-dirty 5cf5ab99-dirty
+Applied providers: AXCLRTExecutionProvider, with options: None
+find model: ./models/buffalo_l/genderage.axmodel genderage [1, 96, 96, 3] 0.0 1.0
+ModelRouter ./models/buffalo_l/w600k_r50.axmodel
+[INFO] Using provider: AXCLRTExecutionProvider
+[INFO] SOC Name: AX650N
+[INFO] VNPU type: VNPUType.DISABLED
+[INFO] Compiler version: 5.0-patch1-dirty 5cf5ab99-dirty
+Applied providers: AXCLRTExecutionProvider, with options: None
+find model: ./models/buffalo_l/w600k_r50.axmodel recognition [1, 112, 112, 3] 0.0 1.0
+set det-size: (640, 640)
+warning: det_size is already set in detection model, ignore
+100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 12.49it/s]
+特征库构建完成，包含以下人员： ['wade', 'kobe']
+请输入查询图像路径 (输入 'exit' 退出):
+./query/wade2.jpg
+识别结果: wade, 相似度(0-1): 0.5860
+结果已保存到 ./output/wade2.jpg
+请输入查询图像路径 (输入 'exit' 退出):
+./query/ssd_horse.jpg
+识别结果: Unknown, 相似度(0-1): 0.0135
+结果已保存到 ./output/ssd_horse.jpg
+请输入查询图像路径 (输入 'exit' 退出):
+exit
+
+```
+
+保存结果在 `./output` 目录下：
+![image](./output/wade2.jpg)
+![image](./output/ssd_horse.jpg)

insightface/__init__.py

@@ -0,0 +1,21 @@
+# coding: utf-8
+# pylint: disable=wrong-import-position
+"""InsightFace: A Face Analysis Toolkit."""
+from __future__ import absolute_import
+
+# try:
+#     #import mxnet as mx
+#     import onnxruntime
+# except ImportError:
+#     raise ImportError(
+#         "Unable to import dependency onnxruntime. "
+#     )
+
+__version__ = '0.7.3'
+
+from . import model_zoo
+from . import utils
+from . import app
+from . import data
+from . import thirdparty
+

insightface/app/__init__.py

@@ -0,0 +1,2 @@
+from .face_analysis import *
+# from .mask_renderer import *

insightface/app/common.py

@@ -0,0 +1,49 @@
+import numpy as np
+from numpy.linalg import norm as l2norm
+#from easydict import EasyDict
+
+class Face(dict):
+
+    def __init__(self, d=None, **kwargs):
+        if d is None:
+            d = {}
+        if kwargs:
+            d.update(**kwargs)
+        for k, v in d.items():
+            setattr(self, k, v)
+        # Class attributes
+        #for k in self.__class__.__dict__.keys():
+        #    if not (k.startswith('__') and k.endswith('__')) and not k in ('update', 'pop'):
+        #        setattr(self, k, getattr(self, k))
+
+    def __setattr__(self, name, value):
+        if isinstance(value, (list, tuple)):
+            value = [self.__class__(x)
+                    if isinstance(x, dict) else x for x in value]
+        elif isinstance(value, dict) and not isinstance(value, self.__class__):
+            value = self.__class__(value)
+        super(Face, self).__setattr__(name, value)
+        super(Face, self).__setitem__(name, value)
+
+    __setitem__ = __setattr__
+
+    def __getattr__(self, name):
+        return None
+
+    @property
+    def embedding_norm(self):
+        if self.embedding is None:
+            return None
+        return l2norm(self.embedding)
+
+    @property 
+    def normed_embedding(self):
+        if self.embedding is None:
+            return None
+        return self.embedding / self.embedding_norm
+
+    @property 
+    def sex(self):
+        if self.gender is None:
+            return None
+        return 'M' if self.gender==1 else 'F'

insightface/app/face_analysis.py

@@ -0,0 +1,108 @@
+# -*- coding: utf-8 -*-
+# @Organization  : insightface.ai
+# @Author        : Jia Guo
+# @Time          : 2021-05-04
+# @Function      : 
+
+
+from __future__ import division
+
+import glob
+import os.path as osp
+
+import numpy as np
+from numpy.linalg import norm
+
+from ..model_zoo import model_zoo
+from ..utils import DEFAULT_MP_NAME, ensure_available
+from .common import Face
+
+__all__ = ['FaceAnalysis']
+
+class FaceAnalysis:
+    def __init__(self, name=DEFAULT_MP_NAME, root='./models/axmodel', allowed_modules=None, **kwargs):
+
+        self.models = {}
+        self.model_dir = root
+        onnx_files = glob.glob(osp.join(self.model_dir, '*.axmodel'))
+        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
+        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, det_thresh=det_thresh)
+            else:
+                model.prepare(ctx_id)
+
+    def get(self, img, max_num=0, det_metric='default'):
+        bboxes, kpss = self.det_model.detect(img,
+                                             max_num=max_num,
+                                             metric=det_metric)
+        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)
+                #print(landmark.shape)
+                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)
+
+            #for key, value in face.items():
+            #    if key.startswith('landmark_3d'):
+            #        print(key, value.shape)
+            #        print(value[0:10,:])
+            #        lmk = np.round(value).astype(int)
+            #        for l in range(lmk.shape[0]):
+            #            color = (255, 0, 0)
+            #            cv2.circle(dimg, (lmk[l][0], lmk[l][1]), 1, color,
+            #                       2)
+        return dimg
+

insightface/app/mask_renderer.py

@@ -0,0 +1,232 @@
+import os, sys, datetime
+import numpy as np
+import os.path as osp
+import albumentations as A
+from albumentations.core.transforms_interface import ImageOnlyTransform
+from .face_analysis import FaceAnalysis
+from ..utils import get_model_dir
+from ..thirdparty import face3d
+from ..data import get_image as ins_get_image
+from ..utils import DEFAULT_MP_NAME
+import cv2
+
+class MaskRenderer:
+    def __init__(self, name=DEFAULT_MP_NAME, root='~/.insightface', insfa=None):
+        #if insfa is None, enter render_only mode
+        self.mp_name = name
+        self.root = root
+        self.insfa = insfa
+        model_dir = get_model_dir(name, root)
+        bfm_file = osp.join(model_dir, 'BFM.mat')
+        assert osp.exists(bfm_file), 'should contains BFM.mat in your model directory'
+        self.bfm = face3d.morphable_model.MorphabelModel(bfm_file)
+        self.index_ind = self.bfm.kpt_ind
+        bfm_uv_file = osp.join(model_dir, 'BFM_UV.mat')
+        assert osp.exists(bfm_uv_file), 'should contains BFM_UV.mat in your model directory'
+        uv_coords = face3d.morphable_model.load.load_uv_coords(bfm_uv_file)
+        self.uv_size = (224,224)
+        self.mask_stxr =  0.1
+        self.mask_styr = 0.33
+        self.mask_etxr = 0.9
+        self.mask_etyr =  0.7
+        self.tex_h , self.tex_w, self.tex_c = self.uv_size[1] , self.uv_size[0],3
+        texcoord = np.zeros_like(uv_coords)
+        texcoord[:, 0] = uv_coords[:, 0] * (self.tex_h - 1)
+        texcoord[:, 1] = uv_coords[:, 1] * (self.tex_w - 1)
+        texcoord[:, 1] = self.tex_w - texcoord[:, 1] - 1
+        self.texcoord = np.hstack((texcoord, np.zeros((texcoord.shape[0], 1))))
+        self.X_ind = self.bfm.kpt_ind
+        self.mask_image_names = ['mask_white', 'mask_blue', 'mask_black', 'mask_green']
+        self.mask_aug_probs = [0.4, 0.4, 0.1, 0.1]
+        #self.mask_images = []
+        #self.mask_images_rgb = []
+        #for image_name in mask_image_names:
+        #    mask_image = ins_get_image(image_name)
+        #    self.mask_images.append(mask_image)
+        #    mask_image_rgb = mask_image[:,:,::-1]
+        #    self.mask_images_rgb.append(mask_image_rgb)
+
+
+    def prepare(self, ctx_id=0, det_thresh=0.5, det_size=(128, 128)):
+        self.pre_ctx_id = ctx_id
+        self.pre_det_thresh = det_thresh
+        self.pre_det_size = det_size
+
+    def transform(self, shape3D, R):
+        s = 1.0
+        shape3D[:2, :] = shape3D[:2, :]
+        shape3D = s * np.dot(R, shape3D)
+        return shape3D
+
+    def preprocess(self, vertices, w, h):
+        R1 = face3d.mesh.transform.angle2matrix([0, 180, 180])
+        t = np.array([-w // 2, -h // 2, 0])
+        vertices = vertices.T
+        vertices += t
+        vertices = self.transform(vertices.T, R1).T
+        return vertices
+
+    def project_to_2d(self,vertices,s,angles,t):
+        transformed_vertices = self.bfm.transform(vertices, s, angles, t)
+        projected_vertices = transformed_vertices.copy() # using stantard camera & orth projection
+        return projected_vertices[self.bfm.kpt_ind, :2]
+
+    def params_to_vertices(self,params  , H , W):
+        fitted_sp, fitted_ep, fitted_s, fitted_angles, fitted_t  = params
+        fitted_vertices = self.bfm.generate_vertices(fitted_sp, fitted_ep)
+        transformed_vertices = self.bfm.transform(fitted_vertices, fitted_s, fitted_angles,
+                                                  fitted_t)
+        transformed_vertices = self.preprocess(transformed_vertices.T, W, H)
+        image_vertices = face3d.mesh.transform.to_image(transformed_vertices, H, W)
+        return image_vertices
+
+    def draw_lmk(self, face_image):
+        faces = self.insfa.get(face_image, max_num=1)
+        if len(faces)==0:
+            return face_image
+        return self.insfa.draw_on(face_image, faces)
+
+    def build_params(self, face_image):
+        #landmark = self.if3d68_handler.get(face_image)
+        #if landmark is None:
+        #    return None #face not found
+        if self.insfa is None:
+            self.insfa = FaceAnalysis(name=self.mp_name, root=self.root, allowed_modules=['detection', 'landmark_3d_68'])
+            self.insfa.prepare(ctx_id=self.pre_ctx_id,  det_thresh=self.pre_det_thresh, det_size=self.pre_det_size)
+
+        faces = self.insfa.get(face_image, max_num=1)
+        if len(faces)==0:
+            return None
+        landmark = faces[0].landmark_3d_68[:,:2]
+        fitted_sp, fitted_ep, fitted_s, fitted_angles, fitted_t = self.bfm.fit(landmark, self.X_ind, max_iter = 3)
+        return [fitted_sp, fitted_ep, fitted_s, fitted_angles, fitted_t]
+
+    def generate_mask_uv(self,mask, positions):
+        uv_size = (self.uv_size[1], self.uv_size[0], 3)
+        h, w, c = uv_size
+        uv = np.zeros(shape=(self.uv_size[1],self.uv_size[0], 3), dtype=np.uint8)
+        stxr, styr  = positions[0], positions[1]
+        etxr, etyr = positions[2], positions[3]
+        stx, sty = int(w * stxr), int(h * styr)
+        etx, ety = int(w * etxr), int(h * etyr)
+        height = ety - sty
+        width = etx - stx
+        mask = cv2.resize(mask, (width, height))
+        uv[sty:ety, stx:etx] = mask
+        return uv
+
+    def render_mask(self,face_image, mask_image, params, input_is_rgb=False, auto_blend = True, positions=[0.1, 0.33, 0.9, 0.7]):
+        if isinstance(mask_image, str):
+            to_rgb = True if input_is_rgb else False
+            mask_image = ins_get_image(mask_image, to_rgb=to_rgb)
+        uv_mask_image = self.generate_mask_uv(mask_image, positions)
+        h,w,c = face_image.shape
+        image_vertices = self.params_to_vertices(params ,h,w)
+        output = (1-face3d.mesh.render.render_texture(image_vertices, self.bfm.full_triangles , uv_mask_image, self.texcoord, self.bfm.full_triangles, h , w ))*255
+        output = output.astype(np.uint8)
+        if auto_blend:
+            mask_bd = (output==255).astype(np.uint8)
+            final = face_image*mask_bd + (1-mask_bd)*output
+            return final
+        return output
+
+    #def mask_augmentation(self, face_image, label, input_is_rgb=False, p=0.1):
+    #    if np.random.random()<p:
+    #        assert isinstance(label, (list, np.ndarray)), 'make sure the rec dataset includes mask params'
+    #        assert len(label)==237 or len(lable)==235, 'make sure the rec dataset includes mask params'
+    #        if len(label)==237:
+    #            if label[1]<0.0: #invalid label for mask aug
+    #                return face_image
+    #            label = label[2:]
+    #        params = self.decode_params(label)
+    #        mask_image_name = np.random.choice(self.mask_image_names, p=self.mask_aug_probs)
+    #        pos = np.random.uniform(0.33, 0.5)
+    #        face_image = self.render_mask(face_image, mask_image_name, params, input_is_rgb=input_is_rgb, positions=[0.1, pos, 0.9, 0.7])
+    #    return face_image
+
+    @staticmethod
+    def encode_params(params):
+        p0 = list(params[0])
+        p1 = list(params[1])
+        p2 = [float(params[2])]
+        p3 = list(params[3])
+        p4 = list(params[4])
+        return p0+p1+p2+p3+p4
+
+    @staticmethod
+    def decode_params(params):
+        p0 = params[0:199]
+        p0 = np.array(p0, dtype=np.float32).reshape( (-1, 1))
+        p1 = params[199:228]
+        p1 = np.array(p1, dtype=np.float32).reshape( (-1, 1))
+        p2 = params[228]
+        p3 = tuple(params[229:232])
+        p4 = params[232:235]
+        p4 = np.array(p4, dtype=np.float32).reshape( (-1, 1))
+        return p0, p1, p2, p3, p4
+    
+class MaskAugmentation(ImageOnlyTransform):
+
+    def __init__(
+            self,
+            mask_names=['mask_white', 'mask_blue', 'mask_black', 'mask_green'],
+            mask_probs=[0.4,0.4,0.1,0.1],
+            h_low = 0.33,
+            h_high = 0.35,
+            always_apply=False,
+            p=1.0,
+            ):
+        super(MaskAugmentation, self).__init__(always_apply, p)
+        self.renderer = MaskRenderer()
+        assert len(mask_names)>0
+        assert len(mask_names)==len(mask_probs)
+        self.mask_names = mask_names
+        self.mask_probs = mask_probs
+        self.h_low = h_low
+        self.h_high = h_high
+        #self.hlabel = None
+
+
+    def apply(self, image, hlabel, mask_name, h_pos, **params):
+        #print(params.keys())
+        #hlabel = params.get('hlabel')
+        assert len(hlabel)==237 or len(hlabel)==235, 'make sure the rec dataset includes mask params'
+        if len(hlabel)==237:
+            if hlabel[1]<0.0:
+                return image
+            hlabel = hlabel[2:]
+        #print(len(hlabel))
+        mask_params = self.renderer.decode_params(hlabel)
+        image = self.renderer.render_mask(image, mask_name, mask_params, input_is_rgb=True, positions=[0.1, h_pos, 0.9, 0.7])
+        return image
+
+    @property
+    def targets_as_params(self):
+        return ["image", "hlabel"]
+
+    def get_params_dependent_on_targets(self, params):
+        hlabel = params['hlabel']
+        mask_name = np.random.choice(self.mask_names, p=self.mask_probs)
+        h_pos = np.random.uniform(self.h_low, self.h_high)
+        return {'hlabel': hlabel, 'mask_name': mask_name, 'h_pos': h_pos}
+
+    def get_transform_init_args_names(self):
+        #return ("hlabel", 'mask_names', 'mask_probs', 'h_low', 'h_high')
+        return ('mask_names', 'mask_probs', 'h_low', 'h_high')
+
+
+if __name__ == "__main__":
+    tool = MaskRenderer('antelope')
+    tool.prepare(det_size=(128,128))
+    image = cv2.imread("Tom_Hanks_54745.png")
+    params = tool.build_params(image)
+    #out = tool.draw_lmk(image)
+    #cv2.imwrite('output_lmk.jpg', out)
+    #mask_image  = cv2.imread("masks/mask1.jpg")
+    #mask_image  = cv2.imread("masks/black-mask.png")
+    #mask_image  = cv2.imread("masks/mask2.jpg")
+    mask_out = tool.render_mask(image, 'mask_blue', params)# use single thread to test the time cost
+
+    cv2.imwrite('output_mask.jpg', mask_out)
+
+

insightface/commands/__init__.py

@@ -0,0 +1,13 @@
+from abc import ABC, abstractmethod
+from argparse import ArgumentParser
+
+
+class BaseInsightFaceCLICommand(ABC):
+    @staticmethod
+    @abstractmethod
+    def register_subcommand(parser: ArgumentParser):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def run(self):
+        raise NotImplementedError()

insightface/commands/insightface_cli.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python
+
+from argparse import ArgumentParser
+
+from .model_download import ModelDownloadCommand
+from .rec_add_mask_param import RecAddMaskParamCommand
+
+def main():
+    parser = ArgumentParser("InsightFace CLI tool", usage="insightface-cli <command> [<args>]")
+    commands_parser = parser.add_subparsers(help="insightface-cli command-line helpers")
+
+    # Register commands
+    ModelDownloadCommand.register_subcommand(commands_parser)
+    RecAddMaskParamCommand.register_subcommand(commands_parser)
+
+    args = parser.parse_args()
+
+    if not hasattr(args, "func"):
+        parser.print_help()
+        exit(1)
+
+    # Run
+    service = args.func(args)
+    service.run()
+
+
+if __name__ == "__main__":
+    main()
+

insightface/commands/model_download.py

@@ -0,0 +1,36 @@
+from argparse import ArgumentParser
+
+from . import BaseInsightFaceCLICommand
+import os
+import os.path as osp
+import zipfile
+import glob
+from ..utils import download
+
+
+def model_download_command_factory(args):
+    return ModelDownloadCommand(args.model, args.root, args.force)
+
+
+class ModelDownloadCommand(BaseInsightFaceCLICommand):
+    #_url_format = '{repo_url}models/{file_name}.zip'
+    @staticmethod
+    def register_subcommand(parser: ArgumentParser):
+        download_parser = parser.add_parser("model.download")
+        download_parser.add_argument(
+            "--root", type=str, default='~/.insightface', help="Path to location to store the models"
+        )
+        download_parser.add_argument(
+            "--force", action="store_true", help="Force the model to be download even if already in root-dir"
+        )
+        download_parser.add_argument("model", type=str, help="Name of the model to download")
+        download_parser.set_defaults(func=model_download_command_factory)
+
+    def __init__(self, model: str, root: str, force: bool):
+        self._model = model
+        self._root = root
+        self._force = force
+
+    def run(self):
+        download('models', self._model, force=self._force, root=self._root)
+

insightface/commands/rec_add_mask_param.py

@@ -0,0 +1,94 @@
+
+import numbers
+import os
+from argparse import ArgumentParser, Namespace
+
+import mxnet as mx
+import numpy as np
+
+from ..app import MaskRenderer
+from ..data.rec_builder import RecBuilder
+from . import BaseInsightFaceCLICommand
+
+
+def rec_add_mask_param_command_factory(args: Namespace):
+
+    return RecAddMaskParamCommand(
+        args.input, args.output
+    )
+
+
+class RecAddMaskParamCommand(BaseInsightFaceCLICommand):
+    @staticmethod
+    def register_subcommand(parser: ArgumentParser):
+        _parser = parser.add_parser("rec.addmaskparam")
+        _parser.add_argument("input", type=str, help="input rec")
+        _parser.add_argument("output", type=str, help="output rec, with mask param")
+        _parser.set_defaults(func=rec_add_mask_param_command_factory)
+
+    def __init__(
+        self,
+        input: str,
+        output: str,
+    ):
+        self._input = input
+        self._output = output
+
+
+    def run(self):
+        tool = MaskRenderer()
+        tool.prepare(ctx_id=0, det_size=(128,128))
+        root_dir = self._input
+        path_imgrec = os.path.join(root_dir, 'train.rec')
+        path_imgidx = os.path.join(root_dir, 'train.idx')
+        imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r')
+        save_path = self._output
+        wrec=RecBuilder(path=save_path)
+        s = imgrec.read_idx(0)
+        header, _ = mx.recordio.unpack(s)
+        if header.flag > 0:
+            if len(header.label)==2:
+                imgidx = np.array(range(1, int(header.label[0])))
+            else:
+                imgidx = np.array(list(self.imgrec.keys))
+        else:
+            imgidx = np.array(list(self.imgrec.keys))
+        stat = [0, 0]
+        print('total:', len(imgidx))
+        for iid, idx in enumerate(imgidx):
+            #if iid==500000:
+            #    break
+            if iid%1000==0:
+                print('processing:', iid)
+            s = imgrec.read_idx(idx)
+            header, img = mx.recordio.unpack(s)
+            label = header.label
+            if not isinstance(label, numbers.Number):
+                label = label[0]
+            sample = mx.image.imdecode(img).asnumpy()
+            bgr = sample[:,:,::-1]
+            params = tool.build_params(bgr)
+            #if iid<10:
+            #    mask_out = tool.render_mask(bgr, 'mask_blue', params)
+            #    cv2.imwrite('maskout_%d.jpg'%iid, mask_out)
+            stat[1] += 1
+            if params is None:
+                wlabel = [label] + [-1.0]*236
+                stat[0] += 1
+            else:
+                #print(0, params[0].shape, params[0].dtype)
+                #print(1, params[1].shape, params[1].dtype)
+                #print(2, params[2])
+                #print(3, len(params[3]), params[3][0].__class__)
+                #print(4, params[4].shape, params[4].dtype)
+                mask_label = tool.encode_params(params)
+                wlabel = [label, 0.0]+mask_label # 237 including idlabel, total mask params size is 235
+                if iid==0:
+                    print('param size:', len(mask_label), len(wlabel), label)
+            assert len(wlabel)==237
+            wrec.add_image(img, wlabel)
+            #print(len(params))
+
+        wrec.close()
+        print('finished on', self._output, ', failed:', stat[0])
+

insightface/data/__init__.py

@@ -0,0 +1,2 @@
+from .image import get_image
+from .pickle_object import get_object

insightface/data/image.py

@@ -0,0 +1,28 @@
+import cv2
+import os
+import os.path as osp
+from pathlib import Path
+
+class ImageCache:
+    data = {}
+
+def get_image(name, to_rgb=False, use_cache=True):
+    key = (name, to_rgb)
+    if key in ImageCache.data:
+        return ImageCache.data[key]
+    images_dir = osp.join(Path(__file__).parent.absolute(), 'images')
+    ext_names = ['.jpg', '.png', '.jpeg']
+    image_file = None
+    for ext_name in ext_names:
+        _image_file = osp.join(images_dir, "%s%s"%(name, ext_name))
+        if osp.exists(_image_file):
+            image_file = _image_file
+            break
+    assert image_file is not None, '%s not found'%name
+    img = cv2.imread(image_file)
+    if to_rgb:
+        img = img[:,:,::-1]
+    if use_cache:
+        ImageCache.data[key] = img
+    return img
+

insightface/data/objects/meanshape_68.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:39ffecf84ba73f0d0d7e49380833ba88713c9fcdec51df4f7ac45a48b8f4cc51
+size 974

insightface/data/pickle_object.py

@@ -0,0 +1,27 @@
+import sys
+import os
+import os.path as osp
+from pathlib import Path
+import pickle
+
+def get_object(name):
+    if getattr(sys, 'frozen', False):
+        base_dir = sys._MEIPASS
+    else:
+        base_dir = Path(__file__).parent.absolute()
+
+    objects_dir = osp.join(base_dir, 'objects')
+
+    if not name.endswith('.pkl'):
+        name = name + ".pkl"
+
+    filepath = osp.join(objects_dir, name)
+    
+    if not osp.exists(filepath):
+        print(f"[Error] File not found: {filepath}")
+        return None
+
+    with open(filepath, 'rb') as f:
+        obj = pickle.load(f)
+
+    return obj

insightface/data/rec_builder.py

@@ -0,0 +1,71 @@
+import pickle
+import numpy as np
+import os
+import os.path as osp
+import sys
+import mxnet as mx
+
+
+class RecBuilder():
+    def __init__(self, path, image_size=(112, 112)):
+        self.path = path
+        self.image_size = image_size
+        self.widx = 0
+        self.wlabel = 0
+        self.max_label = -1
+        assert not osp.exists(path), '%s exists' % path
+        os.makedirs(path)
+        self.writer = mx.recordio.MXIndexedRecordIO(os.path.join(path, 'train.idx'), 
+                                                    os.path.join(path, 'train.rec'),
+                                                    'w')
+        self.meta = []
+
+    def add(self, imgs):
+        #!!! img should be BGR!!!!
+        #assert label >= 0
+        #assert label > self.last_label
+        assert len(imgs) > 0
+        label = self.wlabel
+        for img in imgs:
+            idx = self.widx
+            image_meta = {'image_index': idx, 'image_classes': [label]}
+            header = mx.recordio.IRHeader(0, label, idx, 0)
+            if isinstance(img, np.ndarray):
+                s = mx.recordio.pack_img(header,img,quality=95,img_fmt='.jpg')
+            else:
+                s = mx.recordio.pack(header, img)
+            self.writer.write_idx(idx, s)
+            self.meta.append(image_meta)
+            self.widx += 1
+        self.max_label = label
+        self.wlabel += 1
+
+
+    def add_image(self, img, label):
+        #!!! img should be BGR!!!!
+        #assert label >= 0
+        #assert label > self.last_label
+        idx = self.widx
+        header = mx.recordio.IRHeader(0, label, idx, 0)
+        if isinstance(label, list):
+            idlabel = label[0]
+        else:
+            idlabel = label
+        image_meta = {'image_index': idx, 'image_classes': [idlabel]}
+        if isinstance(img, np.ndarray):
+            s = mx.recordio.pack_img(header,img,quality=95,img_fmt='.jpg')
+        else:
+            s = mx.recordio.pack(header, img)
+        self.writer.write_idx(idx, s)
+        self.meta.append(image_meta)
+        self.widx += 1
+        self.max_label = max(self.max_label, idlabel)
+
+    def close(self):
+        with open(osp.join(self.path, 'train.meta'), 'wb') as pfile:
+            pickle.dump(self.meta, pfile, protocol=pickle.HIGHEST_PROTOCOL)
+        print('stat:', self.widx, self.wlabel)
+        with open(os.path.join(self.path, 'property'), 'w') as f:
+            f.write("%d,%d,%d\n" % (self.max_label+1, self.image_size[0], self.image_size[1]))
+            f.write("%d\n" % (self.widx))
+

insightface/model_zoo/__init__.py

@@ -0,0 +1,6 @@
+from .model_zoo import get_model
+from .arcface_onnx import ArcFaceONNX
+from .retinaface import RetinaFace
+# from .scrfd import SCRFD
+from .landmark import Landmark
+from .attribute import Attribute