Files
JustTwirk/main.py
2025-11-28 08:31:35 +00:00

162 lines
4.9 KiB
Python

import pickle
import sys
from ultralytics import YOLO
import cv2
import time
import utils
from calculate import normalize_pose, compare_poses_boolean
from draw import draw_new
from utils import find_closest
from video_methods import initialize_method
model = YOLO("yolo11x-pose.pt")
if len(sys.argv) == 2:
method_type = sys.argv[1]
else:
print("Podaj argument 'cam', albo 'net'.")
exit(1)
method = initialize_method(method_type)
do_pose_shot = False
def click_event(event, x, y, flags, param):
global do_pose_shot
if event == cv2.EVENT_LBUTTONDOWN: # lewy przycisk myszy
do_pose_shot = not do_pose_shot
def main():
last_time = time.time()
currTimeIndex = 0
currIndex = None
currMove = None
currStatus = "Zacznij tanczyc"
mehCount = 0
goodCount = 0
failCount = 0
failRate = 2
moves = []
with open('moves.pkl', 'rb') as f: # 'rb' = read binary
moves = pickle.load(f)
startValue = moves[0][0]
totalCount = len(moves)
for i, move in enumerate(moves):
moves[i] = (move[0] - startValue, move[1], move[2])
print(moves)
while True:
frame = method.receive_frame()
frame = cv2.flip(frame, 1)
results = model(frame, verbose=False)
current_time = time.time()
delta = current_time - last_time
last_time = current_time
fps = 1 / delta if delta > 0 else float('inf')
# print(f"\rDelta: {delta:.4f}s, FPS: {fps:.2f}", end="")
if len(results) == 0:
continue
result = results[0]
kpts = result.keypoints.data[0] if len(result.keypoints.data) else None
if kpts is None:
continue
img = frame
normalized = normalize_pose(result.keypoints.xy.cpu().numpy()[0])
draw = utils.normalize(result.keypoints.xy.cpu().numpy()[0])
cv2.imshow('you', draw_new(draw * 100 + 100))
if currTimeIndex != 0 and moves.index(find_closest(moves, time.time() - currTimeIndex)) == len(moves) - 1:
mehCount = totalCount - failCount - goodCount
print(
f"PODSUMOWANIE: FAIL {failCount} MEH: {mehCount} PERFECT: {goodCount} PERCENTAGE: {(goodCount + (0.95 * mehCount)) / totalCount * 100}%")
exit(1)
if currMove is None:
if compare_poses_boolean(moves[0][1], normalized):
currIndex = 1
currTimeIndex = time.time()
deltaTime = time.time()
currStatus = f"Zaczoles tanczyc {currIndex}"
currMove = moves[0]
# thread = Thread(target=print_animation, args=(moves, False))
# thread.start()
else:
changed = False
closest = find_closest(moves, time.time() - currTimeIndex)
cv2.imshow('Dots', draw_new(closest[2]))
if abs((time.time() - currTimeIndex) - moves[currIndex][0]) > failRate:
currStatus = f"FAIL!"
failCount += 1
if compare_poses_boolean(closest[1], normalized):
# delays += (time.time() - deltaTime - moves[0][0]) * 1000
# delaysCount += 1
currStatus = f"SUPER! {currIndex} Zostalo {len(moves)} Delay {(time.time() - currTimeIndex - closest[0]) / 1000}ms"
deltaTime = time.time()
currIndex = moves.index(closest) + 1
goodCount += 1
changed = True
if not changed and compare_poses_boolean(moves[currIndex][1], normalized):
# delays += (time.time() - deltaTime - moves[0][0]) * 1000
# delaysCount += 1
currStatus = f"SUPER! {currIndex} Zostalo {len(moves)} Delay {(time.time() - currTimeIndex - closest[0]) / 1000}ms"
deltaTime = time.time()
changed = True
currIndex += 1
goodCount += 1
# if do_pose_shot:
# moves.append((time.time() - startTime, normalize_pose(result.keypoints.xy.cpu().numpy()[0]), result.keypoints.xy.cpu()[0]))
# elif len(moves) != 0:
# with open('moves.pkl', 'wb') as f: # 'wb' = write binary
# pickle.dump
# (moves, f)
#
# exit(1)
cv2.putText(
img, # obraz
currStatus, # tekst
(50, 100), # pozycja (x, y) lewego dolnego rogu tekstu
cv2.FONT_HERSHEY_SIMPLEX, # czcionka
1, # rozmiar (skalowanie)
(0, 0, 255), # kolor (BGR) - tutaj czerwony
2, # grubość linii
cv2.LINE_AA # typ antyaliasingu
)
cv2.imshow('Klatka z kamerki', img)
cv2.setMouseCallback('Klatka z kamerki', click_event)
cv2.waitKey(1) # Czekaj na naciśnięcie klawisza
main()