conciomith/RetinaFace_FaceDetector_Extractor

1import os2os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'3 4#---------------------------5 6import numpy as np7import tensorflow as tf8import cv29 10import retinaface_model11import preprocess 12import postprocess13 14#---------------------------15 16import tensorflow as tf17tf_version = int(tf.__version__.split(".")[0])18 19if tf_version == 2:20    import logging21    tf.get_logger().setLevel(logging.ERROR)22 23#---------------------------24 25def build_model():26    27    global model #singleton design pattern28    29    if not "model" in globals():30        31        model = tf.function(32            retinaface_model.build_model(),33            input_signature=(tf.TensorSpec(shape=[None, None, None, 3], dtype=np.float32),)34        )35 36    return model37 38def get_image(img_path):39    if type(img_path) == str:  # Load from file path40        if not os.path.isfile(img_path):41            raise ValueError("Input image file path (", img_path, ") does not exist.")42        img = cv2.imread(img_path)43 44    elif isinstance(img_path, np.ndarray):  # Use given NumPy array45        img = img_path.copy()46 47    else:48        raise ValueError("Invalid image input. Only file paths or a NumPy array accepted.")49 50    # Validate image shape51    if len(img.shape) != 3 or np.prod(img.shape) == 0:52        raise ValueError("Input image needs to have 3 channels at must not be empty.")53 54    return img55 56def detect_faces(img_path, threshold=0.9, model = None, allow_upscaling = True):57    """58    TODO: add function doc here59    """60 61    img = get_image(img_path)62 63    #---------------------------64 65    if model is None:66        model = build_model()67 68    #---------------------------69 70    nms_threshold = 0.4; decay4=0.571 72    _feat_stride_fpn = [32, 16, 8]73 74    _anchors_fpn = {75        'stride32': np.array([[-248., -248.,  263.,  263.], [-120., -120.,  135.,  135.]], dtype=np.float32),76        'stride16': np.array([[-56., -56.,  71.,  71.], [-24., -24.,  39.,  39.]], dtype=np.float32),77        'stride8': np.array([[-8., -8., 23., 23.], [ 0.,  0., 15., 15.]], dtype=np.float32)78    }79 80    _num_anchors = {'stride32': 2, 'stride16': 2, 'stride8': 2}81 82    #---------------------------83 84    proposals_list = []85    scores_list = []86    landmarks_list = []87    im_tensor, im_info, im_scale = preprocess.preprocess_image(img, allow_upscaling)88    net_out = model(im_tensor)89    net_out = [elt.numpy() for elt in net_out]90    sym_idx = 091 92    for _idx, s in enumerate(_feat_stride_fpn):93        _key = 'stride%s'%s94        scores = net_out[sym_idx]95        scores = scores[:, :, :, _num_anchors['stride%s'%s]:]96 97        bbox_deltas = net_out[sym_idx + 1]98        height, width = bbox_deltas.shape[1], bbox_deltas.shape[2]99 100        A = _num_anchors['stride%s'%s]101        K = height * width102        anchors_fpn = _anchors_fpn['stride%s'%s]103        anchors = postprocess.anchors_plane(height, width, s, anchors_fpn)104        anchors = anchors.reshape((K * A, 4))105        scores = scores.reshape((-1, 1))106 107        bbox_stds = [1.0, 1.0, 1.0, 1.0]108        bbox_deltas = bbox_deltas109        bbox_pred_len = bbox_deltas.shape[3]//A110        bbox_deltas = bbox_deltas.reshape((-1, bbox_pred_len))111        bbox_deltas[:, 0::4] = bbox_deltas[:,0::4] * bbox_stds[0]112        bbox_deltas[:, 1::4] = bbox_deltas[:,1::4] * bbox_stds[1]113        bbox_deltas[:, 2::4] = bbox_deltas[:,2::4] * bbox_stds[2]114        bbox_deltas[:, 3::4] = bbox_deltas[:,3::4] * bbox_stds[3]115        proposals = postprocess.bbox_pred(anchors, bbox_deltas)116 117        proposals = postprocess.clip_boxes(proposals, im_info[:2])118 119        if s==4 and decay4<1.0:120            scores *= decay4121 122        scores_ravel = scores.ravel()123        order = np.where(scores_ravel>=threshold)[0]124        proposals = proposals[order, :]125        scores = scores[order]126 127        proposals[:, 0:4] /= im_scale128        proposals_list.append(proposals)129        scores_list.append(scores)130 131        landmark_deltas = net_out[sym_idx + 2]132        landmark_pred_len = landmark_deltas.shape[3]//A133        landmark_deltas = landmark_deltas.reshape((-1, 5, landmark_pred_len//5))134        landmarks = postprocess.landmark_pred(anchors, landmark_deltas)135        landmarks = landmarks[order, :]136 137        landmarks[:, :, 0:2] /= im_scale138        landmarks_list.append(landmarks)139        sym_idx += 3140 141    proposals = np.vstack(proposals_list)142    if proposals.shape[0]==0:143        landmarks = np.zeros( (0,5,2) )144        return np.zeros( (0,5) ), landmarks145    scores = np.vstack(scores_list)146    scores_ravel = scores.ravel()147    order = scores_ravel.argsort()[::-1]148 149    proposals = proposals[order, :]150    scores = scores[order]151    landmarks = np.vstack(landmarks_list)152    landmarks = landmarks[order].astype(np.float32, copy=False)153 154    pre_det = np.hstack((proposals[:,0:4], scores)).astype(np.float32, copy=False)155 156    #nms = cpu_nms_wrapper(nms_threshold)157    #keep = nms(pre_det)158    keep = postprocess.cpu_nms(pre_det, nms_threshold)159 160    det = np.hstack( (pre_det, proposals[:,4:]) )161    det = det[keep, :]162    landmarks = landmarks[keep]163 164    resp = {}165    for idx, face in enumerate(det):166 167        label = 'face_'+str(idx+1)168        resp[label] = {}169        resp[label]["score"] = face[4]170 171        resp[label]["facial_area"] = list(face[0:4].astype(int))172 173        resp[label]["landmarks"] = {}174        resp[label]["landmarks"]["right_eye"] = list(landmarks[idx][0])175        resp[label]["landmarks"]["left_eye"] = list(landmarks[idx][1])176        resp[label]["landmarks"]["nose"] = list(landmarks[idx][2])177        resp[label]["landmarks"]["mouth_right"] = list(landmarks[idx][3])178        resp[label]["landmarks"]["mouth_left"] = list(landmarks[idx][4])179 180    return resp181 182def extract_faces(img_path, threshold=0.9, model = None, align = True, allow_upscaling = True):183 184    resp = []185 186    #---------------------------187 188    img = get_image(img_path)189 190    #---------------------------191 192    obj = detect_faces(img_path = img, threshold = threshold, model = model, allow_upscaling = allow_upscaling)193 194    if type(obj) == dict:195        for key in obj:196            identity = obj[key]197 198            facial_area = identity["facial_area"]199            facial_img = img[facial_area[1]: facial_area[3], facial_area[0]: facial_area[2]]200 201            if align == True:202                landmarks = identity["landmarks"]203                left_eye = landmarks["left_eye"]204                right_eye = landmarks["right_eye"]205                nose = landmarks["nose"]206                mouth_right = landmarks["mouth_right"]207                mouth_left = landmarks["mouth_left"]208 209                facial_img = postprocess.alignment_procedure(facial_img, right_eye, left_eye, nose)210 211            resp.append(facial_img[:, :, ::-1])212    #elif type(obj) == tuple:213 214    return resp215