Inference_from_any_pose.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d96a4911-2cfc-4234-9098-453018ec4436",
   "metadata": {},
   "source": [
    "# Inference from any pose"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6bd975a-3960-43ce-8ede-2ac20fb3af2c",
   "metadata": {},
   "source": [
    "With this notebook you can run inference of HOOD starting with a garment in any pose. For the body sequence you can use either:\n",
    "* Sequence of SMPL parameters (see `$HOOD_DATA/fromanypose/pose_sequence.pkl` as an example) or\n",
    "* Sequence of meshes (see `$HOOD_DATA/fromanypose/mesh_sequence.pkl` as an example)\n",
    "\n",
    "You also need to have a garment mesh (as `.obj` file) which is aligned with the first frame of your body sequence.\n",
    "\n",
    "In this notebook we first show how to convert a garment mesh in `.obj` format into a `.pkl` template used in HOOD. \n",
    "\n",
    "Then, we show how to use the configuration file `aux/from_any_pose.yaml` to run inference with over albitrary SMPL sequence and arbitrary mesh sequence.\n",
    "\n",
    "Note that, pinned vertices are not yet supported in this example, so do not expect it to work with lower body garments (they would slide down :c)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2b141c5-3f89-4d81-90ae-566296bd6e23",
   "metadata": {},
   "source": [
    "## Set environmental variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec571b31-88dc-4d7d-b975-d0c8b8eae87a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "HOOD_PROJECT = \"/path/to/hood/repository\"\n",
    "HOOD_DATA = \"/path/to/hood/data\"\n",
    "\n",
    "os.environ[\"HOOD_PROJECT\"] = HOOD_PROJECT\n",
    "os.environ[\"HOOD_DATA\"] = HOOD_DATA"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a91eba6-68a4-401b-bbfe-33292ca5ba4b",
   "metadata": {},
   "source": [
    "## Create template file from the .obj file:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4b41a41-90b3-48b5-8447-858797c9dd53",
   "metadata": {},
   "source": [
    "Use `utils.mesh_creation::obj2template()` function to convert an `.obj` file into a template dictionary and then save it with `pickle_dump`\n",
    "\n",
    "`$HOOD_DATA/fromanypose/tshirt.obj` is provided as an example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0568c817-32a6-4dfd-9ece-b5005dcc800b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import trimesh\n",
    "from utils.mesh_creation import obj2template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a972f1e0-af66-4b8a-9b4e-f0ab41807a7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "obj_path = Path(DEFAULTS.data_root) / 'fromanypose' / 'tshirt.obj'\n",
    "out_template_path = Path(DEFAULTS.data_root) / 'fromanypose' / 'tshirt.pkl'\n",
    "\n",
    "template_dict = obj2template(obj_path)\n",
    "pickle_dump(template_dict, out_template_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4ec39d2-8dd8-4161-bec6-7a6390ba1ff4",
   "metadata": {},
   "source": [
    "## Inference with a SMPL sequence"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c84b2f1-f186-4559-ae37-b891f2f4b664",
   "metadata": {},
   "source": [
    "Here we use a sequence of SMPL parameters as a body sequence. The garment template needs to be aligned with the first frame of the sequence.\n",
    "\n",
    "The SMPL pose sequence has to be a `.pkl` file containing a dictionary of the following items (see `$HOOD_DATA/fromanypose/pose_sequence.pkl` as an example):\n",
    "* `body_pose`: np.array [N, 69]\n",
    "* `global_orient`: np.array [N, 3]\n",
    "* `transl`: np.array [N, 3]\n",
    "* `betas`: np.array [10,]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "119b712c-2151-4453-98a7-4fada9df4061",
   "metadata": {},
   "source": [
    "### Edit `aux/from_any_pose` config"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f53330ee-1013-4e21-a7b3-e86cab76a7a4",
   "metadata": {},
   "source": [
    "First, change the configuration file `configs/aux/from_any_pose.yaml` to set your garment template and the pose sequence"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "452270f3-e509-47fe-8784-a3de6d498104",
   "metadata": {},
   "source": [
    "Go to `configs/aux/from_any_pose.yaml` and edit following fields there:\n",
    "\n",
    "- `dataloader.from_any_pose.scanreg.pose_sequence_type`, type of the pose sequence can be either `smpl` or `mesh`. For this example, set `smpl`\n",
    "- `dataloader.from_any_pose.scanreg.pose_sequence_path` path to the smpl pose sequence relative to `$HOOD_DATA`. For this example we use `fromanypose/pose_sequence.pkl`\n",
    "- `dataloader.from_any_pose.scanreg.garment_template_path` path to the .pkl file with the garment template  For this example we use `fromanypose/tshirt.pkl`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f760ca04-f5f7-4543-89bc-9a46a01896eb",
   "metadata": {},
   "source": [
    "### Load runner and dataloader from `from_any_pose` config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f676f90-1e9a-446a-85c8-7ccf61fbc53d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.validation import Config as ValidationConfig\n",
    "from utils.arguments import load_params, create_modules\n",
    "from utils.arguments import load_params\n",
    "from utils.common import move2device, pickle_dump\n",
    "from utils.defaults import DEFAULTS\n",
    "from pathlib import Path\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86dc5a9d-6d23-4e70-afff-2530c7b07332",
   "metadata": {},
   "outputs": [],
   "source": [
    "modules, config = load_params('aux/from_any_pose')\n",
    "dataloader_m, runner_module, runner, aux_modules = create_modules(modules, config)\n",
    "dataloader = dataloader_m.create_dataloader()\n",
    "\n",
    "checkpoint_path = Path(DEFAULTS.data_root) / 'trained_models' / 'postcvpr.pth'\n",
    "state_dict =  torch.load(checkpoint_path)\n",
    "runner.load_state_dict(state_dict['training_module'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c952eac2-864b-4a4f-b20b-13bcaafccd26",
   "metadata": {},
   "source": [
    "### load sample, infer and save trajectories\n",
    "\n",
    "To visualise saved trajectories, see `Inference.ipynb::write_video`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b411d7f2-3d83-431d-b8a3-08937646f5bf",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sample = next(iter(dataloader))\n",
    "\n",
    "trajectories_dict = runner.valid_rollout(sample)\n",
    "# Save the sequence to disc\n",
    "out_path = Path(DEFAULTS.data_root) / 'temp' / 'output_f.pkl'\n",
    "print(f\"Rollout saved into {out_path}\")\n",
    "pickle_dump(dict(trajectories_dict), out_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cc63433-3d5a-4583-9636-1562ea1f62f2",
   "metadata": {},
   "source": [
    "## Inference with a mesh sequence"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffabf78b-9afb-409a-8cda-3433e9b93308",
   "metadata": {},
   "source": [
    "Here we use a sequence of arbitrary meshes.  The garment template needs to be aligned with the first frame of the sequence.\n",
    "\n",
    "The mesh sequence has to be a `.pkl` file containing a dictionary of the following items (see `$HOOD_DATA/fromanypose/mesh_sequence.pkl` as an example):\n",
    "* `verts`: np.array [N, 3]\n",
    "* `faces`: np.array [F, 3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44b639c8-1394-422d-af5a-fe0b6b104fb8",
   "metadata": {},
   "source": [
    "### Edit `aux/from_any_pose` config\n",
    "\n",
    "Change the configuration file `configs/aux/from_any_pose.yaml` to set your garment template and the pose sequence\n",
    "\n",
    "Go to `configs/aux/from_any_pose.yaml` and edit following fields there:\n",
    "\n",
    "- `dataloader.from_any_pose.scanreg.pose_sequence_type`, type of the pose sequence can be either `smpl` or `mesh`. For this example, set `mesh`\n",
    "- `dataloader.from_any_pose.scanreg.pose_sequence_path` path to the smpl pose sequence relative to `$HOOD_DATA`. For this example we use `fromanypose/mesh_sequence.pkl`\n",
    "- `dataloader.from_any_pose.scanreg.garment_template_path` path to the .pkl file with the garment template  For this example we use `fromanypose/tshirt.pkl`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "facedbd7-2d8b-4131-bcc9-a9c688c5b5ed",
   "metadata": {},
   "source": [
    "### Load runner and dataloader from `from_any_pose` config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a143d7ec-7e0e-4765-86b4-463eaed6baeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.validation import Config as ValidationConfig\n",
    "from utils.arguments import load_params, create_modules\n",
    "from utils.arguments import load_params\n",
    "from utils.common import move2device, pickle_dump\n",
    "from utils.defaults import DEFAULTS\n",
    "from pathlib import Path\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f07db298-afb1-4b72-af26-7842a6947626",
   "metadata": {},
   "outputs": [],
   "source": [
    "modules, config = load_params('aux/from_any_pose')\n",
    "dataloader_m, runner_module, runner, aux_modules = create_modules(modules, config)\n",
    "dataloader = dataloader_m.create_dataloader()\n",
    "\n",
    "checkpoint_path = '/mnt/sdb1/hood_public/trained_models/postcvpr.pth'\n",
    "state_dict =  torch.load(checkpoint_path)\n",
    "runner.load_state_dict(state_dict['training_module'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9a43ad4-5bd9-4ae9-939b-f804dea17dd0",
   "metadata": {},
   "source": [
    "### load sample, infer and save trajectories\n",
    "\n",
    "To visualise saved trajectories, see `Inference.ipynb::write_video`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d81b569-69d4-4989-9035-6127d39e26b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = next(iter(dataloader))\n",
    "\n",
    "trajectories_dict = runner.valid_rollout(sample)\n",
    "# Save the sequence to disc\n",
    "out_path = Path(DEFAULTS.data_root) / 'temp' / 'output_m.pkl'\n",
    "print(f\"Rollout saved into {out_path}\")\n",
    "pickle_dump(dict(trajectories_dict), out_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f936192-3ff7-4eb5-801e-6c11d1da4e49",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hood",
   "language": "python",
   "name": "hood"
  },
  "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}