pythonmediapipevideo-tracking

How to change colors of the tracking points and connector lines on the output video using Mediapipe Pose?


I am referring to 33 body points and connector lines between them. I'd like to change the colors of those, especially of the white default color of the connector lines.

Here's my code, I have created a class module for mediapipe which I can import and use in my other programs

import cv2
import mediapipe as mp

class poseDetector():

    def __init__(self, mode=False, complex=1, smooth_landmarks=True, segmentation=True, smooth_segmentation=True,
                 detectionCon=0.5, trackCon=0.5):

        self.mode = mode
        self.complex = complex
        self.smooth_landmarks = smooth_landmarks
        self.segmentation = segmentation
        self.smooth_segmentation = smooth_segmentation
        self.detectionCon = detectionCon
        self.trackCon = trackCon

        self.mpDraw = mp.solutions.drawing_utils
        self.mpDrawStyle = mp.solutions.drawing_styles
        self.mpPose = mp.solutions.pose
        self.pose = self.mpPose.Pose(self.mode, self.complex, self.smooth_landmarks, self.segmentation,
                                     self.smooth_segmentation, self.detectionCon, self.trackCon)

    def findPose(self, img, draw=True):
        imgRGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        self.results = self.pose.process(imgRGB)
        if self.results.pose_landmarks:
            if draw:
                self.mpDraw.draw_landmarks(img, self.results.pose_landmarks,
                                           self.mpPose.POSE_CONNECTIONS)
        return img

def main():
    cap = cv2.VideoCapture("..//assets//videos//v4.mp4")
    detector = poseDetector()
    while True:
        success, img = cap.read()
        img = detector.findPose(img)

        cv2.imshow("Image", img)
        cv2.waitKey(1)


if __name__ == "__main__":
    main()


Solution

  • So as per the documentation, this is the code for draw_landmarks

    mp_drawing.draw_landmarks(
    image: numpy.ndarray,
    landmark_list: mediapipe.framework.formats.landmark_pb2.NormalizedLandmarkList,
    connections: Optional[List[Tuple[int, int]]] = None,
    landmark_drawing_spec: mediapipe.python.solutions.drawing_utils.DrawingSpec = DrawingSpec(color=(0, 0, 255), thickness=2, circle_radius=2),
    connection_drawing_spec: mediapipe.python.solutions.drawing_utils.DrawingSpec = DrawingSpec(color=(0, 255, 0), thickness=2, circle_radius=2),
    )
    

    So in your findPose function you need to update only one line of code

    def findPose(self, img, draw=True):
        imgRGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        self.results = self.pose.process(imgRGB)
        if self.results.pose_landmarks:
            if draw:
                self.mpDraw.draw_landmarks(img, self.results.pose_landmarks,
                                           self.mpPose.POSE_CONNECTIONS,
                self.mpDraw.DrawingSpec(color=(245,117,66), thickness=2, circle_radius=2),
                self.mpDraw.DrawingSpec(color=(245,66,230), thickness=2, circle_radius=2))
        return img
    

    The first self.mpDraw.DrawingSpec argument corresponds to the points of the landmark. The second self.mpDraw.DrawingSpec argument corresponds to the COnnection between those landmarks points. The color is in (B, G, R) format