| import os |
| import pathlib |
|
|
| import matplotlib.pyplot as plt |
| import numpy as np |
|
|
|
|
| def project_root_position(position_arr: np.array, file_name: str): |
| """ |
| Take batch of root arrays and porject it on 2D plane |
| |
| N: samples |
| L: trajectory length |
| J: joints |
| |
| position_arr: [N,L,J,3] |
| """ |
|
|
| root_joints = position_arr[:, :, 0] |
|
|
| x_pos = root_joints[:, :, 0] |
| y_pos = root_joints[:, :, 2] |
|
|
| fig = plt.figure() |
|
|
| for i in range(x_pos.shape[1]): |
|
|
| if i == 0: |
| plt.scatter(x_pos[:, i], y_pos[:, i], c="b") |
| elif i == x_pos.shape[1] - 1: |
| plt.scatter(x_pos[:, i], y_pos[:, i], c="r") |
| else: |
| plt.scatter(x_pos[:, i], y_pos[:, i], c="k", marker="*", s=1) |
|
|
| plt.title(f"Root Position: {file_name}") |
| plt.xlabel("X Axis") |
| plt.ylabel("Y Axis") |
| plt.xlim((-300, 300)) |
| plt.ylim((-300, 300)) |
| plt.grid() |
| plt.savefig(f"{file_name}.png", dpi=200) |
|
|
|
|
| def plot_single_pose( |
| pose, |
| frame_idx, |
| skeleton, |
| save_dir, |
| prefix, |
| ): |
| fig = plt.figure() |
| ax = fig.add_subplot(111, projection="3d") |
|
|
| parent_idx = skeleton.parents() |
|
|
| for i, p in enumerate(parent_idx): |
| if i > 0: |
| ax.plot( |
| [pose[i, 0], pose[p, 0]], |
| [pose[i, 2], pose[p, 2]], |
| [pose[i, 1], pose[p, 1]], |
| c="k", |
| ) |
|
|
| x_min = pose[:, 0].min() |
| x_max = pose[:, 0].max() |
|
|
| y_min = pose[:, 1].min() |
| y_max = pose[:, 1].max() |
|
|
| z_min = pose[:, 2].min() |
| z_max = pose[:, 2].max() |
|
|
| ax.set_xlim(x_min, x_max) |
| ax.set_xlabel("$X$ Axis") |
|
|
| ax.set_ylim(z_min, z_max) |
| ax.set_ylabel("$Y$ Axis") |
|
|
| ax.set_zlim(y_min, y_max) |
| ax.set_zlabel("$Z$ Axis") |
|
|
| plt.draw() |
|
|
| title = f"{prefix}: {frame_idx}" |
| plt.title(title) |
| prefix = prefix |
| pathlib.Path(save_dir).mkdir(parents=True, exist_ok=True) |
| plt.savefig(os.path.join(save_dir, prefix + str(frame_idx) + ".png"), dpi=60) |
| plt.close() |
|
|
|
|
| def plot_pose( |
| start_pose, |
| inbetween_pose, |
| target_pose, |
| frame_idx, |
| skeleton, |
| save_dir, |
| prefix, |
| ): |
| fig = plt.figure() |
| ax = fig.add_subplot(111, projection="3d") |
|
|
| parent_idx = skeleton.parents() |
|
|
| for i, p in enumerate(parent_idx): |
| if i > 0: |
| ax.plot( |
| [start_pose[i, 0], start_pose[p, 0]], |
| [start_pose[i, 2], start_pose[p, 2]], |
| [start_pose[i, 1], start_pose[p, 1]], |
| c="b", |
| ) |
| ax.plot( |
| [inbetween_pose[i, 0], inbetween_pose[p, 0]], |
| [inbetween_pose[i, 2], inbetween_pose[p, 2]], |
| [inbetween_pose[i, 1], inbetween_pose[p, 1]], |
| c="k", |
| ) |
| ax.plot( |
| [target_pose[i, 0], target_pose[p, 0]], |
| [target_pose[i, 2], target_pose[p, 2]], |
| [target_pose[i, 1], target_pose[p, 1]], |
| c="r", |
| ) |
|
|
| x_min = np.min( |
| [start_pose[:, 0].min(), inbetween_pose[:, 0].min(), target_pose[:, 0].min()] |
| ) |
| x_max = np.max( |
| [start_pose[:, 0].max(), inbetween_pose[:, 0].max(), target_pose[:, 0].max()] |
| ) |
|
|
| y_min = np.min( |
| [start_pose[:, 1].min(), inbetween_pose[:, 1].min(), target_pose[:, 1].min()] |
| ) |
| y_max = np.max( |
| [start_pose[:, 1].max(), inbetween_pose[:, 1].max(), target_pose[:, 1].max()] |
| ) |
|
|
| z_min = np.min( |
| [start_pose[:, 2].min(), inbetween_pose[:, 2].min(), target_pose[:, 2].min()] |
| ) |
| z_max = np.max( |
| [start_pose[:, 2].max(), inbetween_pose[:, 2].max(), target_pose[:, 2].max()] |
| ) |
|
|
| ax.set_xlim(x_min, x_max) |
| ax.set_xlabel("$X$ Axis") |
|
|
| ax.set_ylim(z_min, z_max) |
| ax.set_ylabel("$Y$ Axis") |
|
|
| ax.set_zlim(y_min, y_max) |
| ax.set_zlabel("$Z$ Axis") |
|
|
| plt.draw() |
|
|
| title = f"{prefix}: {frame_idx}" |
| plt.title(title) |
| prefix = prefix |
| pathlib.Path(save_dir).mkdir(parents=True, exist_ok=True) |
| plt.savefig(os.path.join(save_dir, prefix + str(frame_idx) + ".png"), dpi=60) |
| plt.close() |
|
|
|
|
| def plot_pose_with_stop( |
| start_pose, |
| inbetween_pose, |
| target_pose, |
| stopover, |
| frame_idx, |
| skeleton, |
| save_dir, |
| prefix, |
| ): |
| fig = plt.figure() |
| ax = fig.add_subplot(111, projection="3d") |
|
|
| parent_idx = skeleton.parents() |
|
|
| for i, p in enumerate(parent_idx): |
| if i > 0: |
| ax.plot( |
| [start_pose[i, 0], start_pose[p, 0]], |
| [start_pose[i, 2], start_pose[p, 2]], |
| [start_pose[i, 1], start_pose[p, 1]], |
| c="b", |
| ) |
| ax.plot( |
| [inbetween_pose[i, 0], inbetween_pose[p, 0]], |
| [inbetween_pose[i, 2], inbetween_pose[p, 2]], |
| [inbetween_pose[i, 1], inbetween_pose[p, 1]], |
| c="k", |
| ) |
| ax.plot( |
| [target_pose[i, 0], target_pose[p, 0]], |
| [target_pose[i, 2], target_pose[p, 2]], |
| [target_pose[i, 1], target_pose[p, 1]], |
| c="r", |
| ) |
|
|
| ax.plot( |
| [stopover[i, 0], stopover[p, 0]], |
| [stopover[i, 2], stopover[p, 2]], |
| [stopover[i, 1], stopover[p, 1]], |
| c="indigo", |
| ) |
|
|
| x_min = np.min( |
| [start_pose[:, 0].min(), inbetween_pose[:, 0].min(), target_pose[:, 0].min()] |
| ) |
| x_max = np.max( |
| [start_pose[:, 0].max(), inbetween_pose[:, 0].max(), target_pose[:, 0].max()] |
| ) |
|
|
| y_min = np.min( |
| [start_pose[:, 1].min(), inbetween_pose[:, 1].min(), target_pose[:, 1].min()] |
| ) |
| y_max = np.max( |
| [start_pose[:, 1].max(), inbetween_pose[:, 1].max(), target_pose[:, 1].max()] |
| ) |
|
|
| z_min = np.min( |
| [start_pose[:, 2].min(), inbetween_pose[:, 2].min(), target_pose[:, 2].min()] |
| ) |
| z_max = np.max( |
| [start_pose[:, 2].max(), inbetween_pose[:, 2].max(), target_pose[:, 2].max()] |
| ) |
|
|
| ax.set_xlim(x_min, x_max) |
| ax.set_xlabel("$X$ Axis") |
|
|
| ax.set_ylim(z_min, z_max) |
| ax.set_ylabel("$Y$ Axis") |
|
|
| ax.set_zlim(y_min, y_max) |
| ax.set_zlabel("$Z$ Axis") |
|
|
| plt.draw() |
|
|
| title = f"{prefix}: {frame_idx}" |
| plt.title(title) |
| prefix = prefix |
| pathlib.Path(save_dir).mkdir(parents=True, exist_ok=True) |
| plt.savefig(os.path.join(save_dir, prefix + str(frame_idx) + ".png"), dpi=60) |
| plt.close() |
|
|
