play_bvhio.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import bvhio.bvhio as bvhio\n",
    "from bvhio.bvhio import BvhJoint, Joint, Pose\n",
    "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",
    "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\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",
    "number = Union[int, float]\n",
    "Color = tuple[int, int, int] | str\n",
    "Step = layout.slider.Step\n",
    "Slider = layout.Slider\n",
    "NDArray = np.ndarray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b = bvhio.readAsBvh(\"172_jump_4.bvh\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Name\\tIndex\\tDepth\")\n",
    "for joint, index, depth in b.Root.layout():\n",
    "    n = joint.Name\n",
    "    print(f\"{n}\\t{index}\\t{depth}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```txt\n",
    "Hips\n",
    "├── Spine1\n",
    "│   ├── Spine2\n",
    "│       ├── Spine3\n",
    "│           ├── Spine4\n",
    "│               ├── Neck\n",
    "│               │   ├── Head\n",
    "│               │       └── HeadTip\n",
    "│               ├── RightShoulder\n",
    "│               │   ├── RightArm\n",
    "│               │       ├── RightForeArm\n",
    "│               │           ├── RightHand\n",
    "│               │               ├── RightFinger1Metacarpal\n",
    "│               │               │   ├── RightFinger1Proximal\n",
    "│               │               │       ├── RightFinger1Distal\n",
    "│               │               │           └── RightFinger1Tip\n",
    "│               │               ├── RightFinger5Metacarpal\n",
    "│               │               │   ├── RightFinger5Proximal\n",
    "│               │               │       ├── RightFinger5Medial\n",
    "│               │               │           ├── RightFinger5Distal\n",
    "│               │               │               └── RightFinger5Tip\n",
    "│               │               ├── RightFinger4Metacarpal\n",
    "│               │               │   ├── RightFinger4Proximal\n",
    "│               │               │       ├── RightFinger4Medial\n",
    "│               │               │           ├── RightFinger4Distal\n",
    "│               │               │               └── RightFinger4Tip\n",
    "│               │               ├── RightFinger3Metacarpal\n",
    "│               │               │   ├── RightFinger3Proximal\n",
    "│               │               │       ├── RightFinger3Medial\n",
    "│               │               │           ├── RightFinger3Distal\n",
    "│               │               │               └── RightFinger3Tip\n",
    "│               │               ├── RightFinger2Metacarpal\n",
    "│               │                   ├── RightFinger2Proximal\n",
    "│               │                       ├── RightFinger2Medial\n",
    "│               │                           ├── RightFinger2Distal\n",
    "│               │                               └── RightFinger2Tip\n",
    "│               ├── LeftShoulder\n",
    "│                   ├── LeftArm\n",
    "│                       ├── LeftForeArm\n",
    "│                           ├── LeftHand\n",
    "│                               ├── LeftFinger5Metacarpal\n",
    "│                               │   ├── LeftFinger5Proximal\n",
    "│                               │       ├── LeftFinger5Medial\n",
    "│                               │           ├── LeftFinger5Distal\n",
    "│                               │               └── LeftFinger5Tip\n",
    "│                               ├── LeftFinger4Metacarpal\n",
    "│                               │   ├── LeftFinger4Proximal\n",
    "│                               │       ├── LeftFinger4Medial\n",
    "│                               │           ├── LeftFinger4Distal\n",
    "│                               │               └── LeftFinger4Tip\n",
    "│                               ├── LeftFinger3Metacarpal\n",
    "│                               │   ├── LeftFinger3Proximal\n",
    "│                               │       ├── LeftFinger3Medial\n",
    "│                               │           ├── LeftFinger3Distal\n",
    "│                               │               └── LeftFinger3Tip\n",
    "│                               ├── LeftFinger2Metacarpal\n",
    "│                               │   ├── LeftFinger2Proximal\n",
    "│                               │       ├── LeftFinger2Medial\n",
    "│                               │           ├── LeftFinger2Distal\n",
    "│                               │               └── LeftFinger2Tip\n",
    "│                               ├── LeftFinger1Metacarpal\n",
    "│                                   ├── LeftFinger1Proximal\n",
    "│                                       ├── LeftFinger1Distal\n",
    "│                                           └── LeftFinger1Tip\n",
    "├── RightThigh\n",
    "│   ├── RightShin\n",
    "│       ├── RightFoot\n",
    "│           ├── RightToe\n",
    "│               └── RightToeTip\n",
    "├── LeftThigh\n",
    "    ├── LeftShin\n",
    "        ├── LeftFoot\n",
    "            ├── LeftToe\n",
    "                └── LeftToeTip\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_DEPTH = 8\n",
    "joints: list[BvhJoint] = []\n",
    "for joint, index, depth in b.Root.layout():\n",
    "    if depth <= MAX_DEPTH:\n",
    "        joints.append(joint)\n",
    "joints = list(filter(lambda j: \"Toe\" not in j.Name, joints))\n",
    "for j in joints:\n",
    "    print(j.Name)\n",
    "print(len(joints))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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]})\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Hips\n",
    "- Spine1\n",
    "- Spine2\n",
    "- Spine3\n",
    "- Spine4\n",
    "- Neck\n",
    "- Head\n",
    "- HeadTip\n",
    "- RightShoulder\n",
    "- RightArm\n",
    "- RightForeArm\n",
    "- RightHand\n",
    "- LeftShoulder\n",
    "- LeftArm\n",
    "- LeftForeArm\n",
    "- LeftHand\n",
    "- RightThigh\n",
    "- RightShin\n",
    "- RightFoot\n",
    "- LeftThigh\n",
    "- LeftShin\n",
    "- LeftFoot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_stuff(joint:BvhJoint):\n",
    "    print(joint.Channels)\n",
    "    kf = joint.Keyframes[201]\n",
    "    print(joint.Offset)\n",
    "    print(joint.Name)\n",
    "    print(kf.Position) # if there's no position channel, use offset as position (weird)\n",
    "    print(kf.Up)\n",
    "    print(kf.Rotation)\n",
    "\n",
    "print_stuff(joints[0])\n",
    "print_stuff(joints[1])\n",
    "print_stuff(joints[2])\n",
    "print_stuff(joints[10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "s_bvh_joints = [\n",
    "    Joint(index=0, name=\"Hips\", color=COLOR_SPINE),\n",
    "    Joint(index=1, name=\"Spine1\", color=COLOR_SPINE),\n",
    "    Joint(index=2, name=\"Spine2\", color=COLOR_SPINE),\n",
    "    Joint(index=3, name=\"Spine3\", color=COLOR_SPINE),\n",
    "    Joint(index=4, name=\"Spine4\", color=COLOR_SPINE),\n",
    "    Joint(index=5, name=\"Neck\", color=COLOR_SPINE),\n",
    "    Joint(index=6, name=\"Head\", color=COLOR_SPINE),\n",
    "    Joint(index=7, name=\"HeadTip\", color=COLOR_SPINE),\n",
    "    Joint(index=8, name=\"RightShoulder\", color=COLOR_ARMS),\n",
    "    Joint(index=9, name=\"RightArm\", color=COLOR_ARMS),\n",
    "    Joint(index=10, name=\"RightForeArm\", color=COLOR_ARMS),\n",
    "    Joint(index=11, name=\"RightHand\", color=COLOR_ARMS),\n",
    "    Joint(index=12, name=\"LeftShoulder\", color=COLOR_ARMS),\n",
    "    Joint(index=13, name=\"LeftArm\", color=COLOR_ARMS),\n",
    "    Joint(index=14, name=\"LeftForeArm\", color=COLOR_ARMS),\n",
    "    Joint(index=15, name=\"LeftHand\", color=COLOR_ARMS),\n",
    "    Joint(index=16, name=\"RightThigh\", color=COLOR_LEGS),\n",
    "    Joint(index=17, name=\"RightShin\", color=COLOR_LEGS),\n",
    "    Joint(index=18, name=\"RightFoot\", color=COLOR_LEGS),\n",
    "    Joint(index=19, name=\"LeftThigh\", color=COLOR_LEGS),\n",
    "    Joint(index=20, name=\"LeftShin\", color=COLOR_LEGS),\n",
    "    Joint(index=21, name=\"LeftFoot\", color=COLOR_LEGS),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jaxtyped(typechecker=typechecked)\n",
    "def joint_to_stacked(joint: BvhJoint) -> Num[NDArray, \"N 3\"]:\n",
    "    \"\"\"\n",
    "    Convert a joint to a numpy matrix of shape (N, 3), where N is the number of keyframes.\n",
    "    \"\"\"\n",
    "    return np.stack([np.array(joint.Keyframes[i].Up) for i in range(len(joint.Keyframes))])\n",
    "\n",
    "@jaxtyped(typechecker=typechecked)\n",
    "def stack_many(joints: list[BvhJoint]) -> Num[NDArray, \"N J 3\"]:\n",
    "    \"\"\"\n",
    "    Stack the keyframes of many joints into a single numpy matrix, where N is the number of keyframes and J is the number of joints.\n",
    "\n",
    "    Note that the order of the joints is preserved in the second dimension of the output.\n",
    "    Every joint in the list should have the same number of keyframes.\n",
    "    \"\"\"\n",
    "    return np.stack([joint_to_stacked(j) for j in joints], axis=1)\n",
    "\n",
    "kps = stack_many(joints)\n",
    "display(kps.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fw = go.FigureWidget()\n",
    "\n",
    "\n",
    "def plot_frame(kps: Num[NDArray, \"N J 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",
    "    # 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 s_bvh_joints\n",
    "    ]\n",
    "    fig.add_traces(scatters)\n",
    "    fw = go.FigureWidget(fig)\n",
    "    # if there's a JavaScript error\n",
    "    # restart Visual Studio Code (or use `.show()` method?)\n",
    "    fw.update_layout(height=600)\n",
    "    display(fw)\n",
    "    return fig\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)"
   ]
  }
 ],
 "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
}