PoseEstimationProject/AITrainer.py at master · 1yakub/PoseEstimationProject · GitHub

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
import time

import cv2 as cv
import numpy as np
import PoseModule as pm

import tkinter as tk
from tkinter import ttk

# choose workout
def click(x):
    global flag
    win.destroy()
    flag = x


win = tk.Tk()
win.title("AI Trainer")
Lbl = ttk.Label(win, text="Choose your workout")
Lbl.pack()
flag = ""

action1 = ttk.Button(win, text="Bicep Curl", command=lambda: click("Bicep Curl"))
action2 = ttk.Button(win, text="Squat", command=lambda: click("Squat"))
action1.pack()
action2.pack()

win.mainloop()


def detect_bicep_curls():
    global per, bar
    # left arm
    angle = detector.get_angle(img, 11, 13, 15, draw=True)
    # right arm
    detector.get_angle(img, 12, 14, 16, draw=True)
    per = np.interp(angle, (50, 160), (100, 0))
    bar = np.interp(angle, (50, 160), (300, 600))


def detect_squats():
    global per, bar
    # left side
    angle = detector.get_angle(img, 11, 23, 25, draw=True)
    per = np.interp(angle, (65, 165), (100, 0))
    bar = np.interp(angle, (65, 165), (300, 600))

cap = cv.VideoCapture(0)
address ="http://192.168.0.100:8080/video"
# cap.open(address)
fps = cap.get(cv.CAP_PROP_FPS)
delay = round(1000 / fps)

detector = pm.PoseDetector()
count = 0
direction = 0
sec = 5

while cap.isOpened():

    #countdown before it starts
    prev = time.time()
    while sec > 0:
        ret, img = cap.read()
        cv.putText(img, str(sec), (img.shape[1]//2, img.shape[0]//2), cv.FONT_HERSHEY_SIMPLEX, 4, (0, 0, 255), 5)
        cv.imshow("image", img)
        cur = time.time()
        if cur - prev >= 1:
            prev = cur
            sec -= 1
        cv.waitKey(1)

    # detect workout
    ret, img = cap.read()

    if not ret:
        break

    img = detector.get_pose(img, False)
    lmList = detector.get_position(img,False)

    if len(lmList) != 0:
        if flag == "Bicep Curl":
            detect_bicep_curls()
        elif flag == "Squat":
            detect_squats()

        # counting reps
        color = (255,255,0)
        if per == 100:
            color = (0, 0, 255)
            if direction == 0:
                count += 0.5
                direction = 1
        if per == 0:
            color = (0, 0, 255)
            if direction == 1:
                count += 0.5
                direction = 0
        cv.putText(img, str(int(count)), (50,100), cv.FONT_HERSHEY_SIMPLEX, 4, (0, 0, 0), 5)

        cv.rectangle(img, (50, 300), (150,600),(255,255,255),cv.FILLED)
        cv.rectangle(img, (50, int(bar)), (150,600),color,cv.FILLED)
        cv.putText(img, f'{int(per)}%', (50,280), cv.FONT_HERSHEY_SIMPLEX, 1, color, 2)

    cv.imshow("image", img)

    if cv.waitKey(1) == 27:
        break