npy.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from matplotlib.axes import Axes\n",
    "import matplotlib.gridspec as gridspec\n",
    "import numpy as np\n",
    "from typing import cast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = np.load(\"front.npz\")\n",
    "from numpy.typing import NDArray\n",
    "kps = d[\"reconstruction\"]\n",
    "kps.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```json\n",
    "\"keypoints\": { 0: \"nose\", 1: \"left_eye\", 2: \"right_eye\", 3: \"left_ear\", 4: \"right_ear\", 5: \"left_shoulder\", 6: \"right_shoulder\", 7: \"left_elbow\", 8: \"right_elbow\", 9: \"left_wrist\", 10: \"right_wrist\", 11: \"left_hip\", 12: \"right_hip\", 13: \"left_knee\", 14: \"right_knee\", 15: \"left_ankle\", 16: \"right_ankle\" }\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import plotly.express as px\n",
    "import plotly.graph_objects as go\n",
    "from typeguard import typechecked\n",
    "from jaxtyping import Int, Float, Bool, Num, jaxtyped\n",
    "from pydantic import BaseModel\n",
    "from typing import Optional, Union\n",
    "from plotly.graph_objects import layout\n",
    "from ipywidgets import interact, interactive, fixed, interact_manual, interactive_output\n",
    "import ipywidgets as widgets\n",
    "from IPython.display import clear_output\n",
    "\n",
    "\n",
    "NDArray = np.ndarray\n",
    "number = Union[int, float]\n",
    "Color = tuple[int, int, int] | str\n",
    "Step = layout.slider.Step\n",
    "Slider = layout.Slider\n",
    "\n",
    "\n",
    "class Joint(BaseModel):\n",
    "    index: int\n",
    "    opposite_index: Optional[int] = None\n",
    "    name: str\n",
    "    color: Color\n",
    "\n",
    "\n",
    "class Bone(BaseModel):\n",
    "    joint1: Joint\n",
    "    joint2: Joint\n",
    "    name: str\n",
    "    color: Color\n",
    "\n",
    "    @staticmethod\n",
    "    def from_indexes(joints: list[Joint], idx_1: int, idx_2: int, name: str,\n",
    "                     color: Color) -> \"Bone\":\n",
    "        return Bone(joint1=joints[idx_1],\n",
    "                    joint2=joints[idx_2],\n",
    "                    name=name,\n",
    "                    color=color)\n",
    "\n",
    "\n",
    "# https://plotly.com/python-api-reference/generated/plotly.graph_objects.scatter3d.marker.html\n",
    "# plotly.graph_objects.scatter3d.Marker\n",
    "def to_rgb_str(color: tuple[int, int, int]) -> str:\n",
    "    return f\"rgb({color[0]},{color[1]},{color[2]})\"\n",
    "\n",
    "\n",
    "COLOR_SPINE = to_rgb_str((138, 201, 38))  # green, spine & head\n",
    "COLOR_ARMS = to_rgb_str((255, 202, 58))  # yellow, arms & shoulders\n",
    "COLOR_LEGS = to_rgb_str((25, 130, 196))  # blue, legs & hips\n",
    "CIRCLE_SIZE = 4\n",
    "LINE_WIDTH = 3\n",
    "\n",
    "# for coco\n",
    "coco_joints = [\n",
    "    Joint(index=0, name=\"nose\", color=COLOR_SPINE),\n",
    "    Joint(index=1, name=\"left_eye\", opposite_index=2, color=COLOR_SPINE),\n",
    "    Joint(index=2, name=\"right_eye\", opposite_index=1, color=COLOR_SPINE),\n",
    "    Joint(index=3, name=\"left_ear\", opposite_index=4, color=COLOR_SPINE),\n",
    "    Joint(index=4, name=\"right_ear\", opposite_index=3, color=COLOR_SPINE),\n",
    "    Joint(index=5, name=\"left_shoulder\", opposite_index=6, color=COLOR_ARMS),\n",
    "    Joint(index=6, name=\"right_shoulder\", opposite_index=5, color=COLOR_ARMS),\n",
    "    Joint(index=7, name=\"left_elbow\", opposite_index=8, color=COLOR_ARMS),\n",
    "    Joint(index=8, name=\"right_elbow\", opposite_index=7, color=COLOR_ARMS),\n",
    "    Joint(index=9, name=\"left_wrist\", opposite_index=10, color=COLOR_ARMS),\n",
    "    Joint(index=10, name=\"right_wrist\", opposite_index=9, color=COLOR_ARMS),\n",
    "    Joint(index=11, name=\"left_hip\", opposite_index=12, color=COLOR_LEGS),\n",
    "    Joint(index=12, name=\"right_hip\", opposite_index=11, color=COLOR_LEGS),\n",
    "    Joint(index=13, name=\"left_knee\", opposite_index=14, color=COLOR_LEGS),\n",
    "    Joint(index=14, name=\"right_knee\", opposite_index=13, color=COLOR_LEGS),\n",
    "    Joint(index=15, name=\"left_ankle\", opposite_index=16, color=COLOR_LEGS),\n",
    "    Joint(index=16, name=\"right_ankle\", opposite_index=15, color=COLOR_LEGS),\n",
    "]\n",
    "\n",
    "# https://github.com/lllyasviel/ControlNet/discussions/266\n",
    "# for human 3.6\n",
    "human_36_joints = [\n",
    "    Joint(index=0, name=\"bottom_torso\", color=COLOR_SPINE),\n",
    "    Joint(index=1, name=\"left_hip\", opposite_index=4, color=COLOR_LEGS),\n",
    "    Joint(index=2, name=\"left_knee\", opposite_index=5, color=COLOR_LEGS),\n",
    "    Joint(index=3, name=\"left_foot\", opposite_index=6, color=COLOR_LEGS),\n",
    "    Joint(index=4, name=\"right_hip\", opposite_index=1, color=COLOR_LEGS),\n",
    "    Joint(index=5, name=\"right_knee\", opposite_index=2, color=COLOR_LEGS),\n",
    "    Joint(index=6, name=\"right_foot\", opposite_index=3, color=COLOR_LEGS),\n",
    "    Joint(index=7, name=\"center_torso\", color=COLOR_SPINE),\n",
    "    Joint(index=8, name=\"upper_torso\", color=COLOR_SPINE),\n",
    "    Joint(index=9, name=\"neck_base\", color=COLOR_SPINE),\n",
    "    Joint(index=10, name=\"center_head\", color=COLOR_SPINE),\n",
    "    Joint(index=11, name=\"right_shoulder\", opposite_index=14, color=COLOR_ARMS),\n",
    "    Joint(index=12, name=\"right_elbow\", opposite_index=15, color=COLOR_ARMS),\n",
    "    Joint(index=13, name=\"right_hand\", opposite_index=16, color=COLOR_ARMS),\n",
    "    Joint(index=14, name=\"left_shoulder\", opposite_index=11, color=COLOR_ARMS),\n",
    "    Joint(index=15, name=\"left_elbow\", opposite_index=12, color=COLOR_ARMS),\n",
    "    Joint(index=16, name=\"left_hand\", opposite_index=13, color=COLOR_ARMS),\n",
    "]\n",
    "\n",
    "human_36_bones = [\n",
    "    Bone.from_indexes(human_36_joints, 0, 1, \"left_hip\", COLOR_LEGS),\n",
    "    Bone.from_indexes(human_36_joints, 1, 2, \"left_knee\", COLOR_LEGS),\n",
    "    Bone.from_indexes(human_36_joints, 2, 3, \"left_foot\", COLOR_LEGS),\n",
    "    Bone.from_indexes(human_36_joints, 0, 4, \"right_hip\", COLOR_LEGS),\n",
    "    Bone.from_indexes(human_36_joints, 4, 5, \"right_knee\", COLOR_LEGS),\n",
    "    Bone.from_indexes(human_36_joints, 5, 6, \"right_foot\", COLOR_LEGS),\n",
    "    Bone.from_indexes(human_36_joints, 0, 7, \"center_torso\", COLOR_SPINE),\n",
    "    Bone.from_indexes(human_36_joints, 7, 8, \"upper_torso\", COLOR_SPINE),\n",
    "    Bone.from_indexes(human_36_joints, 8, 9, \"neck_base\", COLOR_SPINE),\n",
    "    Bone.from_indexes(human_36_joints, 9, 10, \"center_head\", COLOR_SPINE),\n",
    "    Bone.from_indexes(human_36_joints, 7, 11, \"right_shoulder\", COLOR_ARMS),\n",
    "    Bone.from_indexes(human_36_joints, 11, 12, \"right_elbow\", COLOR_ARMS),\n",
    "    Bone.from_indexes(human_36_joints, 12, 13, \"right_hand\", COLOR_ARMS),\n",
    "    Bone.from_indexes(human_36_joints, 7, 14, \"left_shoulder\", COLOR_ARMS),\n",
    "    Bone.from_indexes(human_36_joints, 14, 15, \"left_elbow\", COLOR_ARMS),\n",
    "    Bone.from_indexes(human_36_joints, 15, 16, \"left_hand\", COLOR_ARMS),\n",
    "]\n",
    "\n",
    "# note that the keypoints are in the format of (x, y, z)\n",
    "# where x is the horizontal axis, z is the depth axis, and y is the vertical axis (upwards)\n",
    "# i.e. the x-z plane is the ground plane\n",
    "\n",
    "\n",
    "@jaxtyped(typechecker=typechecked)\n",
    "def xz_ground_2_xy_ground(xyz: Num[NDArray, \"17 3\"]) -> Num[NDArray, \"17 3\"]:\n",
    "    \"\"\"\n",
    "    essentially make xyz -> xzy\n",
    "    \"\"\"\n",
    "    return np.array([xyz[:, 0], xyz[:, 2], xyz[:, 1]]).T\n",
    "\n",
    "\n",
    "def swap_axes_direction(arr: Num[NDArray, \"17 1\"]) -> Num[NDArray, \"17 1\"]:\n",
    "    return -arr\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, DisplayHandle\n",
    "\n",
    "hdl:Optional[DisplayHandle] = None\n",
    "is_first = True\n",
    "\n",
    "# https://stackoverflow.com/questions/63716543/plotly-how-to-update-redraw-a-plotly-express-figure-with-new-data\n",
    "fw = go.FigureWidget()\n",
    "def plot_frame(kps: Num[NDArray, \"N 17 3\"], index: int):\n",
    "    global fw\n",
    "    global is_first\n",
    "    assert 0 <= index < kps.shape[0]\n",
    "    fig = go.Figure()\n",
    "    sel = kps[index]\n",
    "    sel = xz_ground_2_xy_ground(sel)\n",
    "    # reverse the upright axis\n",
    "    sel[:, 2] = -sel[:, 2]\n",
    "\n",
    "    scatters = [\n",
    "        go.Scatter3d(x=[sel[j.index, 0]],\n",
    "                     y=[sel[j.index, 1]],\n",
    "                     z=[sel[j.index, 2]],\n",
    "                     mode='markers',\n",
    "                     marker=dict(size=CIRCLE_SIZE, color=j.color),\n",
    "                     name=j.name) for j in human_36_joints\n",
    "    ]\n",
    "    lines = [\n",
    "        go.Scatter3d(x=[sel[b.joint1.index, 0], sel[b.joint2.index, 0]],\n",
    "                     y=[sel[b.joint1.index, 1], sel[b.joint2.index, 1]],\n",
    "                     z=[sel[b.joint1.index, 2], sel[b.joint2.index, 2]],\n",
    "                     mode='lines',\n",
    "                     line=dict(color=b.color, width=LINE_WIDTH),\n",
    "                     name=b.name) for b in human_36_bones\n",
    "    ]\n",
    "    fig.add_traces(scatters + lines)\n",
    "    fw = go.FigureWidget(fig)\n",
    "    if is_first:\n",
    "        is_first = False\n",
    "    else:\n",
    "        assert hdl is not None, \"hdl must be initialized\"\n",
    "        fw = go.FigureWidget(fig)\n",
    "        fw.update_layout(height=600, scene=dict(aspectmode=\"data\"))\n",
    "        fw.show()\n",
    "        hdl.update(None)\n",
    "\n",
    "\n",
    "slider = widgets.IntSlider(min=0,\n",
    "                           max=kps.shape[0] - 1,\n",
    "                           step=1,\n",
    "                           value=0,\n",
    "                           continue_update=False)\n",
    "\n",
    "p = interactive(plot_frame, kps=fixed(kps), index=slider)\n",
    "display(p, clear=True)\n",
    "hdl = display(fw, display_id=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}