Shellbrady/LivePortrait5

1# -*- coding: utf-8 -*-2# @Organization  : insightface.ai3# @Author        : Jia Guo4# @Time          : 2021-05-045# @Function      :6 7 8from __future__ import division9 10import glob11import os.path as osp12 13import numpy as np14import onnxruntime15from numpy.linalg import norm16 17from ..model_zoo import model_zoo18from ..utils import ensure_available19from .common import Face20 21 22DEFAULT_MP_NAME = 'buffalo_l'23__all__ = ['FaceAnalysis']24 25class FaceAnalysis:26    def __init__(self, name=DEFAULT_MP_NAME, root='~/.insightface', allowed_modules=None, **kwargs):27        onnxruntime.set_default_logger_severity(3)28        self.models = {}29        self.model_dir = ensure_available('models', name, root=root)30        onnx_files = glob.glob(osp.join(self.model_dir, '*.onnx'))31        onnx_files = sorted(onnx_files)32        for onnx_file in onnx_files:33            model = model_zoo.get_model(onnx_file, **kwargs)34            if model is None:35                print('model not recognized:', onnx_file)36            elif allowed_modules is not None and model.taskname not in allowed_modules:37                print('model ignore:', onnx_file, model.taskname)38                del model39            elif model.taskname not in self.models and (allowed_modules is None or model.taskname in allowed_modules):40                # print('find model:', onnx_file, model.taskname, model.input_shape, model.input_mean, model.input_std)41                self.models[model.taskname] = model42            else:43                print('duplicated model task type, ignore:', onnx_file, model.taskname)44                del model45        assert 'detection' in self.models46        self.det_model = self.models['detection']47 48 49    def prepare(self, ctx_id, det_thresh=0.5, det_size=(640, 640)):50        self.det_thresh = det_thresh51        assert det_size is not None52        # print('set det-size:', det_size)53        self.det_size = det_size54        for taskname, model in self.models.items():55            if taskname=='detection':56                model.prepare(ctx_id, input_size=det_size, det_thresh=det_thresh)57            else:58                model.prepare(ctx_id)59 60    def get(self, img, max_num=0):61        bboxes, kpss = self.det_model.detect(img,62                                             max_num=max_num,63                                             metric='default')64        if bboxes.shape[0] == 0:65            return []66        ret = []67        for i in range(bboxes.shape[0]):68            bbox = bboxes[i, 0:4]69            det_score = bboxes[i, 4]70            kps = None71            if kpss is not None:72                kps = kpss[i]73            face = Face(bbox=bbox, kps=kps, det_score=det_score)74            for taskname, model in self.models.items():75                if taskname=='detection':76                    continue77                model.get(img, face)78            ret.append(face)79        return ret80 81    def draw_on(self, img, faces):82        import cv283        dimg = img.copy()84        for i in range(len(faces)):85            face = faces[i]86            box = face.bbox.astype(np.int)87            color = (0, 0, 255)88            cv2.rectangle(dimg, (box[0], box[1]), (box[2], box[3]), color, 2)89            if face.kps is not None:90                kps = face.kps.astype(np.int)91                #print(landmark.shape)92                for l in range(kps.shape[0]):93                    color = (0, 0, 255)94                    if l == 0 or l == 3:95                        color = (0, 255, 0)96                    cv2.circle(dimg, (kps[l][0], kps[l][1]), 1, color,97                               2)98            if face.gender is not None and face.age is not None:99                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)100 101            #for key, value in face.items():102            #    if key.startswith('landmark_3d'):103            #        print(key, value.shape)104            #        print(value[0:10,:])105            #        lmk = np.round(value).astype(np.int)106            #        for l in range(lmk.shape[0]):107            #            color = (255, 0, 0)108            #            cv2.circle(dimg, (lmk[l][0], lmk[l][1]), 1, color,109            #                       2)110        return dimg111