File size: 3,733 Bytes
4ba516b
cb0074d
 
 
 
 
4ba516b
cb0074d
 
 
 
 
2a4205b
 
 
cb0074d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2a4205b
cb0074d
 
 
 
 
 
 
 
 
 
 
 
 
2a4205b
 
 
cb0074d
 
 
 
 
 
 
 
 
 
 
 
 
e2dff2b
cb0074d
2a4205b
cb0074d
 
2a4205b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cb0074d
 
 
 
 
 
2a4205b
cb0074d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2a4205b
 
 
 
cb0074d
 
 
 
 
 
 
2a4205b
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
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)