JustTwirk/main.py

import pickle
import sys

import regex as re
import numpy as np
from ultralytics import YOLO
import cv2
import time

from calculate import normalize_pose, compare_poses_boolean
from draw import draw_new

model = YOLO("yolo11x-pose.pt")

if len(sys.argv) == 1:
    method = sys.argv[1]
else:
    print("Podaj argument 'cam', albo 'network'.")
    exit(1)

cap = cv2.VideoCapture(0)

if not cap.isOpened():
    print("Nie można otworzyć kamerki")
    exit(1)

last_time = time.time()


startTime = time.time()
stage = 0

pose = normalize_pose(np.array([[     353.17,      107.28],
 [      363.3, 96.435],
 [      347.1, 98.647],
 [     390.12, 99.096],
 [     346.09, 103.63],
 [     425.77, 149.95],
 [     327.92, 153.12],
 [     495.16, 169.82],
 [     260.86,  166.3],
 [     565.53, 182.68],
 [     192.34, 170.58],
 [     393.83, 316.99],
 [     337.05, 316.41],
 [     388.36, 433.37],
 [     333.88, 431.89],
 [     383.58, 479.16],
 [     330.89,    480]]))

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

currTimeIndex = 0
currIndex = None
currMove = None
currStatus = "Zacznij tanczyc"

mehCount = 0
goodCount = 0
failCount = 0
failRate = 2

def wczytaj_dane_z_txt(sciezka):
    wynik = []
    with open(sciezka, "r") as f:
        zawartosc = f.read()

    # Znajdź wszystkie krotki w formacie (float, array([...]))
    pattern = re.compile(r"\(([^,]+),\s*array\((\[.*?\]),\s*dtype=float32\)\)")
    matches = pattern.findall(zawartosc)

    for m in matches:
        liczba = float(m[0])
        tablica = np.array(eval(m[1]), dtype=np.float32)
        wynik.append((liczba, tablica))

    return wynik

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])

def find_closest(target):
    global moves
    return min(moves, key=lambda t: abs(t[0] - target))

while True:
    ret, frame = cap.read()
    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="")

    for result in results:
        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])
        cv2.imshow('you', draw_new(result.keypoints.xy.cpu()[0]))

        if currTimeIndex != 0 and moves.index(find_closest(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(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


# Access the results
for result in results:
    annotated_frame = result.plot()  # zwraca obraz z naniesionymi keypoints

    # Wyświetlenie obrazu przy użyciu OpenCV
    cv2.imshow("Pose", annotated_frame)
    cv2.waitKey(0)
cv2.destroyAllWindows()