models/pose_estimation_mediapipe/mp_pose.py

import numpy as np
import cv2 as cv

class MPPose:
    def __init__(self, modelPath, confThreshold=0.5, backendId=0, targetId=0):
        self.model_path = modelPath
        self.conf_threshold = confThreshold
        self.backend_id = backendId
        self.target_id = targetId

        self.input_size = np.array([256, 256])  # wh
        # RoI will be larger so the performance will be better, but preprocess will be slower. Default to 1.
        self.PERSON_BOX_PRE_ENLARGE_FACTOR = 1
        self.PERSON_BOX_ENLARGE_FACTOR = 1.25

        self.model = cv.dnn.readNet(self.model_path)
        self.model.setPreferableBackend(self.backend_id)
        self.model.setPreferableTarget(self.target_id)

    @property
    def name(self):
        return self.__class__.__name__

    def setBackendAndTarget(self, backendId, targetId):
        self._backendId = backendId
        self._targetId = targetId
        self.model.setPreferableBackend(self.backend_id)
        self.model.setPreferableTarget(self.target_id)

    def _preprocess(self, image, person):
        '''
        Rotate input for inference.
        Parameters:
          image - input image of BGR channel order
          face_bbox - human face bounding box found in image of format [[x1, y1], [x2, y2]] (top-left and bottom-right points)
          person_landmarks - 4 landmarks (2 full body points, 2 upper body points) of shape [4, 2]
        Returns:
          rotated_person - rotated person image for inference
          rotate_person_bbox - person box of interest range
          angle - rotate angle for person
          rotation_matrix - matrix for rotation and de-rotation
          pad_bias - pad pixels of interest range
        '''
        # crop and pad image to interest range
        pad_bias = np.array([0, 0], dtype=np.int32)  # left, top
        person_keypoints = person[4: 12].reshape(-1, 2)
        mid_hip_point = person_keypoints[0]
        full_body_point = person_keypoints[1]
        # get RoI
        full_dist = np.linalg.norm(mid_hip_point - full_body_point)
        full_bbox = np.array([mid_hip_point - full_dist, mid_hip_point + full_dist], np.int32)
        # enlarge to make sure full body can be cover
        center_bbox = np.sum(full_bbox, axis=0) / 2
        wh_bbox = full_bbox[1] - full_bbox[0]
        new_half_size = wh_bbox * self.PERSON_BOX_PRE_ENLARGE_FACTOR / 2
        full_bbox = np.array([
            center_bbox - new_half_size,
            center_bbox + new_half_size], np.int32)

        person_bbox = full_bbox.copy()
        # refine person bbox
        person_bbox[:, 0] = np.clip(person_bbox[:, 0], 0, image.shape[1])
        person_bbox[:, 1] = np.clip(person_bbox[:, 1], 0, image.shape[0])
        # crop to the size of interest
        image = image[person_bbox[0][1]:person_bbox[1][1], person_bbox[0][0]:person_bbox[1][0], :]
        # pad to square
        left, top = person_bbox[0] - full_bbox[0]
        right, bottom = full_bbox[1] - person_bbox[1]
        image = cv.copyMakeBorder(image, top, bottom, left, right, cv.BORDER_CONSTANT, None, (0, 0, 0))
        pad_bias += person_bbox[0] - [left, top]
        # compute rotation
        mid_hip_point -= pad_bias
        full_body_point -= pad_bias
        radians = np.pi / 2 - np.arctan2(-(full_body_point[1] - mid_hip_point[1]), full_body_point[0] - mid_hip_point[0])
        radians = radians - 2 * np.pi * np.floor((radians + np.pi) / (2 * np.pi))
        angle = np.rad2deg(radians)
        #  get rotation matrix
        rotation_matrix = cv.getRotationMatrix2D(mid_hip_point, angle, 1.0)
        #  get rotated image
        rotated_image = cv.warpAffine(image, rotation_matrix, (image.shape[1], image.shape[0]))
        #  get landmark bounding box
        blob = cv.resize(rotated_image, dsize=self.input_size, interpolation=cv.INTER_AREA).astype(np.float32)
        rotated_person_bbox = np.array([[0, 0], [image.shape[1], image.shape[0]]], dtype=np.int32)
        blob = cv.cvtColor(blob, cv.COLOR_BGR2RGB)
        blob = blob / 255. # [0, 1]
        return blob[np.newaxis, :, :, :], rotated_person_bbox, angle, rotation_matrix, pad_bias

    def infer(self, image, person):
        h, w, _ = image.shape
        # Preprocess
        input_blob, rotated_person_bbox, angle, rotation_matrix, pad_bias = self._preprocess(image, person)

        # Forward
        self.model.setInput(input_blob)
        output_blob = self.model.forward(self.model.getUnconnectedOutLayersNames())

        # Postprocess
        results = self._postprocess(output_blob, rotated_person_bbox, angle, rotation_matrix, pad_bias, np.array([w, h]))
        return results # [bbox_coords, landmarks_coords, conf]

    def _postprocess(self, blob, rotated_person_bbox, angle, rotation_matrix, pad_bias, img_size):
        landmarks, conf, mask, heatmap, landmarks_word = blob

        conf = conf[0][0]
        if conf < self.conf_threshold:
            return None

        landmarks = landmarks[0].reshape(-1, 5)  # shape: (1, 195) -> (39, 5)
        landmarks_word = landmarks_word[0].reshape(-1, 3) # shape: (1, 117) -> (39, 3)

        # recover sigmoid score
        landmarks[:, 3:] = 1 / (1 + np.exp(-landmarks[:, 3:]))
        # TODO: refine landmarks with heatmap. reference: https://github.com/tensorflow/tfjs-models/blob/master/pose-detection/src/blazepose_tfjs/detector.ts#L577-L582
        heatmap = heatmap[0]

        # transform coords back to the input coords
        wh_rotated_person_bbox = rotated_person_bbox[1] - rotated_person_bbox[0]
        scale_factor = wh_rotated_person_bbox / self.input_size
        landmarks[:, :2] = (landmarks[:, :2] - self.input_size / 2) * scale_factor
        landmarks[:, 2] = landmarks[:, 2] * max(scale_factor) # depth scaling
        coords_rotation_matrix = cv.getRotationMatrix2D((0, 0), angle, 1.0)
        rotated_landmarks = np.dot(landmarks[:, :2], coords_rotation_matrix[:, :2])
        rotated_landmarks = np.c_[rotated_landmarks, landmarks[:, 2:]]
        rotated_landmarks_world = np.dot(landmarks_word[:, :2], coords_rotation_matrix[:, :2])
        rotated_landmarks_world = np.c_[rotated_landmarks_world, landmarks_word[:, 2]]
        # invert rotation
        rotation_component = np.array([
            [rotation_matrix[0][0], rotation_matrix[1][0]],
            [rotation_matrix[0][1], rotation_matrix[1][1]]])
        translation_component = np.array([
            rotation_matrix[0][2], rotation_matrix[1][2]])
        inverted_translation = np.array([
            -np.dot(rotation_component[0], translation_component),
            -np.dot(rotation_component[1], translation_component)])
        inverse_rotation_matrix = np.c_[rotation_component, inverted_translation]
        #  get box center
        center = np.append(np.sum(rotated_person_bbox, axis=0) / 2, 1)
        original_center = np.array([
            np.dot(center, inverse_rotation_matrix[0]),
            np.dot(center, inverse_rotation_matrix[1])])
        landmarks[:, :2] = rotated_landmarks[:, :2] + original_center + pad_bias

        # get bounding box from rotated_landmarks
        bbox = np.array([
            np.amin(landmarks[:, :2], axis=0),
            np.amax(landmarks[:, :2], axis=0)])  # [top-left, bottom-right]
        center_bbox = np.sum(bbox, axis=0) / 2
        wh_bbox = bbox[1] - bbox[0]
        new_half_size = wh_bbox * self.PERSON_BOX_ENLARGE_FACTOR / 2
        bbox = np.array([
            center_bbox - new_half_size,
            center_bbox + new_half_size])

        # invert rotation for mask
        mask = mask[0].reshape(256, 256) # shape: (1, 256, 256, 1) -> (256, 256)
        invert_rotation_matrix = cv.getRotationMatrix2D((mask.shape[1]/2, mask.shape[0]/2), -angle, 1.0)
        invert_rotation_mask = cv.warpAffine(mask, invert_rotation_matrix, (mask.shape[1], mask.shape[0]))
        # enlarge mask
        invert_rotation_mask = cv.resize(invert_rotation_mask, wh_rotated_person_bbox)
        # crop and pad mask
        min_w, min_h = -np.minimum(pad_bias, 0)
        left, top = np.maximum(pad_bias, 0)
        pad_over = img_size - [invert_rotation_mask.shape[1], invert_rotation_mask.shape[0]] - pad_bias
        max_w, max_h = np.minimum(pad_over, 0) + [invert_rotation_mask.shape[1], invert_rotation_mask.shape[0]]
        right, bottom = np.maximum(pad_over, 0)
        invert_rotation_mask = invert_rotation_mask[min_h:max_h, min_w:max_w]
        invert_rotation_mask = cv.copyMakeBorder(invert_rotation_mask, top, bottom, left, right, cv.BORDER_CONSTANT, None, 0)
        # binarize mask
        invert_rotation_mask = np.where(invert_rotation_mask > 0, 255, 0).astype(np.uint8)

        # 2*2 person bbox: [[x1, y1], [x2, y2]]
        # 39*5 screen landmarks: 33 keypoints and 6 auxiliary points with [x, y, z, visibility, presence], z value is relative to HIP
        # Visibility is probability that a keypoint is located within the frame and not occluded by another bigger body part or another object
        # Presence is probability that a keypoint is located within the frame
        # 39*3 world landmarks: 33 keypoints and 6 auxiliary points with [x, y, z] 3D metric x, y, z coordinate
        # img_height*img_width mask: gray mask, where 255 indicates the full body of a person and 0 means background
        # 64*64*39 heatmap: currently only used for refining landmarks, requires sigmod processing before use
        # conf: confidence of prediction
        return [bbox, landmarks, rotated_landmarks_world, invert_rotation_mask, heatmap, conf]