Add face detection

alphapose/face/centerface.py

@@ -0,0 +1,124 @@
+import datetime
+import os
+import cv2
+import numpy as np
+
+current_path = os.path.dirname(__file__)
+class CenterFace(object):
+    def __init__(self, model_path, landmarks=True):
+        self.landmarks = landmarks
+        if self.landmarks:
+            model_path = current_path + '/models/onnx/centerface.onnx'
+            self.net = cv2.dnn.readNetFromONNX(model_path)
+        else:
+            self.net = cv2.dnn.readNetFromONNX('cface.1k.onnx')
+
+    def __call__(self, img, threshold=0.5):
+        blob = cv2.dnn.blobFromImage(img, scalefactor=1.0, size=(self.img_w_new, self.img_h_new), mean=(0, 0, 0), swapRB=True, crop=False)
+        self.net.setInput(blob)
+        begin = datetime.datetime.now()
+        if self.landmarks:
+            heatmap, scale, offset, lms = self.net.forward(["537", "538", "539", '540'])
+        else:
+            heatmap, scale, offset = self.net.forward(["535", "536", "537"])
+
+        end = datetime.datetime.now()
+        print("cpu times = ", end - begin)
+        if self.landmarks:
+            dets, lms = self.decode(heatmap, scale, offset, lms, (self.img_h_new, self.img_w_new), threshold=threshold)
+        else:
+            dets = self.decode(heatmap, scale, offset, None, (self.img_h_new, self.img_w_new), threshold=threshold)
+
+        if len(dets) > 0:
+            dets[:, 0:4:2], dets[:, 1:4:2] = dets[:, 0:4:2] / self.scale_w, dets[:, 1:4:2] / self.scale_h
+            if self.landmarks:
+                lms[:, 0:10:2], lms[:, 1:10:2] = lms[:, 0:10:2] / self.scale_w, lms[:, 1:10:2] / self.scale_h
+        else:
+            dets = np.empty(shape=[0, 5], dtype=np.float32)
+            if self.landmarks:
+                lms = np.empty(shape=[0, 10], dtype=np.float32)
+        if self.landmarks:
+            return dets, lms
+        else:
+            return dets
+
+    def transform(self, h, w):
+        img_h_new, img_w_new = int(np.ceil(h / 32) * 32), int(np.ceil(w / 32) * 32)
+        scale_h, scale_w = img_h_new / h, img_w_new / w
+        self.img_h_new, self.img_w_new, self.scale_h, self.scale_w = img_h_new, img_w_new, scale_h, scale_w
+
+    def decode(self, heatmap, scale, offset, landmark, size, threshold=0.1):
+        heatmap = np.squeeze(heatmap)
+        scale0, scale1 = scale[0, 0, :, :], scale[0, 1, :, :]
+        offset0, offset1 = offset[0, 0, :, :], offset[0, 1, :, :]
+        c0, c1 = np.where(heatmap > threshold)
+        if self.landmarks:
+            boxes, lms = [], []
+        else:
+            boxes = []
+        if len(c0) > 0:
+            for i in range(len(c0)):
+                s0, s1 = np.exp(scale0[c0[i], c1[i]]) * 4, np.exp(scale1[c0[i], c1[i]]) * 4
+                o0, o1 = offset0[c0[i], c1[i]], offset1[c0[i], c1[i]]
+                s = heatmap[c0[i], c1[i]]
+                x1, y1 = max(0, (c1[i] + o1 + 0.5) * 4 - s1 / 2), max(0, (c0[i] + o0 + 0.5) * 4 - s0 / 2)
+                x1, y1 = min(x1, size[1]), min(y1, size[0])
+                boxes.append([x1, y1, min(x1 + s1, size[1]), min(y1 + s0, size[0]), s])
+                if self.landmarks:
+                    lm = []
+                    for j in range(5):
+                        lm.append(landmark[0, j * 2 + 1, c0[i], c1[i]] * s1 + x1)
+                        lm.append(landmark[0, j * 2, c0[i], c1[i]] * s0 + y1)
+                    lms.append(lm)
+            boxes = np.asarray(boxes, dtype=np.float32)
+            keep = self.nms(boxes[:, :4], boxes[:, 4], 0.3)
+            boxes = boxes[keep, :]
+            if self.landmarks:
+                lms = np.asarray(lms, dtype=np.float32)
+                lms = lms[keep, :]
+        if self.landmarks:
+            return boxes, lms
+        else:
+            return boxes
+
+    def nms(self, boxes, scores, nms_thresh):
+        x1 = boxes[:, 0]
+        y1 = boxes[:, 1]
+        x2 = boxes[:, 2]
+        y2 = boxes[:, 3]
+        areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+        order = np.argsort(scores)[::-1]
+        num_detections = boxes.shape[0]
+        suppressed = np.zeros((num_detections,), dtype=np.bool)
+
+        keep = []
+        for _i in range(num_detections):
+            i = order[_i]
+            if suppressed[i]:
+                continue
+            keep.append(i)
+
+            ix1 = x1[i]
+            iy1 = y1[i]
+            ix2 = x2[i]
+            iy2 = y2[i]
+            iarea = areas[i]
+
+            for _j in range(_i + 1, num_detections):
+                j = order[_j]
+                if suppressed[j]:
+                    continue
+
+                xx1 = max(ix1, x1[j])
+                yy1 = max(iy1, y1[j])
+                xx2 = min(ix2, x2[j])
+                yy2 = min(iy2, y2[j])
+                w = max(0, xx2 - xx1 + 1)
+                h = max(0, yy2 - yy1 + 1)
+
+                inter = w * h
+                ovr = inter / (iarea + areas[j] - inter)
+                if ovr >= nms_thresh:
+                    suppressed[j] = True
+
+        return keep

alphapose/face/face.py

@@ -0,0 +1,156 @@
+import os
+import time
+
+import cv2
+import numpy as np
+
+
+from alphapose.face.centerface import CenterFace
+from alphapose.face.prnet import PRN
+from alphapose.face.utils.cv_plot import plot_kpt, plot_pose_box, plot_vertices
+from alphapose.face.utils.render_app import get_visibility, get_uv_mask, get_depth_image
+from alphapose.face.utils.estimate_pose import estimate_pose
+
+current_path = os.path.dirname(__file__)
+
+
+#useless path, enter anything
+face_model_path = '../face/models/onnx/centerface.onnx'
+face_engine = CenterFace(model_path=face_model_path, landmarks=True)
+
+#useless path, enter anything
+face_3d_model_path = '../face/models/prnet.pth'
+
+face_3d_model = PRN(face_3d_model_path, '../face')
+
+colors = [tuple(np.random.choice(np.arange(256).astype(np.int32), size=3)) for i in range(100)]
+
+
+def kp_connections(keypoints):
+    kp_lines = [
+        [keypoints.index('nose'), keypoints.index('left_eye')],
+        [keypoints.index('left_eye'), keypoints.index('left_ear')],
+        [keypoints.index('nose'), keypoints.index('right_eye')],
+        [keypoints.index('right_eye'), keypoints.index('right_ear')],
+        [keypoints.index('right_shoulder'), keypoints.index('right_elbow')],
+        [keypoints.index('right_elbow'), keypoints.index('right_wrist')],
+        [keypoints.index('right_shoulder'), keypoints.index('right_hip')],
+        [keypoints.index('right_hip'), keypoints.index('right_knee')],
+        [keypoints.index('right_knee'), keypoints.index('right_ankle')],
+        [keypoints.index('left_shoulder'), keypoints.index('left_elbow')],
+        [keypoints.index('left_elbow'), keypoints.index('left_wrist')],
+        [keypoints.index('left_shoulder'), keypoints.index('left_hip')],
+        [keypoints.index('left_hip'), keypoints.index('left_knee')],
+        [keypoints.index('left_knee'), keypoints.index('left_ankle')],
+    ]
+    return kp_lines
+
+
+def get_keypoints():
+    """Get the COCO keypoints and their left/right flip coorespondence map."""
+    keypoints = [
+        'nose',  # 1
+        'left_eye',  # 2
+        'right_eye',  # 3
+        'left_ear',  # 4
+        'right_ear',  # 5
+        'left_shoulder',  # 6
+        'right_shoulder',  # 7
+        'left_elbow',  # 8
+        'right_elbow',  # 9
+        'left_wrist',  # 10
+        'right_wrist',  # 11
+        'left_hip',  # 12
+        'right_hip',  # 13
+        'left_knee',  # 14
+        'right_knee',  # 15
+        'left_ankle',  # 16
+        'right_ankle',  # 17
+    ]
+
+    return keypoints
+
+
+_kp_connections = kp_connections(get_keypoints())
+
+
+def face_process(result, rgb_img, orig_img, boxes, scores, ids, preds_img, preds_scores):
+    boxes = boxes.numpy()
+
+    i = 0
+    face_engine.transform(orig_img.shape[0], orig_img.shape[1])
+    face_dets, lms = face_engine(orig_img, threshold=0.35)
+
+    bbox_xywh = []
+    cls_conf = []
+
+    for person in result:
+
+        keypoints = person['keypoints']
+
+        keypoints = keypoints.numpy()
+
+        bbox = boxes[i]
+        color = colors[i]
+
+        body_prob = scores.numpy()
+
+        body_bbox = np.array(bbox[:4], dtype=np.int32)
+        w = body_bbox[2] - body_bbox[0]
+        h = body_bbox[3] - body_bbox[1]
+        bbox_xywh.append([body_bbox[0], body_bbox[1], w, h])
+        cls_conf.append(body_prob)
+
+        center_of_the_face = np.mean(keypoints[:7, :], axis=0)
+
+        image = orig_img
+
+        if len(face_dets) != 0:
+            face_min_dis = np.argmin(
+                np.sum(((face_dets[:, 2:4] + face_dets[:, :2]) / 2. - center_of_the_face) ** 2, axis=1))
+
+            face_bbox = face_dets[face_min_dis][:4]
+            face_prob = face_dets[face_min_dis][4]
+
+
+            face_image = rgb_img[int(face_bbox[1]): int(face_bbox[3]), int(face_bbox[0]): int(face_bbox[2])]
+
+            #cv2.imwrite('/home/jiasong/centerface/prj-python/' + '%d.jpg' % i, face_image)
+
+            [h, w, c] = face_image.shape
+
+            box = np.array(
+                [0, face_image.shape[1] - 1, 0, face_image.shape[0] - 1])  # cropped with bounding box
+
+            pos = face_3d_model.process(face_image, box)
+
+            vertices = face_3d_model.get_vertices(pos)
+            save_vertices = vertices.copy()
+            save_vertices[:, 1] = h - 1 - save_vertices[:, 1]
+
+            kpt = face_3d_model.get_landmarks(pos)
+
+            camera_matrix, pose = estimate_pose(vertices)
+
+            bgr_face_image = cv2.cvtColor(face_image, cv2.COLOR_BGR2RGB)
+            image_pose = plot_pose_box(bgr_face_image, camera_matrix, kpt)
+            sparse_face = plot_kpt(bgr_face_image, kpt)
+
+            dense_face = plot_vertices(bgr_face_image, vertices)
+            image[int(face_bbox[1]): int(face_bbox[3]), int(face_bbox[0]): int(face_bbox[2])] = cv2.resize(
+                sparse_face, (w, h))
+
+
+            for kpt_elem in kpt:
+                kpt_elem[0] +=face_bbox[0]
+                kpt_elem[1] +=face_bbox[1]
+
+            face_keypoints = kpt[:,:2]
+
+            person['FaceKeypoint'] = face_keypoints 
+
+            bgr_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+
+        i += 1
+    return result