Spaces:
Runtime error
Runtime error
File size: 2,581 Bytes
c7de15e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
# -*- 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
|