| import cv2, io |
| import numpy as np |
| import torch |
| import matplotlib.pyplot as plt |
| import matplotlib |
| from mpl_toolkits.mplot3d.art3d import Poly3DCollection |
| import mpl_toolkits.mplot3d.axes3d as p3 |
| from textwrap import wrap |
| from ..utils.rotation_conversions import rotation_6d_to_matrix |
|
|
| def _vis_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] |
| |
|
|
| 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): |
| |
| 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, auto_add_to_figure=False) |
| fig.add_axes(ax) |
| |
| init() |
| |
| for coll in ax.lines: |
| coll.remove() |
| for coll in ax.collections: |
| coll.remove() |
| ax.view_init(elev=110, azim=-90) |
| ax.dist = 7.5 |
| plot_xzPlane(MINS[0] - trajec[index, 0], MAXS[0] - trajec[index, 0], 0, |
| MINS[2] - trajec[index, 1], MAXS[2] - trajec[index, 1]) |
|
|
| 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') |
|
|
| for i, (chain, color) in enumerate(zip(smpl_kinetic_chain, colors)): |
| linewidth = 4.0 if i < 5 else 2.0 |
| ax.plot3D(data[index, chain, 0], data[index, chain, 1], data[index, chain, 2], |
| linewidth=linewidth, color=color) |
|
|
| 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) |
| 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 vis_3d_motion(smpl_joints_batch, title_batch=None, outname=None, fps=30) : |
| out = [] |
| for i in range(len(smpl_joints_batch)) : |
| out.append(_vis_3d_motion([smpl_joints_batch[i], None, title_batch[i] if title_batch is not None else None])) |
| if outname is not None: |
| import imageio.v3 as iio |
| images = np.array(out[-1]).astype(np.uint8)[..., :3] |
|
|
| with iio.imopen(outname[i], "w", plugin="pyav") as writer: |
| writer.init_video_stream("libx264", fps=int(fps)) |
| writer._video_stream.options = {"crf": str(17)} |
| writer.write(images) |
|
|
| def accumulate_rotations(relative_rotations): |
| R_total = [relative_rotations[0]] |
| for R_rel in relative_rotations[1:]: |
| R_total.append(np.matmul(R_rel, R_total[-1])) |
| return np.array(R_total) |
|
|
| def recover_from_local_position(final_x, njoint): |
|
|
| if final_x.ndim == 3: |
| bs, nfrm, _ = final_x.shape |
| is_batched = True |
| else: |
| nfrm, _ = final_x.shape |
| bs = 1 |
| is_batched = False |
| final_x = final_x.reshape(1, *final_x.shape) |
|
|
| |
| positions_no_heading = final_x[:,:,8:8+3*njoint].reshape(bs, nfrm, njoint, 3) |
| velocities_root_xy_no_heading = final_x[:,:,:2] |
| global_heading_diff_rot = final_x[:,:,2:8] |
|
|
| positions_with_heading = [] |
| for b in range(bs): |
| |
| global_heading_rot = accumulate_rotations(rotation_6d_to_matrix(torch.from_numpy(global_heading_diff_rot[b])).numpy()) |
| inv_global_heading_rot = np.transpose(global_heading_rot, (0, 2, 1)) |
| |
| |
| curr_pos_with_heading = np.matmul(np.repeat(inv_global_heading_rot[:, None,:, :], njoint, axis=1), |
| positions_no_heading[b][...,None]).squeeze(-1) |
|
|
| |
| velocities_root_xyz_no_heading = np.zeros((velocities_root_xy_no_heading[b].shape[0], 3)) |
| velocities_root_xyz_no_heading[:, 0] = velocities_root_xy_no_heading[b, :, 0] |
| velocities_root_xyz_no_heading[:, 2] = velocities_root_xy_no_heading[b, :, 1] |
| velocities_root_xyz_no_heading[1:, :] = np.matmul(inv_global_heading_rot[:-1], |
| velocities_root_xyz_no_heading[1:, :,None]).squeeze(-1) |
|
|
| root_translation = np.cumsum(velocities_root_xyz_no_heading, axis=0) |
|
|
| |
| curr_pos_with_heading[:, :, 0] += root_translation[:, 0:1] |
| curr_pos_with_heading[:, :, 2] += root_translation[:, 2:] |
| |
| positions_with_heading.append(curr_pos_with_heading) |
|
|
| positions_with_heading = np.stack(positions_with_heading, axis=0) |
|
|
| if not is_batched: |
| positions_with_heading = positions_with_heading.squeeze(0) |
|
|
| return positions_with_heading |
