61 lines
1.6 KiB
Python
61 lines
1.6 KiB
Python
import pickle
|
|
import time
|
|
|
|
import cv2
|
|
from ultralytics import YOLO
|
|
|
|
from calculate import normalize_pose
|
|
from utils import normalize
|
|
|
|
# Wczytanie wideo
|
|
cap = cv2.VideoCapture("input.mp4")
|
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
|
width = 1920
|
|
height = 1080
|
|
|
|
# Ustawienia zapisu wideo
|
|
fourcc = cv2.VideoWriter_fourcc(*"avc1")
|
|
out = cv2.VideoWriter("output.mp4", fourcc, fps, (width, height))
|
|
|
|
# Wczytanie modelu YOLOv8 Pose
|
|
model = YOLO("yolo11x-pose.pt", verbose=False) # Twój model pose
|
|
|
|
moves = []
|
|
started = False
|
|
frame_id = 0
|
|
|
|
while True:
|
|
ok, frame = cap.read()
|
|
if not ok:
|
|
break
|
|
|
|
# Skalowanie do 640x640
|
|
frame_resized = cv2.resize(frame, (width, height))
|
|
|
|
# Wykrywanie poz
|
|
results = model.predict(frame_resized, verbose=False)
|
|
|
|
# Rysowanie punktów 2D
|
|
for result in results:
|
|
if result.keypoints is not None and len(result.keypoints.xy) > 0:
|
|
for keypoint in result.keypoints.xy[0]: # keypoints[0] bo dla jednej osoby
|
|
if not started:
|
|
frame_id = 0
|
|
started = True
|
|
x, y = keypoint # współrzędne + confidence
|
|
x = int(x)
|
|
y = int(y)
|
|
cv2.circle(frame_resized, (x, y), 5, (0, 255, 0), -1) # zielone kropki
|
|
|
|
moves.append((frame_id * 1 / fps, normalize_pose(result.keypoints.xy.cpu().numpy()[0]), normalize(result.keypoints.xy.cpu()[0]) * 160 + 250))
|
|
|
|
out.write(frame_resized)
|
|
frame_id += 1
|
|
|
|
with open('moves2.pkl', 'wb') as f: # 'wb' = write binary
|
|
pickle.dump(moves, f)
|
|
|
|
cap.release()
|
|
out.release()
|
|
print("Zapisano: output.mp4")
|