Files
JustTwirk/3ddisplay_replay_smoothed.py
2025-12-08 20:25:20 +01:00

59 lines
1.5 KiB
Python

from scipy.signal import savgol_filter
import numpy as np
import pickle
from matplotlib import pyplot as plt
with open("replay_xyz.pkl", "rb") as f:
points3DList = pickle.load(f)
skeleton = [
[0, 1], [0, 2],
[1, 3], [2, 4],
[5, 7], [7, 9],
[6, 8], [8, 10],
[5, 6],
[5, 11], [6, 12],
[11, 12],
[11, 13], [13, 15],
[12, 14], [14, 16]
]
keys_sorted = sorted(points3DList.keys())
points_sequence = np.array([points3DList[k] for k in keys_sorted]) # (frames, points, 3)
# --- Filtr Savitzky-Golaya ---
window_length = 7 # musi być nieparzyste
polyorder = 2
smoothed_sequence = savgol_filter(points_sequence, window_length=window_length,
polyorder=polyorder, axis=0, mode='nearest')
plt.ion()
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
points_plot = ax.scatter([], [], [], c='r', marker='o', s=50)
lines_plot = [ax.plot([0,0],[0,0],[0,0], c='b')[0] for _ in skeleton]
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
ax.set_xlim(-0.6, 0.4)
ax.set_ylim(1.2, 2.2)
ax.set_zlim(-0.5, 1.1)
ax.view_init(elev=20, azim=-60)
for frame_points in smoothed_sequence:
X = frame_points[:,0] - 0.25
Z = -frame_points[:,1] + 0.5
Y = frame_points[:,2]
points_plot._offsets3d = (X, Y, Z)
for idx, (i, j) in enumerate(skeleton):
lines_plot[idx].set_data([X[i], X[j]], [Y[i], Y[j]])
lines_plot[idx].set_3d_properties([Z[i], Z[j]])
fig.canvas.draw()
fig.canvas.flush_events()
plt.pause(0.001)