Spaces:
Runtime error
Runtime error
import cv2 | |
from time import time | |
import pickle as pk | |
import mediapipe as mp | |
import pandas as pd | |
import multiprocessing as mtp | |
from recommendations import check_pose_angle | |
from landmarks import extract_landmarks | |
from calc_angles import rangles | |
def init_cam(): | |
cam = cv2.VideoCapture(0) # Use 0 for default camera | |
if not cam.isOpened(): | |
raise ValueError("Could not open camera.") | |
cam.set(cv2.CAP_PROP_AUTOFOCUS, 0) | |
cam.set(cv2.CAP_PROP_FOCUS, 360) | |
cam.set(cv2.CAP_PROP_BRIGHTNESS, 130) | |
cam.set(cv2.CAP_PROP_SHARPNESS, 125) | |
cam.set(cv2.CAP_PROP_FRAME_WIDTH, 1280) | |
cam.set(cv2.CAP_PROP_FRAME_HEIGHT, 720) | |
return cam | |
def get_pose_name(index): | |
names = { | |
0: "Adho Mukha Svanasana", | |
1: "Phalakasana", | |
2: "Utkata Konasana", | |
3: "Vrikshasana", | |
} | |
return str(names.get(index, "Unknown Pose")) | |
def init_dicts(): | |
landmarks_points = { | |
"nose": 0, | |
"left_shoulder": 11, "right_shoulder": 12, | |
"left_elbow": 13, "right_elbow": 14, | |
"left_wrist": 15, "right_wrist": 16, | |
"left_hip": 23, "right_hip": 24, | |
"left_knee": 25, "right_knee": 26, | |
"left_ankle": 27, "right_ankle": 28, | |
"left_heel": 29, "right_heel": 30, | |
"left_foot_index": 31, "right_foot_index": 32, | |
} | |
landmarks_points_array = {key: [] for key in landmarks_points.keys()} | |
col_names = [f"{name}_{axis}" for name in landmarks_points.keys() for axis in ["x", "y", "z", "v"]] | |
return col_names, landmarks_points_array | |
def cv2_put_text(image, message): | |
cv2.putText( | |
image, | |
message, | |
(50, 50), | |
cv2.FONT_HERSHEY_SIMPLEX, | |
2, | |
(255, 0, 0), | |
5, | |
cv2.LINE_AA | |
) | |
def destroy(cam): | |
cam.release() | |
cv2.destroyAllWindows() | |
if __name__ == "__main__": | |
try: | |
cam = init_cam() | |
model = pk.load(open("./models/4_poses.model", "rb")) | |
cols, landmarks_points_array = init_dicts() | |
angles_df = pd.read_csv("./csv_files/4_angles_poses_angles.csv") | |
mp_drawing = mp.solutions.drawing_utils | |
mp_pose = mp.solutions.pose | |
while True: | |
result, image = cam.read() | |
if not result: | |
print("Failed to grab frame.") | |
break | |
flipped = cv2.flip(image, 1) | |
resized_image = cv2.resize( | |
flipped, | |
(640, 360), | |
interpolation=cv2.INTER_AREA | |
) | |
err, df, landmarks = extract_landmarks( | |
resized_image, | |
mp_pose, | |
cols | |
) | |
if not err: | |
prediction = model.predict(df) | |
probabilities = model.predict_proba(df) | |
mp_drawing.draw_landmarks( | |
flipped, | |
landmarks, | |
mp_pose.POSE_CONNECTIONS | |
) | |
if probabilities[0, prediction[0]] > 0.85: | |
cv2_put_text( | |
flipped, | |
get_pose_name(prediction[0]) | |
) | |
angles = rangles(df, landmarks_points_array) | |
suggestions = check_pose_angle( | |
prediction[0], angles, angles_df) | |
cv2_put_text( | |
flipped, | |
suggestions[0] | |
) | |
else: | |
cv2_put_text( | |
flipped, | |
"No Pose Detected" | |
) | |
cv2.imshow("Frame", flipped) | |
if cv2.waitKey(1) & 0xFF == ord('q'): | |
break | |
finally: | |
destroy(cam) | |