Spaces:
Sleeping
Sleeping
| import torch | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| import io | |
| import matplotlib | |
| from mpl_toolkits.mplot3d.art3d import Poly3DCollection | |
| import mpl_toolkits.mplot3d.axes3d as p3 | |
| from textwrap import wrap | |
| import imageio | |
| def plot_3d_motion(args, figsize=(10, 10), fps=120, radius=4): | |
| matplotlib.use('Agg') | |
| joints, out_name, title = args | |
| data = joints.copy().reshape(len(joints), -1, 3) | |
| nb_joints = joints.shape[1] | |
| smpl_kinetic_chain = [[0, 11, 12, 13, 14, 15], [0, 16, 17, 18, 19, 20], [0, 1, 2, 3, 4], [3, 5, 6, 7], [3, 8, 9, 10]] if nb_joints == 21 else [[0, 2, 5, 8, 11], [0, 1, 4, 7, 10], [0, 3, 6, 9, 12, 15], [9, 14, 17, 19, 21], [9, 13, 16, 18, 20]] | |
| limits = 1000 if nb_joints == 21 else 2 | |
| 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] | |
| def update(index): | |
| def init(): | |
| ax.set_xlim(-limits, limits) | |
| ax.set_ylim(-limits, limits) | |
| ax.set_zlim(0, limits) | |
| 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) | |
| fig = plt.figure(figsize=(480/96., 320/96.), dpi=96) if nb_joints == 21 else plt.figure(figsize=(10, 10), dpi=96) | |
| if title is not None : | |
| wraped_title = '\n'.join(wrap(title, 40)) | |
| fig.suptitle(wraped_title, fontsize=16) | |
| ax = p3.Axes3D(fig) | |
| init() | |
| ax.lines = [] | |
| ax.collections = [] | |
| ax.view_init(elev=110, 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(smpl_kinetic_chain, 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([]) | |
| if out_name is not None : | |
| plt.savefig(out_name, dpi=96) | |
| plt.close() | |
| else : | |
| io_buf = io.BytesIO() | |
| fig.savefig(io_buf, format='raw', dpi=96) | |
| io_buf.seek(0) | |
| # print(fig.bbox.bounds) | |
| arr = np.reshape(np.frombuffer(io_buf.getvalue(), dtype=np.uint8), | |
| newshape=(int(fig.bbox.bounds[3]), int(fig.bbox.bounds[2]), -1)) | |
| io_buf.close() | |
| plt.close() | |
| return arr | |
| out = [] | |
| for i in range(frame_number) : | |
| out.append(update(i)) | |
| out = np.stack(out, axis=0) | |
| return torch.from_numpy(out) | |
| def draw_to_batch(smpl_joints_batch, title_batch=None, outname=None) : | |
| batch_size = len(smpl_joints_batch) | |
| out = [] | |
| for i in range(batch_size) : | |
| out.append(plot_3d_motion([smpl_joints_batch[i], None, title_batch[i] if title_batch is not None else None])) | |
| if outname is not None: | |
| imageio.mimsave(outname[i], np.array(out[-1]), fps=20) | |
| out = torch.stack(out, axis=0) | |
| return out | |