vace/annotators/face.py

# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.

import numpy as np
import torch

from .utils import convert_to_numpy


class FaceAnnotator:
    def __init__(self, cfg, device=None):
        from insightface.app import FaceAnalysis
        self.return_raw = cfg.get('RETURN_RAW', True)
        self.return_mask = cfg.get('RETURN_MASK', False)
        self.return_dict = cfg.get('RETURN_DICT', False)
        self.multi_face = cfg.get('MULTI_FACE', True)
        pretrained_model = cfg['PRETRAINED_MODEL']
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if device is None else device
        self.device_id = self.device.index if self.device.type == 'cuda' else None
        ctx_id = self.device_id if self.device_id is not None else 0
        self.model = FaceAnalysis(name=cfg.MODEL_NAME, root=pretrained_model, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
        self.model.prepare(ctx_id=ctx_id, det_size=(640, 640))

    def forward(self, image=None, return_mask=None, return_dict=None):
        return_mask = return_mask if return_mask is not None else self.return_mask
        return_dict = return_dict if return_dict is not None else self.return_dict
        image = convert_to_numpy(image)
        # [dict_keys(['bbox', 'kps', 'det_score', 'landmark_3d_68', 'pose', 'landmark_2d_106', 'gender', 'age', 'embedding'])]
        faces = self.model.get(image)
        if self.return_raw:
            return faces
        else:
            crop_face_list, mask_list = [], []
            if len(faces) > 0:
                if not self.multi_face:
                    faces = faces[:1]
                for face in faces:
                    x_min, y_min, x_max, y_max = face['bbox'].tolist()
                    crop_face = image[int(y_min): int(y_max) + 1, int(x_min): int(x_max) + 1]
                    crop_face_list.append(crop_face)
                    mask = np.zeros_like(image[:, :, 0])
                    mask[int(y_min): int(y_max) + 1, int(x_min): int(x_max) + 1] = 255
                    mask_list.append(mask)
                if not self.multi_face:
                    crop_face_list = crop_face_list[0]
                    mask_list = mask_list[0]
                if return_mask:
                    if return_dict:
                        return {'image': crop_face_list, 'mask': mask_list}
                    else:
                        return crop_face_list, mask_list
                else:
                    return crop_face_list
            else:
                return None