notebooks/SMPLX_basic.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SMPLX-basic\n",
    "\n",
    "SMPLX extends SMPL with fully articulated hands and an expressive face.\n",
    "\n",
    "In this notebook, we will introduce the basic usage of `smplx.SMPLX` model, and some common problems you may encounter.\n",
    "\n",
    "And there is a \"SMPL family\", which is a set of models related to SMPL. You can have a overview of them in [SMPL wiki](https://meshcapade.wiki/SMPL).\n",
    "\n",
    "The basic usage of SMPLX is quite like SMPL, so we will be brief in the basic concept, and mainly focus on the differences between them. However, the **data of SMPL and SMPLX are quite different** (from the inputs to the outputs), you shouldn't mix them up.\n",
    "\n",
    "## Preparation\n",
    "\n",
    "You should also prepare the environments for SMPL, check the [SMPL-basic](./SMPL_basic.ipynb) for more details.\n",
    "\n",
    "We use three different checkpoints: `SMPLX_NEUTRAL.npz`, `SMPLX_MALE.npz` and `SMPLX_FEMALE.npz`. You can get them from [SMPL eXpressive](https://smpl-x.is.tue.mpg.de/).\n",
    "\n",
    "After downloading the SMPLX's checkpoints, you should put them to `data_inputs/body_models/smplx`, your directory tree should look like this:\n",
    "\n",
    "```\n",
    ".\n",
    "├── SMPLX_FEMALE.npz\n",
    "├── SMPLX_MALE.npz\n",
    "└── SMPLX_NEUTRAL.npz\n",
    "```\n",
    "\n",
    "## Tutorials"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Environment Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Packages you may use very often.\n",
    "import torch\n",
    "import numpy as np\n",
    "from smplx import SMPL, SMPLX\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.utils.path_manager import PathManager\n",
    "from lib.viewer.wis3d_utils import HWis3D as Wis3D\n",
    "from lib.skeleton import Skeleton_SMPL24, Skeleton_SMPL22\n",
    "\n",
    "pm = PathManager()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load SMPLX model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "B = 150\n",
    "\n",
    "body_models = {}\n",
    "genders = ['neutral', 'female', 'male']  # case insensitive\n",
    "\n",
    "for gender in genders:\n",
    "    body_models[gender] = SMPLX(\n",
    "            model_path = pm.inputs / 'body_models' / 'smplx',\n",
    "            gender     = gender,\n",
    "            batch_size = 150,\n",
    "        )\n",
    "\n",
    "# Prepare some parameters for later inference.\n",
    "body_model : SMPLX = body_models['neutral']  # use neutral for example\n",
    "\n",
    "# Prepare mesh template for later visualization.\n",
    "# Tips: mesh = vertices + faces, and the faces are the indices of vertices, which won't change across SMPL's outputs.\n",
    "mesh_temp : np.ndarray = body_model.faces  # (20908, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can find that, the shape of `faces` changes. Which means the **vertices of SMPLX model is different from SMPL model's**. We will dive into this later.\n",
    "\n",
    "Another thing you should notice is that, the implement of `smplx.SMPLX` does not support dynamic batch inference in through it's API. You need to define the batch size while defining the model. \n",
    "\n",
    "However, you still can make batch inference possible through modify the implements of the API, but those are not covered in this notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SMPL-X Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([150, 127, 3]) torch.Size([150, 10475, 3])\n"
     ]
    }
   ],
   "source": [
    "# Inference.\n",
    "smplx_out = body_model(\n",
    "        betas         = torch.zeros(B, 10),\n",
    "        global_orient = torch.zeros(B, 3),\n",
    "        body_pose     = torch.zeros(B, 63),\n",
    "        transl        = torch.zeros(B, 3),\n",
    "    )\n",
    "\n",
    "# Check output.\n",
    "joints : torch.Tensor = smplx_out.joints    # (B, 127, 3)\n",
    "verts  : torch.Tensor = smplx_out.vertices  # (B, 10475, 3)\n",
    "print(joints.shape, verts.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I want to highlight that, the input format of `smplx.SMPLX` is not similar to `smplx.SMPL`'s. For example, the shape of the tensor inputted to `body_pose`, changes from `(B, J, 3)` to `(B, J'*3)`, where `J = 23` and `J' = 21`, so does the `global_orient`, which changes from `(B, 1, 3)` to `(B, 1*3)`.\n",
    "\n",
    "> The reason why `J'` decreases to 21 is that, SMPL-X decouples the pose of hands from body pose to `left_hand_pose` and `right_hand_pose`, so that it can express more complex hand movements. Check the paper for details if you are interested. Hands are not what we are going to talk about here.\n",
    "\n",
    "One parameter I want to address is `transl`. You shouldn't share translation between SMPL and SMPLX, just like other parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "smplx_wis3d = Wis3D(\n",
    "        pm.outputs / 'wis3d',\n",
    "        'SMPLX',\n",
    "    )\n",
    "\n",
    "smplx_wis3d.add_motion_verts(\n",
    "        verts  = verts[:1],\n",
    "        name   = f'smplx_T_pose',\n",
    "        offset = 0,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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": [
    "### SMPL-X v.s. SMPL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not only the input format changes, but also the output format changes. The `vertices` of SMPL-X model is different from SMPL model's. The number of vertices of SMPL-X model is `10475`, while SMPL model's is `6890`.\n",
    "\n",
    "And you can see that the T-Pose of SMPL-X model is more 'relax' than SMPL model's. Now let's put them together to see the differences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING: You are using a SMPL model, with only 10 shape coefficients.\n"
     ]
    }
   ],
   "source": [
    "# Generate the SMPL T-pose.\n",
    "body_model_smpl = SMPL( model_path=pm.inputs/'body_models'/'smpl', gender=\"neutral\")\n",
    "smpl_out = body_model_smpl()\n",
    "smpl_verts = smpl_out.vertices\n",
    "smpl_joints = smpl_out.joints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "smpl_vs_smplx_wis3d = Wis3D(\n",
    "        pm.outputs / 'wis3d',\n",
    "        'SMPL v.s. SMPLX',\n",
    "    )\n",
    "\n",
    "# Visualize SMPL results.\n",
    "smpl_vs_smplx_wis3d.add_motion_verts(\n",
    "    verts  = smpl_verts[:1],\n",
    "    name   = f'smpl_T_pose',\n",
    "    offset = 0,\n",
    ")\n",
    "\n",
    "smpl_vs_smplx_wis3d.add_motion_skel(\n",
    "    joints = smpl_joints[:1, :22],\n",
    "    bones  = Skeleton_SMPL22.bones,\n",
    "    colors = Skeleton_SMPL22.bone_colors,\n",
    "    name   = f'smpl_T_pose',\n",
    "    offset = 0,\n",
    ")\n",
    "\n",
    "# Visualize SMPL-X results.\n",
    "smpl_vs_smplx_wis3d.add_motion_verts(\n",
    "    verts  = verts[:1],\n",
    "    name   = f'smplx_T_pose',\n",
    "    offset = 0,\n",
    ")\n",
    "\n",
    "smpl_vs_smplx_wis3d.add_motion_skel(\n",
    "    joints = joints[:1, :22],\n",
    "    bones  = Skeleton_SMPL22.bones,\n",
    "    colors = Skeleton_SMPL22.bone_colors,\n",
    "    name   = f'smplx_T_pose',\n",
    "    offset = 0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "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": [
    "## Additional Resources / Questions & Answers\n",
    "\n",
    "1. Can I transfer the parameters between SMPL, SMPL-H and SMPL-X?\n",
    "   - Yes, check the URLs below:\n",
    "     - https://github.com/vchoutas/smplx/blob/main/transfer_model/README.md"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lh",
   "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
}