mediapipe_test.py

import cv2
import mediapipe as mp
import SquatPosture as sp
import numpy as np
from utils import *

mp_drawing = mp.solutions.drawing_utils
mp_pose = mp.solutions.pose

# For video input:
# cap = cv2.VideoCapture("data/processed/024_squat.mp4")
cap = cv2.VideoCapture(0)

with mp_pose.Pose(
        min_detection_confidence=0.5,
        min_tracking_confidence=0.5) as pose:
    while cap.isOpened():
        success, image = cap.read()
        if not success:
            print("Ignoring empty camera frame.")
            # If loading a video, use 'break' instead of 'continue'.
            # continue
            break

        # Flip the image horizontally for a later selfie-view display, and convert
        # the BGR image to RGB.
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        # To improve performance, optionally mark the image as not writeable to
        # pass by reference.
        image.flags.writeable = False
        results = pose.process(image)

        # Draw the pose annotation on the image.
        image.flags.writeable = True
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

        image_hight, image_width, _ = image.shape

        params = sp.get_params(results, all=True)
        print(params)

        mp_drawing.draw_landmarks(
            image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)

        coords = landmarks_list_to_array(results.pose_landmarks, image.shape)
        # label_params(image, params, coords)

        cv2.imshow('MediaPipe Pose', image)

        if cv2.waitKey(5) & 0xFF == 27:
            break

cap.release()
cv2.destroyAllWindows()