Shellbrady/LivePortrait5

1# coding: utf-82 3"""4face detectoin and alignment using InsightFace5"""6 7import numpy as np8from .rprint import rlog as log9from .dependencies.insightface.app import FaceAnalysis10from .dependencies.insightface.app.common import Face11from .timer import Timer12 13 14def sort_by_direction(faces, direction: str = 'large-small', face_center=None):15    if len(faces) <= 0:16        return faces17 18    if direction == 'left-right':19        return sorted(faces, key=lambda face: face['bbox'][0])20    if direction == 'right-left':21        return sorted(faces, key=lambda face: face['bbox'][0], reverse=True)22    if direction == 'top-bottom':23        return sorted(faces, key=lambda face: face['bbox'][1])24    if direction == 'bottom-top':25        return sorted(faces, key=lambda face: face['bbox'][1], reverse=True)26    if direction == 'small-large':27        return sorted(faces, key=lambda face: (face['bbox'][2] - face['bbox'][0]) * (face['bbox'][3] - face['bbox'][1]))28    if direction == 'large-small':29        return sorted(faces, key=lambda face: (face['bbox'][2] - face['bbox'][0]) * (face['bbox'][3] - face['bbox'][1]), reverse=True)30    if direction == 'distance-from-retarget-face':31        return sorted(faces, key=lambda face: (((face['bbox'][2]+face['bbox'][0])/2-face_center[0])**2+((face['bbox'][3]+face['bbox'][1])/2-face_center[1])**2)**0.5)32    return faces33 34 35class FaceAnalysisDIY(FaceAnalysis):36    def __init__(self, name='buffalo_l', root='~/.insightface', allowed_modules=None, **kwargs):37        super().__init__(name=name, root=root, allowed_modules=allowed_modules, **kwargs)38 39        self.timer = Timer()40 41    def get(self, img_bgr, **kwargs):42        max_num = kwargs.get('max_num', 0)  # the number of the detected faces, 0 means no limit43        flag_do_landmark_2d_106 = kwargs.get('flag_do_landmark_2d_106', True)  # whether to do 106-point detection44        direction = kwargs.get('direction', 'large-small')  # sorting direction45        face_center = None46 47        bboxes, kpss = self.det_model.detect(img_bgr, max_num=max_num, metric='default')48        if bboxes.shape[0] == 0:49            return []50        ret = []51        for i in range(bboxes.shape[0]):52            bbox = bboxes[i, 0:4]53            det_score = bboxes[i, 4]54            kps = None55            if kpss is not None:56                kps = kpss[i]57            face = Face(bbox=bbox, kps=kps, det_score=det_score)58            for taskname, model in self.models.items():59                if taskname == 'detection':60                    continue61 62                if (not flag_do_landmark_2d_106) and taskname == 'landmark_2d_106':63                    continue64 65                # print(f'taskname: {taskname}')66                model.get(img_bgr, face)67            ret.append(face)68 69        ret = sort_by_direction(ret, direction, face_center)70        return ret71 72    def warmup(self):73        self.timer.tic()74 75        img_bgr = np.zeros((512, 512, 3), dtype=np.uint8)76        self.get(img_bgr)77 78        elapse = self.timer.toc()79        log(f'FaceAnalysisDIY warmup time: {elapse:.3f}s')80