File size: 4,016 Bytes
b887ad8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
import math
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation, FFMpegFileWriter
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import mpl_toolkits.mplot3d.axes3d as p3
# import cv2
def list_cut_average(ll, intervals):
if intervals == 1:
return ll
bins = math.ceil(len(ll) * 1.0 / intervals)
ll_new = []
for i in range(bins):
l_low = intervals * i
l_high = l_low + intervals
l_high = l_high if l_high < len(ll) else len(ll)
ll_new.append(np.mean(ll[l_low:l_high]))
return ll_new
def plot_3d_motion(save_path, kinematic_tree, joints, title, figsize=(10, 10), fps=120, radius=4):
matplotlib.use('Agg')
title_sp = title.split(' ')
if len(title_sp) > 20:
title = '\n'.join([' '.join(title_sp[:10]), ' '.join(title_sp[10:20]), ' '.join(title_sp[20:])])
elif len(title_sp) > 10:
title = '\n'.join([' '.join(title_sp[:10]), ' '.join(title_sp[10:])])
def init():
ax.set_xlim3d([-radius / 4, radius / 4])
ax.set_ylim3d([0, radius / 2])
ax.set_zlim3d([0, radius / 2])
# print(title)
fig.suptitle(title, fontsize=20)
ax.grid(b=False)
def plot_xzPlane(minx, maxx, miny, minz, maxz):
## Plot a plane XZ
verts = [
[minx, miny, minz],
[minx, miny, maxz],
[maxx, miny, maxz],
[maxx, miny, minz]
]
xz_plane = Poly3DCollection([verts])
xz_plane.set_facecolor((0.5, 0.5, 0.5, 0.5))
ax.add_collection3d(xz_plane)
# return ax
# (seq_len, joints_num, 3)
data = joints.copy().reshape(len(joints), -1, 3)
fig = plt.figure(figsize=figsize)
ax = p3.Axes3D(fig)
init()
MINS = data.min(axis=0).min(axis=0)
MAXS = data.max(axis=0).max(axis=0)
colors = ['red', 'blue', 'black', 'red', 'blue',
'darkblue', 'darkblue', 'darkblue', 'darkblue', 'darkblue',
'darkred', 'darkred', 'darkred', 'darkred', 'darkred']
frame_number = data.shape[0]
# print(data.shape)
height_offset = MINS[1]
data[:, :, 1] -= height_offset
trajec = data[:, 0, [0, 2]]
data[..., 0] -= data[:, 0:1, 0]
data[..., 2] -= data[:, 0:1, 2]
# print(trajec.shape)
def update(index):
# print(index)
ax.lines = []
ax.collections = []
ax.view_init(elev=120, azim=-90)
ax.dist = 7.5
# ax =
plot_xzPlane(MINS[0] - trajec[index, 0], MAXS[0] - trajec[index, 0], 0, MINS[2] - trajec[index, 1],
MAXS[2] - trajec[index, 1])
# ax.scatter(data[index, :22, 0], data[index, :22, 1], data[index, :22, 2], color='black', s=3)
# if index > 1:
# ax.plot3D(trajec[:index, 0] - trajec[index, 0], np.zeros_like(trajec[:index, 0]),
# trajec[:index, 1] - trajec[index, 1], linewidth=1.0,
# color='blue')
# ax = plot_xzPlane(ax, MINS[0], MAXS[0], 0, MINS[2], MAXS[2])
for i, (chain, color) in enumerate(zip(kinematic_tree, colors)):
# print(color)
if i < 5:
linewidth = 4.0
else:
linewidth = 2.0
ax.plot3D(data[index, chain, 0], data[index, chain, 1], data[index, chain, 2], linewidth=linewidth,
color=color)
# print(trajec[:index, 0].shape)
plt.axis('off')
ax.set_xticklabels([])
ax.set_yticklabels([])
ax.set_zticklabels([])
ani = FuncAnimation(fig, update, frames=frame_number, interval=1000 / fps, repeat=False)
from matplotlib import rcParams
rcParams['animation.ffmpeg_path'] = '/home/chenlinghao/software/miniconda3/envs/py10/bin/ffmpeg'
writer = FFMpegFileWriter(fps=fps)
ani.save(save_path, writer=writer)
plt.close()
|