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,  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,86 @@
+import os
+import time
+
+import cv2
+import numpy as np
+
+
+from alphapose.face.centerface import CenterFace
+from alphapose.face.prnet import PRN
+
+#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, device, '../face')
+
+
+def face_process(face_3d_model, result, orig_img):
+    face_engine = CenterFace( landmarks=True)
+    rgb_img = orig_img[:, :, ::-1]
+    [H, W, _] = rgb_img.shape
+    face_engine.transform(orig_img.shape[0], orig_img.shape[1])
+    face_dets, lms = face_engine(orig_img, threshold=0.35)
+
+    result_new = []
+
+    for person in result:
+
+        keypoints = person['keypoints']
+        kp_score = person['kp_score']
+
+        keypoints = keypoints.numpy()
+        kp_score = kp_score.numpy()
+
+
+        center_of_the_face = np.mean(keypoints[:5, :], axis=0)
+        face_conf = np.mean(kp_score[:5, :], axis=0)
+
+        face_keypoints = -1*np.ones((68,3))
+        if face_conf > 0.5 and 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]
+            if center_of_the_face[0] < face_bbox[0] or center_of_the_face[1] < face_bbox[1] or center_of_the_face[0] > face_bbox[2] or center_of_the_face[1] > face_bbox[3]:
+                continue 
+            if face_prob < 0.5:
+                continue
+
+            ## below is by intuitive box
+            # wid = max(keypoints[:5, 0]) - min(keypoints[:5, 0])
+            # hgt = max(keypoints[:5, 1]) - min(keypoints[:5, 1])
+            # face_bbox = [max(0,min(keypoints[:5, 0])-0.1*wid), max(0,min(keypoints[:5, 1])-2*hgt), min(max(keypoints[:5, 0])+0.1*wid,W),  min(H,max(keypoints[:5, 1])+2.5*hgt)]
+            # print(face_bbox)
+
+
+
+            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)
+
+            for kpt_elem in kpt:
+                kpt_elem[0] +=face_bbox[0]
+                kpt_elem[1] +=face_bbox[1]
+
+            face_keypoints = kpt[:,:3]
+        # print('face', face_keypoints)
+        person['FaceKeypoint'] = face_keypoints 
+        result_new.append(person)            
+    return result_new

alphapose/face/prnet.py

@@ -0,0 +1,174 @@
+import os
+from time import time
+
+import numpy as np
+import torch
+from skimage.io import imread, imsave
+from skimage.transform import estimate_transform, warp
+
+from .resfcn256 import PRNet
+
+current_path = os.path.dirname(__file__)
+class PRN:
+    ''' Joint 3D Face Reconstruction and Dense Alignment with Position Map Regression Network
+    Args:
+        is_dlib(bool, optional): If true, dlib is used for detecting faces.
+        prefix(str, optional): If run at another folder, the absolute path is needed to load the data.
+    '''
+    def __init__(self, device):
+
+        # resolution of input and output image size.
+        self.resolution_inp = 256
+        self.resolution_op = 256
+        self.device = device
+
+
+        # 1) load model.
+        self.pos_predictor = PRNet(3, 3)
+        model_path = current_path + '/models/prnet.pth'
+
+        state = torch.load(model_path)
+        self.pos_predictor.load_state_dict(state)
+        self.pos_predictor.eval().to(self.device)  # inference stage only.
+
+
+        # uv file
+        self.uv_kpt_ind = np.loadtxt(os.path.join(current_path, 'utils/uv_data/uv_kpt_ind.txt')).astype(
+            np.int32)  # 2 x 68 get kpt
+        self.face_ind = np.loadtxt(os.path.join(current_path, 'utils/uv_data/face_ind.txt')).astype(
+            np.int32)  # get valid vertices in the pos map
+        self.triangles = np.loadtxt(os.path.join(current_path, 'utils/uv_data/triangles.txt')).astype(
+            np.int32)  # ntri x 3
+
+        self.uv_coords = self.generate_uv_coords()
+
+    def generate_uv_coords(self):
+        resolution = self.resolution_op
+        uv_coords = np.meshgrid(range(resolution), range(resolution))
+        uv_coords = np.transpose(np.array(uv_coords), [1, 2, 0])
+        uv_coords = np.reshape(uv_coords, [resolution ** 2, -1]);
+        uv_coords = uv_coords[self.face_ind, :]
+        uv_coords = np.hstack((uv_coords[:, :2], np.zeros([uv_coords.shape[0], 1])))
+        return uv_coords
+
+    def net_forward(self, image):
+        ''' The core of out method: regress the position map of a given image.
+        Args:
+            image: (256,256,3) array. value range: 0~1
+        Returns:
+            pos: the 3D position map. (256, 256, 3) array.
+        '''
+        out = self.pos_predictor(image).cpu().detach().numpy()
+        out = np.transpose(out, (0, 2, 3, 1)).squeeze()
+
+        return out
+
+    def process(self, input, image_info = None):
+        ''' process image with crop operation.
+        Args:
+            input: (h,w,3) array or str(image path). image value range:1~255. 
+            image_info(optional): the bounding box information of faces. if None, will use dlib to detect face. 
+
+        Returns:
+            pos: the 3D position map. (256, 256, 3).
+        '''
+        if isinstance(input, str):
+            try:
+                image = imread(input)
+            except IOError:
+                print("error opening file: ", input)
+                return None
+        else:
+            image = input
+
+        if image.ndim < 3:
+            image = np.tile(image[:,:,np.newaxis], [1,1,3])
+
+
+        if np.max(image_info.shape) > 4: # key points to get bounding box
+            kpt = image_info
+            if kpt.shape[0] > 3:
+                kpt = kpt.T
+            left = np.min(kpt[0, :]); right = np.max(kpt[0, :]); 
+            top = np.min(kpt[1,:]); bottom = np.max(kpt[1,:])
+        else:  # bounding box
+            bbox = image_info
+            left = bbox[0]; right = bbox[1]; top = bbox[2]; bottom = bbox[3]
+        old_size = (right - left + bottom - top)/2
+        center = np.array([right - (right - left) / 2.0, bottom - (bottom - top) / 2.0])
+        size = int(old_size*1.6)
+
+        # crop image
+        src_pts = np.array([[center[0]-size/2, center[1]-size/2], [center[0] - size/2, center[1]+size/2], [center[0]+size/2, center[1]-size/2]])
+        DST_PTS = np.array([[0,0], [0,self.resolution_inp - 1], [self.resolution_inp - 1, 0]])
+        tform = estimate_transform('similarity', src_pts, DST_PTS)
+
+        image = image/255.
+        cropped_image = warp(image, tform.inverse, output_shape=(self.resolution_inp, self.resolution_inp))
+        # run our net
+        #st = time()
+        cropped_image = torch.from_numpy(cropped_image[np.newaxis, ...].transpose(0,3,1,2).astype(np.float32)).to(self.device)
+
+        cropped_pos = self.net_forward(cropped_image)*self.resolution_inp*1.1
+        #print 'net time:', time() - st
+
+        # restore 
+        cropped_vertices = np.reshape(cropped_pos, [-1, 3]).T
+        z = cropped_vertices[2,:].copy()/tform.params[0,0]
+        cropped_vertices[2,:] = 1
+        vertices = np.dot(np.linalg.inv(tform.params), cropped_vertices)
+        vertices = np.vstack((vertices[:2,:], z))
+        pos = np.reshape(vertices.T, [self.resolution_op, self.resolution_op, 3])
+
+        return pos
+
+    def get_landmarks(self, pos):
+        '''
+        Args:
+            pos: the 3D position map. shape = (256, 256, 3).
+        Returns:
+            kpt: 68 3D landmarks. shape = (68, 3).
+        '''
+        kpt = pos[self.uv_kpt_ind[1,:], self.uv_kpt_ind[0,:], :]
+        return kpt
+
+
+    def get_vertices(self, pos):
+        '''
+        Args:
+            pos: the 3D position map. shape = (256, 256, 3).
+        Returns:
+            vertices: the vertices(point cloud). shape = (num of points, 3). n is about 40K here.
+        '''
+        all_vertices = np.reshape(pos, [self.resolution_op**2, -1])
+        vertices = all_vertices[self.face_ind, :]
+
+        return vertices
+
+    def get_colors_from_texture(self, texture):
+        '''
+        Args:
+            texture: the texture map. shape = (256, 256, 3).
+        Returns:
+            colors: the corresponding colors of vertices. shape = (num of points, 3). n is 45128 here.
+        '''
+        all_colors = np.reshape(texture, [self.resolution_op**2, -1])
+        colors = all_colors[self.face_ind, :]
+
+        return colors
+
+
+    def get_colors(self, image, vertices):
+        '''
+        Args:
+            pos: the 3D position map. shape = (256, 256, 3).
+        Returns:
+            colors: the corresponding colors of vertices. shape = (num of points, 3). n is 45128 here.
+        '''
+        [h, w, _] = image.shape
+        vertices[:,0] = np.minimum(np.maximum(vertices[:,0], 0), w - 1)  # x
+        vertices[:,1] = np.minimum(np.maximum(vertices[:,1], 0), h - 1)  # y
+        ind = np.round(vertices).astype(np.int32)
+        colors = image[ind[:,1], ind[:,0], :] # n x 3
+
+        return colors