64 lines
2.3 KiB
Python
64 lines
2.3 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import os
|
|
from abc import ABCMeta
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from scepter.modules.annotator.base_annotator import BaseAnnotator
|
|
from scepter.modules.annotator.registry import ANNOTATORS
|
|
from scepter.modules.utils.distribute import we
|
|
from scepter.modules.utils.file_system import FS
|
|
|
|
|
|
@ANNOTATORS.register_class()
|
|
class FaceAnnotator(BaseAnnotator, metaclass=ABCMeta):
|
|
para_dict = {}
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super().__init__(cfg, logger=logger)
|
|
from insightface.app import FaceAnalysis
|
|
local_path = FS.map_to_local(cfg.PRETRAINED_MODEL)[0]
|
|
local_model_path = os.path.join(local_path, 'models', cfg.MODEL_NAME)
|
|
FS.get_dir_to_local_dir(cfg.PRETRAINED_MODEL, local_model_path)
|
|
self.model = FaceAnalysis(name=cfg.MODEL_NAME, root=local_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
|
|
self.model.prepare(ctx_id=we.device_id, det_size=(640, 640))
|
|
|
|
def forward(self, image=None):
|
|
|
|
if isinstance(image, Image.Image):
|
|
image = np.array(image)
|
|
elif isinstance(image, torch.Tensor):
|
|
image = image.detach().cpu().numpy()
|
|
elif isinstance(image, np.ndarray):
|
|
image = image.copy()
|
|
else:
|
|
raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
|
|
|
|
# [dict_keys(['bbox', 'kps', 'det_score', 'landmark_3d_68', 'pose', 'landmark_2d_106', 'gender', 'age', 'embedding'])]
|
|
faces = self.model.get(image)
|
|
return faces
|
|
|
|
|
|
@ANNOTATORS.register_class()
|
|
class FaceMaskAnnotator(FaceAnnotator):
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super().__init__(cfg, logger=logger)
|
|
self.multi_face = cfg.get('MULTI_FACE', True)
|
|
|
|
def forward(self, image=None):
|
|
faces = super().forward(image)
|
|
if len(faces) > 0:
|
|
if not self.multi_face:
|
|
faces = faces[:1]
|
|
mask = np.zeros_like(image[:, :, 0])
|
|
for face in faces:
|
|
x_min, y_min, x_max, y_max = face['bbox'].tolist()
|
|
mask[int(y_min): int(y_max) + 1, int(x_min): int(x_max) + 1] = 255
|
|
return mask
|
|
else:
|
|
return np.zeros_like(image[:, :, 0])
|