notebooks/SMPL_details.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SMPL-details\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Environment Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Packages you may use very often.\n",
    "import torch\n",
    "import numpy as np\n",
    "from smplx import SMPL\n",
    "from smplx.vertex_ids import vertex_ids as VERTEX_IDS\n",
    "from pytorch3d import transforms  # You may use this package when performing rotation representation transformation.\n",
    "\n",
    "# Things you don't need to care about. They are just for driving the tutorials.\n",
    "from lib.logger.look_tool import look_tensor\n",
    "from lib.utils.path_manager import PathManager\n",
    "from lib.viewer.wis3d_utils import HWis3D as Wis3D\n",
    "from lib.skeleton import Skeleton_SMPL24\n",
    "\n",
    "pm = PathManager()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING: You are using a SMPL model, with only 10 shape coefficients.\n"
     ]
    }
   ],
   "source": [
    "body_model_smpl = SMPL(\n",
    "        model_path = pm.inputs / 'body_models' / 'smpl',\n",
    "        gender     = 'neutral',\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Understand the LBS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Understand the Joint Regressor\n",
    "\n",
    "You can get J_regressor from `smplx.SMPL(...).J_regressor`. It's used to regress the joints from the vertices. To be specific, it's a matrix of size `(24, 6890)`, and the joint positions can be got from the linear combination of the relative vertices positions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([24, 6890])\n",
      "v shape = torch.Size([1, 6890, 3]), j shape = torch.Size([1, 24, 3])\n",
      "j_by_hand shape = torch.Size([1, 24, 3])\n",
      "shape = (1, 24, 3)\tdtype = torch.float32\tdevice = cpu\tmin/max/mean/std = [ -0.000000 -> 0.000000 ] ~ ( 0.000000, 0.000000 )\n"
     ]
    }
   ],
   "source": [
    "J_regressor = body_model_smpl.J_regressor  # (24, 6890)\n",
    "print(J_regressor.shape)\n",
    "\n",
    "# Get 'g.t.' of the SMPL's output.\n",
    "smpl_output = body_model_smpl()\n",
    "v = smpl_output.vertices.detach()\n",
    "j = smpl_output.joints[:, :24, :].detach()\n",
    "print(f'v shape = {v.shape}, j shape = {j.shape}')\n",
    "\n",
    "# Calculate the joints through J_regressor.\n",
    "j_by_hand = torch.matmul(J_regressor[None], v)\n",
    "print(f'j_by_hand shape = {j_by_hand.shape}')\n",
    "\n",
    "# Check the difference.\n",
    "delta = j - j_by_hand\n",
    "_ = look_tensor(delta)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, they have tiny difference (for the first 24 joints). But if you can get joints from the standard output object, you are still suggested to get the joints from that.  \n",
    "\n",
    "> Sometimes, we want the SMPL style 24 joints, but we only have the SMPL-X style parameters. As we have already shown, you can't use SMPL-X parameters to get SMPL's output directly. But you can first use a regressor to get the SMPL style 6890 vertices from SMPL-X style 10475 vertices, and then use the `J_regressor` to get the SMPL style 24 joints from the regressed 6890 vertices. We won't talk about these details here, but you can keep this in mind.\n",
    "\n",
    "Another thing you may need to know is that, the regressor is sparse, which means that most of the elements are zeros. Now let's get into the details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Among the total 165360, there are 165124 zeros, zero rate: 0.998573%.\n"
     ]
    }
   ],
   "source": [
    "# Statistics of the J_regressor.\n",
    "zero_cnt = (J_regressor == 0).sum().item()\n",
    "total_cnt = J_regressor.numel()\n",
    "print(f'Among the total {total_cnt}, there are {zero_cnt} zeros, zero rate: {zero_cnt / total_cnt:.6f}%.')\n",
    "\n",
    "# Visualize the joints.\n",
    "active_masks = (J_regressor != 0)  # (24, 6890)\n",
    "J_regressor_wis3d = Wis3D(\n",
    "        pm.outputs / 'wis3d',\n",
    "        'SMPL-J_regressor',\n",
    "    )\n",
    "# Add first reference skeleton and mesh.\n",
    "J_regressor_wis3d.add_motion_verts(verts=v.repeat(25, 1, 1), name=f'vertices', offset=0)\n",
    "J_regressor_wis3d.add_motion_skel(joints=j.repeat(25, 1, 1), bones=Skeleton_SMPL24.bones, colors=Skeleton_SMPL24.bone_colors, name=f'skeleton', offset=0)\n",
    "\n",
    "# Visualize each part of the J_regressor.\n",
    "for i in range(24):\n",
    "    mask = active_masks[i]\n",
    "    v_masked = v[0, mask]\n",
    "    # Visualize the things of interest.\n",
    "    J_regressor_wis3d.set_scene_id(i+1)\n",
    "    J_regressor_wis3d.add_point_cloud(vertices=v_masked, name=f'VOI-{i}')  # Vertices of interest used to regress i-th joint.\n",
    "    J_regressor_wis3d.add_spheres(centers=v_masked, radius=0.01, name=f'VOI-{i}')  # VOI used to regress i-th joint.\n",
    "    J_regressor_wis3d.add_spheres(centers=j[:1, i], radius=0.02, name=f'joint-{i}')  # i-th joint."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now visualize the active elements of the regressor to show exactly which vertices influence the joints regressed. You are supposed to interact with the Wis3D UI to check the things you are interested in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Serving on http://0.0.0.0:19090\n",
      "^C\n"
     ]
    }
   ],
   "source": [
    "# Start the server. (Remember to terminate the cell before going on.)\n",
    "!wis3d --vis_dir {pm.outputs / 'wis3d'} --host 0.0.0.0 --port 19090"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Understand the Joints 45\n",
    "\n",
    "You may notice that, when I mention \"joints\", I always refer to the first 24. But there are actually 45 joints returned by the SMPL model by default. So what are they?\n",
    "\n",
    "Actually, the 45 joints can be divided into two parts:\n",
    "\n",
    "1. The common SMPL joints, which are first regressed from shapped vertices, and then transformed by the pose parameters. They are the first 24 joints.\n",
    "2. Joints selected (by hand, in advance) from vertices. There are 21 of them by default. You can get the name and the vertices index of them from `smplx.vertex_ids.vertex_ids['smplh']`, I print them below.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Selected joints from vertices.\n",
      "---------------------------------\n",
      "|  joint_name \t|\tvid\t|\n",
      "---------------------------------\n",
      "|  nose  \t|\t332\t|\n",
      "|  reye  \t|\t6260\t|\n",
      "|  leye  \t|\t2800\t|\n",
      "|  rear  \t|\t4071\t|\n",
      "|  lear  \t|\t583\t|\n",
      "|  rthumb  \t|\t6191\t|\n",
      "|  rindex  \t|\t5782\t|\n",
      "|  rmiddle  \t|\t5905\t|\n",
      "|  rring  \t|\t6016\t|\n",
      "|  rpinky  \t|\t6133\t|\n",
      "|  lthumb  \t|\t2746\t|\n",
      "|  lindex  \t|\t2319\t|\n",
      "|  lmiddle  \t|\t2445\t|\n",
      "|  lring  \t|\t2556\t|\n",
      "|  lpinky  \t|\t2673\t|\n",
      "|  LBigToe  \t|\t3216\t|\n",
      "|  LSmallToe  \t|\t3226\t|\n",
      "|  LHeel  \t|\t3387\t|\n",
      "|  RBigToe  \t|\t6617\t|\n",
      "|  RSmallToe  \t|\t6624\t|\n",
      "|  RHeel  \t|\t6787\t|\n",
      "---------------------------------\n"
     ]
    }
   ],
   "source": [
    "print(f'Selected joints from vertices.')\n",
    "print(f'---------------------------------')\n",
    "print(f'|  joint_name \\t|\\tvid\\t|')\n",
    "print(f'---------------------------------')\n",
    "for k, v in VERTEX_IDS['smplh'].items():\n",
    "    print(f'|  {k}  \\t|\\t{v}\\t|')\n",
    "print(f'---------------------------------')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can check the implementation at `smplx.SMPL` for more details, SMPL provide some APIs to customize the outputs. They will be useful if you need to change the definition of the output joints (in order, or maybe select more joints from vertices)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Become Lighter\n",
    "\n",
    "The sparsity of the regressor tells that, if we only need the joint positions, we don't have to get all the vertices. Recall the inference process of SMPL, we found that we can ignore a lot of calculation while performing the Linear Blend Skinning, LBS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cv",
   "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",
   "version": "3.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}