| | |
| | |
| |
|
| |
|
| |
|
| | import numpy as np |
| | import cv2 |
| | from cv2box.utils.math import Normalize |
| | from cv2box import CVImage |
| |
|
| | from .scrfd_insightface import SCRFD |
| | from face_detect.face_align_utils import norm_crop, apply_roi_func |
| |
|
| | |
| | SCRFD_MODEL_PATH = '../pretrain_models/' |
| |
|
| | class FaceDetect5Landmarks: |
| | def __init__(self, mode='scrfd_500m', tracking=False): |
| | self.mode = mode |
| | self.tracking = tracking |
| | self.dis_list = [] |
| | self.last_bboxes_ = [] |
| | assert self.mode in ['scrfd', 'scrfd_500m', 'mtcnn'] |
| | self.bboxes = self.kpss = self.image = None |
| | if 'scrfd' in self.mode: |
| | if self.mode == 'scrfd_500m': |
| | scrfd_model_path = SCRFD_MODEL_PATH + 'scrfd_500m_bnkps_shape640x640.onnx' |
| | else: |
| | scrfd_model_path = SCRFD_MODEL_PATH + 'scrfd_10g_bnkps.onnx' |
| | self.det_model_scrfd = SCRFD(scrfd_model_path) |
| | self.det_model_scrfd.prepare(ctx_id=0, input_size=(640, 640)) |
| |
|
| |
|
| | def get_bboxes(self, image, nms_thresh=0.5, max_num=0, min_bbox_size=None): |
| | """ |
| | Args: |
| | image: RGB image path or Numpy array load by cv2 |
| | nms_thresh: |
| | max_num: |
| | min_bbox_size: |
| | Returns: |
| | """ |
| | self.image = CVImage(image).rgb() |
| |
|
| | if self.tracking: |
| | if len(self.last_bboxes_) == 0: |
| | self.bboxes, self.kpss = self.det_model_scrfd.detect(image, thresh=nms_thresh, max_num=1, |
| | metric='default') |
| | self.last_bboxes_ = self.bboxes |
| | |
| | else: |
| | self.bboxes, self.kpss = self.det_model_scrfd.detect(image, thresh=nms_thresh, max_num=0, |
| | metric='default') |
| | self.bboxes, self.kpss = self.tracking_filter() |
| | else: |
| | if 'scrfd' in self.mode: |
| | self.bboxes, self.kpss = self.det_model_scrfd.detect(self.image, thresh=nms_thresh, |
| | max_num=max_num, |
| | metric='default') |
| |
|
| | return self.bboxes, self.kpss |
| |
|
| | def tracking_filter(self): |
| | for i in range(len(self.bboxes)): |
| | self.dis_list.append(np.linalg.norm(Normalize(self.bboxes[i]).np_norm() - Normalize(self.last_bboxes_[0]).np_norm())) |
| | if not self.dis_list: |
| | return [], [] |
| | best_index = np.argmin(np.array(self.dis_list)) |
| | self.dis_list = [] |
| | self.last_bboxes_ = [self.bboxes[best_index]] |
| | return self.last_bboxes_, [self.kpss[best_index]] |
| |
|
| | def bboxes_filter(self, min_bbox_size): |
| | min_area = np.power(min_bbox_size, 2) |
| | area_list = (self.bboxes[:, 2] - self.bboxes[:, 0]) * (self.bboxes[:, 3] - self.bboxes[:, 1]) |
| | min_index = np.where(area_list < min_area) |
| | self.bboxes = np.delete(self.bboxes, min_index, axis=0) |
| | self.kpss = np.delete(self.kpss, min_index, axis=0) |
| |
|
| | def get_single_face(self, crop_size, mode='mtcnn_512', apply_roi=False): |
| | """ |
| | Args: |
| | crop_size: |
| | mode: default mtcnn_512 arcface_512 arcface default_95 |
| | Returns: cv2 image |
| | """ |
| | assert mode in ['default', 'mtcnn_512', 'mtcnn_256', 'arcface_512', 'arcface', 'default_95'] |
| | if self.bboxes.shape[0] == 0: |
| | return None, None |
| | det_score = self.bboxes[..., 4] |
| | if self.tracking: |
| | best_index = np.argmax(np.array(self.dis_list)) |
| | kpss = None |
| | if self.kpss is not None: |
| | kpss = self.kpss[best_index] |
| | else: |
| | best_index = np.argmax(det_score) |
| | kpss = None |
| | if self.kpss is not None: |
| | kpss = self.kpss[best_index] |
| | if apply_roi: |
| | roi, roi_box, roi_kpss = apply_roi_func(self.image, self.bboxes[best_index], kpss) |
| | align_img, mat_rev = norm_crop(roi, roi_kpss, crop_size, mode=mode) |
| | align_img = cv2.cvtColor(align_img, cv2.COLOR_RGB2BGR) |
| | return align_img, mat_rev, roi_box |
| | else: |
| | align_img, M = norm_crop(self.image, kpss, crop_size, mode=mode) |
| | align_img = cv2.cvtColor(align_img, cv2.COLOR_RGB2BGR) |
| | return align_img, M |
| |
|
| | def get_multi_face(self, crop_size, mode='mtcnn_512'): |
| | """ |
| | Args: |
| | crop_size: |
| | mode: default mtcnn_512 arcface_512 arcface |
| | Returns: |
| | """ |
| | if self.bboxes.shape[0] == 0: |
| | return None |
| | align_img_list = [] |
| | M_list = [] |
| | for i in range(self.bboxes.shape[0]): |
| | kps = None |
| | if self.kpss is not None: |
| | kps = self.kpss[i] |
| | align_img, M = norm_crop(self.image, kps, crop_size, mode=mode) |
| | align_img_list.append(align_img) |
| | M_list.append(M) |
| | return align_img_list, M_list |
| |
|
| | def draw_face(self): |
| | for i_ in range(self.bboxes.shape[0]): |
| | bbox = self.bboxes[i_] |
| | x1, y1, x2, y2, score = bbox.astype(int) |
| | cv2.rectangle(self.image, (x1, y1), (x2, y2), (255, 0, 0), 2) |
| | if self.kpss is not None: |
| | kps = self.kpss[i_] |
| | for kp in kps: |
| | kp = kp.astype(int) |
| | cv2.circle(self.image, tuple(kp), 1, (0, 0, 255), 2) |
| | CVImage(self.image, image_format='cv2').show() |
| |
|
