v2 working

moves_3d_mp4.py

@@ -1,33 +1,27 @@
-import cv2
-import mediapipe as mp
-import matplotlib
-matplotlib.use("Agg")  # <-- ważne: wyłącza GUI
-import matplotlib.pyplot as plt
-from mpl_toolkits.mplot3d import Axes3D
-import numpy as np
+import pickle
+import time
 
-# ---------------------
-# Wideo wejściowe
-# ---------------------
+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 = 640
-height = 640
-
-# ---------------------
-# Wideo wyjściowe
-# ---------------------
-fourcc = cv2.VideoWriter_fourcc(*"MJPG")
-
+width = 1920
+height = 1080
 
+# Ustawienia zapisu wideo
+fourcc = cv2.VideoWriter_fourcc(*"avc1")
 out = cv2.VideoWriter("output.mp4", fourcc, fps, (width, height))
 
-# ---------------------
-# MediaPipe Pose
-# ---------------------
-mp_pose = mp.solutions.pose
-pose = mp_pose.Pose(static_image_mode=False, model_complexity=1)
+# Wczytanie modelu YOLOv8 Pose
+model = YOLO("yolo11x-pose.pt", verbose=False)  # Twój model pose
 
+moves = []
+started = False
 frame_id = 0
 
 while True:
@@ -35,58 +29,32 @@ while True:
     if not ok:
         break
 
-    rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-    results = pose.process(rgb)
+    # Skalowanie do 640x640
+    frame_resized = cv2.resize(frame, (width, height))
 
-    # -----------------------------------------
-    # 3D landmarki: pose_world_landmarks
-    # -----------------------------------------
-    if results.pose_world_landmarks:
-        lm = results.pose_world_landmarks.landmark
+    # Wykrywanie poz
+    results = model.predict(frame_resized, verbose=False)
 
-        xs = np.array([p.x for p in lm])
-        ys = np.array([p.y for p in lm])
-        zs = np.array([p.z for p in lm])
+    # 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
 
-        # -----------------------------
-        # RYSOWANIE 3D w Matplotlib
-        # -----------------------------
-        fig = plt.figure(figsize=(6.4, 6.4), dpi=100)
-        ax = fig.add_subplot(111, projection="3d")
-
-        ax.scatter(xs, zs, ys, s=20)
-
-        ax.set_xlim([-1, 1])
-        ax.set_ylim([-1, 1])
-        ax.set_zlim([-1, 1])
-
-        ax.set_xlabel("X")
-        ax.set_ylabel("Y")
-        ax.set_zlabel("Z")
-        ax.invert_zaxis()
-
-        # -----------------------------------------
-        # Konwersja wykresu Matplotlib → klatka do MP4
-        # -----------------------------------------
-        fig.canvas.draw()
-        renderer = fig.canvas.get_renderer()
-
-        w, h = fig.canvas.get_width_height()
-
-        buf = renderer.buffer_rgba()
-        plot_img = np.frombuffer(buf, dtype=np.uint8).reshape((h, w, 4))[:, :, :3]
-
-
-        plt.close(fig)
-
-        # Dopasowanie rozmiaru do wideo
-        plot_img = cv2.resize(plot_img, (width, height))
-        plot_img = cv2.cvtColor(plot_img, cv2.COLOR_RGB2BGR)
-
-        out.write(plot_img)
+            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")