93 lines
2.3 KiB
Python
93 lines
2.3 KiB
Python
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
|
|
|
|
# ---------------------
|
|
# Wideo wejściowe
|
|
# ---------------------
|
|
cap = cv2.VideoCapture("input.mp4")
|
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
|
width = 640
|
|
height = 640
|
|
|
|
# ---------------------
|
|
# Wideo wyjściowe
|
|
# ---------------------
|
|
fourcc = cv2.VideoWriter_fourcc(*"MJPG")
|
|
|
|
|
|
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)
|
|
|
|
frame_id = 0
|
|
|
|
while True:
|
|
ok, frame = cap.read()
|
|
if not ok:
|
|
break
|
|
|
|
rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
results = pose.process(rgb)
|
|
|
|
# -----------------------------------------
|
|
# 3D landmarki: pose_world_landmarks
|
|
# -----------------------------------------
|
|
if results.pose_world_landmarks:
|
|
lm = results.pose_world_landmarks.landmark
|
|
|
|
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 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)
|
|
|
|
frame_id += 1
|
|
|
|
cap.release()
|
|
out.release()
|
|
print("Zapisano: output.mp4")
|