Search code examples
pythonmachine-learningcomputer-visionmediapipe

Mediapipe Display Body Landmarks Only


I have installed Mediapipe (0.9.0.1) using Python (3.7.0) on windows 11. I have been able to successfully get Mediapipe to generate landmarks (for face and body); for an image, video, and webcam stream.

I would like to now get Mediapipe to only draw body specific landmarks (i.e. exclude facial landmarks).

I understand that I may use OpenCV (or Czone) to accomplish this goal, however, I am looking to achieve my objective using Mediapipe (i.e. using the draw_landmarks function in the MediaPipe library).

The specific bit of code I am trying (but with errors) is the following:

#Initialize a list to store the detected landmarks.
    landmarks = []

    # Iterate over the Mediapipe detected landmarks.
    for landmark in results.pose_landmarks.landmark:
        
        # Append the Mediapipe landmark into the list.
        landmarks.append((int(landmark.x * width), int(landmark.y * height),
                                (landmark.z * width)))
    
    #create index list for specific landmarks 
    body_landmark_indices = [11,12,13,14,15,16,23,24,25,26,27,28,29,30,31,32]

    landmark_list_body = []

    #Create a list which only has the required landmarks
    for index in body_landmark_indices:
        landmark_list_body.append(landmarks[index - 1])

    mp_drawing.draw_landmarks(
                        image=output_image,
                        landmark_list=landmark_list_body.pose_landmarks,
                        connections=mp_pose.POSE_CONNECTIONS,
                        landmark_drawing_spec=landmark_drawing_spec,
                        connection_drawing_spec=connection_drawing_spec)`

Executing the above I get the error `'list' object has no attribute 'pose_landmarks'

I have replaced landmark_list=landmark_list_body.pose_landmarks, with landmark_list=landmark_list_body but with errors.

I am now very tiered and out of ideas. Is there a capeless hero out there?

Thanks.


Solution

  • You can try the following approach:

    import cv2
    import mediapipe as mp
    import numpy as np
    from mediapipe.python.solutions.pose import PoseLandmark
    from mediapipe.python.solutions.drawing_utils import DrawingSpec
    mp_drawing = mp.solutions.drawing_utils
    mp_drawing_styles = mp.solutions.drawing_styles
    mp_pose = mp.solutions.pose
    
    
    custom_style = mp_drawing_styles.get_default_pose_landmarks_style()
    custom_connections = list(mp_pose.POSE_CONNECTIONS)
    
    # list of landmarks to exclude from the drawing
    excluded_landmarks = [
        PoseLandmark.LEFT_EYE, 
        PoseLandmark.RIGHT_EYE, 
        PoseLandmark.LEFT_EYE_INNER, 
        PoseLandmark.RIGHT_EYE_INNER, 
        PoseLandmark.LEFT_EAR,
        PoseLandmark.RIGHT_EAR,
        PoseLandmark.LEFT_EYE_OUTER,
        PoseLandmark.RIGHT_EYE_OUTER,
        PoseLandmark.NOSE,
        PoseLandmark.MOUTH_LEFT,
        PoseLandmark.MOUTH_RIGHT ]
    
    for landmark in excluded_landmarks:
        # we change the way the excluded landmarks are drawn
        custom_style[landmark] = DrawingSpec(color=(255,255,0), thickness=None) 
        # we remove all connections which contain these landmarks
        custom_connections = [connection_tuple for connection_tuple in custom_connections 
                                if landmark.value not in connection_tuple]
    
    
    IMAGE_FILES = ["test.jpg"]
    BG_COLOR = (192, 192, 192) 
    with mp_pose.Pose(
        static_image_mode=True,
        model_complexity=2,
        enable_segmentation=True,
        min_detection_confidence=0.5) as pose:
        for idx, file in enumerate(IMAGE_FILES):
            image = cv2.imread(file)
            image_height, image_width, _ = image.shape
            results = pose.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
            annotated_image = image.copy()
    
            mp_drawing.draw_landmarks(
                annotated_image,
                results.pose_landmarks,
                connections = custom_connections, #  passing the modified connections list
                landmark_drawing_spec=custom_style) # and drawing style 
    
            cv2.imshow('landmarks', annotated_image)
            cv2.waitKey(0)
    

    It modifies the DrawingSpec and POSE_CONNECTIONS to "hide" a subset of landmarks. However, due to the way the draw_landmarks() function is implemented in Mediapipe, it is also required to add a condition in drawing_utils.py (located in site-packages/mediapipe/python/solutions):

    if drawing_spec.thickness == None: continue
    

    Add it before Line 188 (see MediaPipe GitHub repo, # White circle border). The result should look like this:

    ...
          drawing_spec = landmark_drawing_spec[idx] if isinstance(
              landmark_drawing_spec, Mapping) else landmark_drawing_spec
          if drawing_spec.thickness == None: continue
          # White circle border
          circle_border_radius = max(drawing_spec.circle_radius + 1,
                                     int(drawing_spec.circle_radius * 1.2))
    ...
    

    This change is required in order to completely eliminate the white border that is drawn around landmarks regardless of their drawing specification.

    Hope it helps.