npy_bvh.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from npybvh.bvh import Bvh\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",
    "from ipywidgets import interact, interactive, fixed, interact_manual, interactive_output\n",
    "import ipywidgets as widgets\n",
    "from IPython.display import clear_output\n",
    "from dataclasses import dataclass\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anim = Bvh()\n",
    "anim.parse_file(\"172_jump_4.bvh\")\n",
    "joint = anim.joints\n",
    "for k, v in joint.items():\n",
    "    print(k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "class Joint(BaseModel):\n",
    "    index: int\n",
    "    opposite_index: Optional[int] = None\n",
    "    name: str\n",
    "    color: Color\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",
    "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",
    "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",
    "COLOR_FINGERS = to_rgb_str((255, 0, 0))  # red, fingers\n",
    "COLOR_HANDS = COLOR_FINGERS\n",
    "CIRCLE_SIZE = 2\n",
    "LINE_WIDTH = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "joints_map_list = [\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=\"HeadTip_end\", color=COLOR_SPINE),\n",
    "    Joint(index=9, name=\"RightShoulder\", opposite_index=42, color=COLOR_ARMS),\n",
    "    Joint(index=10, name=\"RightArm\", opposite_index=43, color=COLOR_ARMS),\n",
    "    Joint(index=11, name=\"RightForeArm\", opposite_index=44, color=COLOR_ARMS),\n",
    "    Joint(index=12, name=\"RightHand\", opposite_index=45, color=COLOR_ARMS),\n",
    "    Joint(index=13, name=\"RightFinger1Metacarpal\", opposite_index=71, color=COLOR_HANDS),\n",
    "    Joint(index=14, name=\"RightFinger1Proximal\", opposite_index=72, color=COLOR_HANDS),\n",
    "    Joint(index=15, name=\"RightFinger1Distal\", opposite_index=73, color=COLOR_HANDS),\n",
    "    Joint(index=16, name=\"RightFinger1Tip\", opposite_index=74, color=COLOR_HANDS),\n",
    "    Joint(index=17, name=\"RightFinger1Tip_end\", opposite_index=75, color=COLOR_HANDS),\n",
    "    Joint(index=18, name=\"RightFinger5Metacarpal\", opposite_index=46, color=COLOR_HANDS),\n",
    "    Joint(index=19, name=\"RightFinger5Proximal\", opposite_index=47, color=COLOR_HANDS),\n",
    "    Joint(index=20, name=\"RightFinger5Medial\", opposite_index=48, color=COLOR_HANDS),\n",
    "    Joint(index=21, name=\"RightFinger5Distal\", opposite_index=49, color=COLOR_HANDS),\n",
    "    Joint(index=22, name=\"RightFinger5Tip\", opposite_index=50, color=COLOR_HANDS),\n",
    "    Joint(index=23, name=\"RightFinger5Tip_end\", opposite_index=51, color=COLOR_HANDS),\n",
    "    Joint(index=24, name=\"RightFinger4Metacarpal\", opposite_index=52, color=COLOR_HANDS),\n",
    "    Joint(index=25, name=\"RightFinger4Proximal\", opposite_index=53, color=COLOR_HANDS),\n",
    "    Joint(index=26, name=\"RightFinger4Medial\", opposite_index=54, color=COLOR_HANDS),\n",
    "    Joint(index=27, name=\"RightFinger4Distal\", opposite_index=55, color=COLOR_HANDS),\n",
    "    Joint(index=28, name=\"RightFinger4Tip\", opposite_index=56, color=COLOR_HANDS),\n",
    "    Joint(index=29, name=\"RightFinger4Tip_end\", opposite_index=57, color=COLOR_HANDS),\n",
    "    Joint(index=30, name=\"RightFinger3Metacarpal\", opposite_index=58, color=COLOR_HANDS),\n",
    "    Joint(index=31, name=\"RightFinger3Proximal\", opposite_index=59, color=COLOR_HANDS),\n",
    "    Joint(index=32, name=\"RightFinger3Medial\", opposite_index=60, color=COLOR_HANDS),\n",
    "    Joint(index=33, name=\"RightFinger3Distal\", opposite_index=61, color=COLOR_HANDS),\n",
    "    Joint(index=34, name=\"RightFinger3Tip\", opposite_index=62, color=COLOR_HANDS),\n",
    "    Joint(index=35, name=\"RightFinger3Tip_end\", opposite_index=63, color=COLOR_HANDS),\n",
    "    Joint(index=36, name=\"RightFinger2Metacarpal\", opposite_index=64, color=COLOR_HANDS),\n",
    "    Joint(index=37, name=\"RightFinger2Proximal\", opposite_index=65, color=COLOR_HANDS),\n",
    "    Joint(index=38, name=\"RightFinger2Medial\", opposite_index=66, color=COLOR_HANDS),\n",
    "    Joint(index=39, name=\"RightFinger2Distal\", opposite_index=67, color=COLOR_HANDS),\n",
    "    Joint(index=40, name=\"RightFinger2Tip\", opposite_index=68, color=COLOR_HANDS),\n",
    "    Joint(index=41, name=\"RightFinger2Tip_end\", opposite_index=69, color=COLOR_HANDS),\n",
    "    Joint(index=42, name=\"LeftShoulder\", opposite_index=9, color=COLOR_ARMS),\n",
    "    Joint(index=43, name=\"LeftArm\", opposite_index=10, color=COLOR_ARMS),\n",
    "    Joint(index=44, name=\"LeftForeArm\", opposite_index=11, color=COLOR_ARMS),\n",
    "    Joint(index=45, name=\"LeftHand\", opposite_index=12, color=COLOR_ARMS),\n",
    "    Joint(index=46, name=\"LeftFinger5Metacarpal\", opposite_index=18, color=COLOR_HANDS),\n",
    "    Joint(index=47, name=\"LeftFinger5Proximal\", opposite_index=19, color=COLOR_HANDS),\n",
    "    Joint(index=48, name=\"LeftFinger5Medial\", opposite_index=20, color=COLOR_HANDS),\n",
    "    Joint(index=49, name=\"LeftFinger5Distal\", opposite_index=21, color=COLOR_HANDS),\n",
    "    Joint(index=50, name=\"LeftFinger5Tip\", opposite_index=22, color=COLOR_HANDS),\n",
    "    Joint(index=51, name=\"LeftFinger5Tip_end\", opposite_index=23, color=COLOR_HANDS),\n",
    "    Joint(index=52, name=\"LeftFinger4Metacarpal\", opposite_index=24, color=COLOR_HANDS),\n",
    "    Joint(index=53, name=\"LeftFinger4Proximal\", opposite_index=25, color=COLOR_HANDS),\n",
    "    Joint(index=54, name=\"LeftFinger4Medial\", opposite_index=26, color=COLOR_HANDS),\n",
    "    Joint(index=55, name=\"LeftFinger4Distal\", opposite_index=27, color=COLOR_HANDS),\n",
    "    Joint(index=56, name=\"LeftFinger4Tip\", opposite_index=28, color=COLOR_HANDS),\n",
    "    Joint(index=57, name=\"LeftFinger4Tip_end\", opposite_index=29, color=COLOR_HANDS),\n",
    "    Joint(index=58, name=\"LeftFinger3Metacarpal\", opposite_index=30, color=COLOR_HANDS),\n",
    "    Joint(index=59, name=\"LeftFinger3Proximal\", opposite_index=31, color=COLOR_HANDS),\n",
    "    Joint(index=60, name=\"LeftFinger3Medial\", opposite_index=32, color=COLOR_HANDS),\n",
    "    Joint(index=61, name=\"LeftFinger3Distal\", opposite_index=33, color=COLOR_HANDS),\n",
    "    Joint(index=62, name=\"LeftFinger3Tip\", opposite_index=34, color=COLOR_HANDS),\n",
    "    Joint(index=63, name=\"LeftFinger3Tip_end\", opposite_index=35, color=COLOR_HANDS),\n",
    "    Joint(index=64, name=\"LeftFinger2Metacarpal\", opposite_index=36, color=COLOR_HANDS),\n",
    "    Joint(index=65, name=\"LeftFinger2Proximal\", opposite_index=37, color=COLOR_HANDS),\n",
    "    Joint(index=66, name=\"LeftFinger2Medial\", opposite_index=38, color=COLOR_HANDS),\n",
    "    Joint(index=67, name=\"LeftFinger2Distal\", opposite_index=39, color=COLOR_HANDS),\n",
    "    Joint(index=68, name=\"LeftFinger2Tip\", opposite_index=40, color=COLOR_HANDS),\n",
    "    Joint(index=69, name=\"LeftFinger2Tip_end\", opposite_index=41, color=COLOR_HANDS),\n",
    "    Joint(index=70, name=\"LeftFinger1Metacarpal\", opposite_index=13, color=COLOR_HANDS),\n",
    "    Joint(index=71, name=\"LeftFinger1Proximal\", opposite_index=14, color=COLOR_HANDS),\n",
    "    Joint(index=72, name=\"LeftFinger1Distal\", opposite_index=15, color=COLOR_HANDS),\n",
    "    Joint(index=73, name=\"LeftFinger1Tip\", opposite_index=16, color=COLOR_HANDS),\n",
    "    Joint(index=74, name=\"LeftFinger1Tip_end\", opposite_index=17, color=COLOR_HANDS),\n",
    "    Joint(index=75, name=\"RightThigh\", opposite_index=81, color=COLOR_LEGS),\n",
    "    Joint(index=76, name=\"RightShin\", opposite_index=82, color=COLOR_LEGS),\n",
    "    Joint(index=77, name=\"RightFoot\", opposite_index=83, color=COLOR_LEGS),\n",
    "    Joint(index=78, name=\"RightToe\", opposite_index=84, color=COLOR_LEGS),\n",
    "    Joint(index=79, name=\"RightToeTip\", opposite_index=85, color=COLOR_LEGS),\n",
    "    Joint(index=80, name=\"RightToeTip_end\", opposite_index=86, color=COLOR_LEGS),\n",
    "    Joint(index=81, name=\"LeftThigh\", opposite_index=75, color=COLOR_LEGS),\n",
    "    Joint(index=82, name=\"LeftShin\", opposite_index=76, color=COLOR_LEGS),\n",
    "    Joint(index=83, name=\"LeftFoot\", opposite_index=77, color=COLOR_LEGS),\n",
    "    Joint(index=84, name=\"LeftToe\", opposite_index=78, color=COLOR_LEGS),\n",
    "    Joint(index=85, name=\"LeftToeTip\", opposite_index=79, color=COLOR_LEGS),\n",
    "    Joint(index=86, name=\"LeftToeTip_end\", opposite_index=80, color=COLOR_LEGS),\n",
    "]\n",
    "\n",
    "bone_map_lists = [\n",
    "    # Spine and Head\n",
    "    Bone.from_indexes(joints_map_list, 0, 1, \"spine1\", COLOR_SPINE),\n",
    "    Bone.from_indexes(joints_map_list, 1, 2, \"spine2\", COLOR_SPINE),\n",
    "    Bone.from_indexes(joints_map_list, 2, 3, \"spine3\", COLOR_SPINE),\n",
    "    Bone.from_indexes(joints_map_list, 3, 4, \"spine4\", COLOR_SPINE),\n",
    "    Bone.from_indexes(joints_map_list, 4, 5, \"neck\", COLOR_SPINE),\n",
    "    Bone.from_indexes(joints_map_list, 5, 6, \"head\", COLOR_SPINE),\n",
    "    Bone.from_indexes(joints_map_list, 6, 7, \"head_tip\", COLOR_SPINE),\n",
    "    Bone.from_indexes(joints_map_list, 7, 8, \"head_tip_end\", COLOR_SPINE),\n",
    "\n",
    "    # Right Arm and Hand\n",
    "    Bone.from_indexes(joints_map_list, 4, 9, \"right_shoulder\", COLOR_ARMS),\n",
    "    Bone.from_indexes(joints_map_list, 9, 10, \"right_upper_arm\", COLOR_ARMS),\n",
    "    Bone.from_indexes(joints_map_list, 10, 11, \"right_forearm\", COLOR_ARMS),\n",
    "    Bone.from_indexes(joints_map_list, 11, 12, \"right_hand\", COLOR_ARMS),\n",
    "\n",
    "    # Right Fingers\n",
    "    Bone.from_indexes(joints_map_list, 12, 13, \"right_thumb_metacarpal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 13, 14, \"right_thumb_proximal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 14, 15, \"right_thumb_distal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 15, 16, \"right_thumb_tip\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 16, 17, \"right_thumb_tip_end\", COLOR_HANDS),\n",
    "\n",
    "    Bone.from_indexes(joints_map_list, 12, 18, \"right_pinky_metacarpal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 18, 19, \"right_pinky_proximal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 19, 20, \"right_pinky_medial\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 20, 21, \"right_pinky_distal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 21, 22, \"right_pinky_tip\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 22, 23, \"right_pinky_tip_end\", COLOR_HANDS),\n",
    "\n",
    "    Bone.from_indexes(joints_map_list, 12, 24, \"right_ring_metacarpal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 24, 25, \"right_ring_proximal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 25, 26, \"right_ring_medial\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 26, 27, \"right_ring_distal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 27, 28, \"right_ring_tip\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 28, 29, \"right_ring_tip_end\", COLOR_HANDS),\n",
    "\n",
    "    Bone.from_indexes(joints_map_list, 12, 30, \"right_middle_metacarpal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 30, 31, \"right_middle_proximal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 31, 32, \"right_middle_medial\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 32, 33, \"right_middle_distal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 33, 34, \"right_middle_tip\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 34, 35, \"right_middle_tip_end\", COLOR_HANDS),\n",
    "\n",
    "    Bone.from_indexes(joints_map_list, 12, 36, \"right_index_metacarpal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 36, 37, \"right_index_proximal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 37, 38, \"right_index_medial\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 38, 39, \"right_index_distal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 39, 40, \"right_index_tip\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 40, 41, \"right_index_tip_end\", COLOR_HANDS),\n",
    "\n",
    "    # Left Arm and Hand\n",
    "    Bone.from_indexes(joints_map_list, 4, 42, \"left_shoulder\", COLOR_ARMS),\n",
    "    Bone.from_indexes(joints_map_list, 42, 43, \"left_upper_arm\", COLOR_ARMS),\n",
    "    Bone.from_indexes(joints_map_list, 43, 44, \"left_forearm\", COLOR_ARMS),\n",
    "    Bone.from_indexes(joints_map_list, 44, 45, \"left_hand\", COLOR_ARMS),\n",
    "\n",
    "    # Left Fingers\n",
    "    Bone.from_indexes(joints_map_list, 45, 46, \"left_pinky_metacarpal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 46, 47, \"left_pinky_proximal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 47, 48, \"left_pinky_medial\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 48, 49, \"left_pinky_distal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 49, 50, \"left_pinky_tip\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 50, 51, \"left_pinky_tip_end\", COLOR_HANDS),\n",
    "\n",
    "    Bone.from_indexes(joints_map_list, 45, 52, \"left_ring_metacarpal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 52, 53, \"left_ring_proximal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 53, 54, \"left_ring_medial\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 54, 55, \"left_ring_distal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 55, 56, \"left_ring_tip\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 56, 57, \"left_ring_tip_end\", COLOR_HANDS),\n",
    "\n",
    "    Bone.from_indexes(joints_map_list, 45, 58, \"left_middle_metacarpal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 58, 59, \"left_middle_proximal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 59, 60, \"left_middle_medial\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 60, 61, \"left_middle_distal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 61, 62, \"left_middle_tip\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 62, 63, \"left_middle_tip_end\", COLOR_HANDS),\n",
    "\n",
    "    Bone.from_indexes(joints_map_list, 45, 64, \"left_index_metacarpal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 64, 65, \"left_index_proximal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 65, 66, \"left_index_medial\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 66, 67, \"left_index_distal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 67, 68, \"left_index_tip\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 68, 69, \"left_index_tip_end\", COLOR_HANDS),\n",
    "\n",
    "    Bone.from_indexes(joints_map_list, 45, 70, \"left_thumb_metacarpal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 70, 71, \"left_thumb_proximal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 71, 72, \"left_thumb_distal\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 72, 73, \"left_thumb_tip\", COLOR_HANDS),\n",
    "    Bone.from_indexes(joints_map_list, 73, 74, \"left_thumb_tip_end\", COLOR_HANDS),\n",
    "\n",
    "    # Right Leg\n",
    "    Bone.from_indexes(joints_map_list, 0, 75, \"right_thigh\", COLOR_LEGS),\n",
    "    Bone.from_indexes(joints_map_list, 75, 76, \"right_shin\", COLOR_LEGS),\n",
    "    Bone.from_indexes(joints_map_list, 76, 77, \"right_foot\", COLOR_LEGS),\n",
    "    Bone.from_indexes(joints_map_list, 77, 78, \"right_toe\", COLOR_LEGS),\n",
    "    Bone.from_indexes(joints_map_list, 78, 79, \"right_toe_tip\", COLOR_LEGS),\n",
    "    Bone.from_indexes(joints_map_list, 79, 80, \"right_toe_tip_end\", COLOR_LEGS),\n",
    "\n",
    "    # Left Leg\n",
    "    Bone.from_indexes(joints_map_list, 0, 81, \"left_thigh\", COLOR_LEGS),\n",
    "    Bone.from_indexes(joints_map_list, 81, 82, \"left_shin\", COLOR_LEGS),\n",
    "    Bone.from_indexes(joints_map_list, 82, 83, \"left_foot\", COLOR_LEGS),\n",
    "    Bone.from_indexes(joints_map_list, 83, 84, \"left_toe\", COLOR_LEGS),\n",
    "    Bone.from_indexes(joints_map_list, 84, 85, \"left_toe_tip\", COLOR_LEGS),\n",
    "    Bone.from_indexes(joints_map_list, 85, 86, \"left_toe_tip_end\", COLOR_LEGS),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Callable\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class PayloadJoint:\n",
    "    joint: Joint\n",
    "    payload: Num[NDArray, \"F 1 3\"]\n",
    "\n",
    "    @staticmethod\n",
    "    def from_stacked(joint: Joint, payload: Num[NDArray,\n",
    "                                                \"F J 3\"]) -> \"PayloadJoint\":\n",
    "        assert payload.shape[2] == 3, \"must be 3D coordinates\"\n",
    "        assert len(payload.shape) == 3, \"must be (frames, joints, coordinates)\"\n",
    "        idx = joint.index\n",
    "        p = payload[:, idx].reshape(-1, 1, 3)\n",
    "        return PayloadJoint(joint=joint, payload=p)\n",
    "\n",
    "    def create_scatter(self, frame: int) -> go.Scatter3d:\n",
    "        total = self.payload.shape[0]\n",
    "        assert frame < total, f\"frame {frame} out of range {total}\"\n",
    "        return go.Scatter3d(x=[self.payload[frame, 0, 0]],\n",
    "                            y=[self.payload[frame, 0, 1]],\n",
    "                            z=[self.payload[frame, 0, 2]],\n",
    "                            mode=\"markers\",\n",
    "                            marker=dict(size=CIRCLE_SIZE,\n",
    "                                        color=self.joint.color),\n",
    "                            name=self.joint.name)\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class PayloadBone:\n",
    "    bone: Bone\n",
    "    payload: Num[NDArray, \"F 1 2 3\"]\n",
    "\n",
    "    @staticmethod\n",
    "    def from_stacked(bone: Bone, payload: Num[NDArray,\n",
    "                                              \"F J 3\"]) -> \"PayloadBone\":\n",
    "        assert payload.shape[2] == 3, \"must be 3D coordinates\"\n",
    "        assert len(payload.shape) == 3, \"must be (frames, joints, coordinates)\"\n",
    "        idx1 = bone.joint1.index\n",
    "        idx2 = bone.joint2.index\n",
    "        payload_1 = payload[:, idx1].reshape(-1, 1, 1, 3)\n",
    "        payload_2 = payload[:, idx2].reshape(-1, 1, 1, 3)\n",
    "        return PayloadBone(bone=bone,\n",
    "                           payload=np.concatenate([payload_1, payload_2],\n",
    "                                                  axis=2))\n",
    "\n",
    "    def create_lines(self, frame: int) -> go.Scatter3d:\n",
    "        total = self.payload.shape[0]\n",
    "        assert frame < total, f\"frame {frame} out of range {total}\"\n",
    "        return go.Scatter3d(x=self.payload[frame, 0, :, 0].T,\n",
    "                            y=self.payload[frame, 0, :, 1].T,\n",
    "                            z=self.payload[frame, 0, :, 2].T,\n",
    "                            mode=\"lines\",\n",
    "                            line=dict(width=LINE_WIDTH, color=self.bone.color),\n",
    "                            name=self.bone.name)\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Skeleton:\n",
    "    bone_schema: list[Bone]\n",
    "    joint_schema: list[Joint]\n",
    "    joints: list[PayloadJoint]\n",
    "    bones: list[PayloadBone]\n",
    "\n",
    "    @staticmethod\n",
    "    def from_stacked(joints: list[Joint], bones: list[Bone],\n",
    "                     payload: Num[NDArray, \"F J 3\"]) -> \"Skeleton\":\n",
    "        return Skeleton(\n",
    "            bone_schema=bones,\n",
    "            joint_schema=joints,\n",
    "            joints=[\n",
    "                PayloadJoint.from_stacked(joint, payload) for joint in joints\n",
    "            ],\n",
    "            bones=[PayloadBone.from_stacked(bone, payload) for bone in bones])\n",
    "\n",
    "    @jaxtyped(typechecker=typechecked)\n",
    "    def to_stacked_joints(self) -> Num[NDArray, \"F J 3\"]:\n",
    "        return np.stack([joint.payload for joint in self.joints], axis=1)\n",
    "\n",
    "    @jaxtyped(typechecker=typechecked)\n",
    "    def to_stacked_bones(self) -> Num[NDArray, \"F J 2 3\"]:\n",
    "        return np.stack([bone.payload for bone in self.bones], axis=1)\n",
    "\n",
    "    @property\n",
    "    def total_frames(self) -> int:\n",
    "        return self.joints[0].payload.shape[0]\n",
    "\n",
    "    def filter(self, predicate: Callable[[str], bool]) -> \"Skeleton\":\n",
    "        \"\"\"\n",
    "        filter joints and bones by name\n",
    "        \"\"\"\n",
    "        joints = [joint for joint in self.joints if predicate(joint.joint.name)]\n",
    "        bones = [bone for bone in self.bones if predicate(bone.bone.name)]\n",
    "        bs = [bone for bone in self.bone_schema if predicate(bone.name)]\n",
    "        js = [joint for joint in self.joint_schema if predicate(joint.name)]\n",
    "        return Skeleton(joints=joints,\n",
    "                        bones=bones,\n",
    "                        joint_schema=js,\n",
    "                        bone_schema=bs)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pos, rot = anim.all_frame_poses()\n",
    "display(pos.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jaxtyped(typechecker=typechecked)\n",
    "def preprocess_data(pos: Num[NDArray, \"F J 3\"]) -> Num[NDArray, \"F J 3\"]:\n",
    "    pre = pos.copy()\n",
    "    pre = np.concatenate([pre[:, :, 0:1], pre[:, :, 2:3], pre[:, :, 1:2]], axis=2)\n",
    "    return pre"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pre = preprocess_data(pos)\n",
    "sk = Skeleton.from_stacked(joints_map_list, bone_map_lists, pre)\n",
    "\n",
    "def filter_by_name(name: str) -> bool:\n",
    "    if \"finger\" in name.lower():\n",
    "        return False\n",
    "    if \"index\" in name.lower():\n",
    "        return False\n",
    "    if \"middle\" in name.lower():\n",
    "        return False\n",
    "    if \"ring\" in name.lower():\n",
    "        return False\n",
    "    if \"thumb\" in name.lower():\n",
    "        return False\n",
    "    if \"pinky\" in name.lower():\n",
    "        return False\n",
    "    if \"tip\" in name.lower():\n",
    "        return False\n",
    "    return True\n",
    "\n",
    "sk_f = sk.filter(filter_by_name)"
   ]
  },
  {
   "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",
    "fw = go.FigureWidget()\n",
    "def plot_frame(sk: Skeleton, index: int):\n",
    "    global fw\n",
    "    global is_first\n",
    "    frames = sk.total_frames\n",
    "    assert 0 <= index < frames, f\"index must be between 0 and {frames - 1} inclusive but got {index}; frames={frames}\"\n",
    "    fig = go.Figure()\n",
    "    scatters = [j.create_scatter(index) for j in sk.joints]\n",
    "    lines = [b.create_lines(index) for b in sk.bones]\n",
    "    fig.add_traces(scatters + lines)\n",
    "    # if there's a JavaScript error\n",
    "    # restart Visual Studio Code (or use `.show()` method?)\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",
    "        # https://stackoverflow.com/questions/52863305/plotly-scatter3d-how-can-i-force-3d-axes-to-have-the-same-scale-aspect-ratio\n",
    "        fw.update_layout(height=600, scene=dict(aspectmode=\"data\"))\n",
    "        fw.show()\n",
    "        hdl.update(None)\n",
    "\n",
    "slider = widgets.IntSlider(min=0,\n",
    "                           max=sk.total_frames - 1,\n",
    "                           step=1,\n",
    "                           value=0,\n",
    "                           continue_update=False)\n",
    "\n",
    "p = interactive(plot_frame, sk=fixed(sk_f), 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
}