visualize_mesh.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualize Mesh\n",
    "\n",
    "This notebook provides a demo of how to generate SMPL mesh from 3D skeletons of MUGL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import ipyvolume as ipv\n",
    "import h5py\n",
    "import os\n",
    "\n",
    "from smplx import SMPL\n",
    "import pickle\n",
    "\n",
    "import torch \n",
    "import torch.nn as nn\n",
    "\n",
    "from model import *\n",
    "from rotation.rotation import rot6d_to_rotmat, batch_rigid_transform\n",
    "from torch.autograd import Variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "skeleton = np.load('./files/skeleton.npy')\n",
    "viewpoint = np.load('./files/viewpoint.npy')\n",
    "beta = torch.tensor([-0.1474,  0.0632,  0.7616,  2.9261,  0.3609,  0.2267, -0.3828,  0.3000,\n",
    "         0.5667,  0.0230])\n",
    "\n",
    "latent_dim = 352\n",
    "device = torch.device('cuda:0')\n",
    "num_class = 120\n",
    "main_path = \"/ssd_scratch/cvit/debtanu.gupta/\"\n",
    "\n",
    "smpl = SMPL('./files',batch_size=32)\n",
    "\n",
    "with open(\"./files/SMPL_NEUTRAL.pkl\",'rb') as f:\n",
    "    smpl_data = pickle.load(f,encoding='latin1')\n",
    "\n",
    "parent_array = list(smpl_data['kintree_table'][0][:24])\n",
    "parent_array[0] = -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def infer(model, label,v=0):\n",
    "    '''\n",
    "        this function generates samples\n",
    "        model: model object as input\n",
    "        label: calss label\n",
    "        v: viewpoint\n",
    "    '''\n",
    "    model.eval()\n",
    "    # z = torch.randn(6, latent_dim).to(device).float()\n",
    "    \n",
    "    # y = np.repeat(np.arange(3),2)\n",
    "    # y = np.arange(num_class)\n",
    "    y = np.repeat(label,10)\n",
    "#     rot_list = []\n",
    "#     for i in y:\n",
    "#         idx = np.where(label==i)\n",
    "#         rot_lbl = rot[idx]\n",
    "#         rand = np.random.randint(rot_lbl.shape[0])\n",
    "#         rot_list.append(rot_lbl[rand])\n",
    "    rot = np.repeat(viewpoint[v:v+1], 10, axis=0)\n",
    "    # rot = torch.tensor(rot[:,0,:]).to(device).float()\n",
    "    rot = rot[:,0,:].reshape((rot.shape[0],48,6))\n",
    "    rot = rot[:,0,:]\n",
    "    rot = torch.tensor(rot).to(device).float()\n",
    "\n",
    "    label = np.zeros((y.shape[0], num_class))\n",
    "    label[np.arange(y.shape[0]), y] = 1\n",
    "    label = torch.tensor(label).to(device).float()\n",
    "    with torch.no_grad():\n",
    "        m, v = model.gaussian_parameters(model.z_pre.squeeze(0), dim=0)\n",
    "        idx = torch.distributions.categorical.Categorical(model.pi).sample((label.shape[0],))\n",
    "        m, v = m[idx], v[idx]\n",
    "        z = model.sample_gaussian(m, v)\n",
    "        \n",
    "        z = torch.cat((z,label,rot), dim=1)\n",
    "        z = model.latent2hidden(z)\n",
    "        z = z.reshape((z.shape[0], 4, -1))\n",
    "        pred = model.decoder_net(z)\n",
    "        root_pred = model.root_traj(z).unsqueeze(2)\n",
    "\n",
    "        N,T,_ = pred.shape\n",
    "        alpha = rot6d_to_rotmat(torch.tensor(pred.reshape((N,T,2,144)))).view((N,T,2,24,3,3)).float()\n",
    "        betas = beta.view((1,1,1,10)).repeat((N,T,2,1))\n",
    "        alpha = alpha.to(betas.device)\n",
    "        root_pred = root_pred.to(betas.device)\n",
    "\n",
    "        # pred_3d1 = fkt(pred[:,:,:144].contiguous(), skeleton)\n",
    "        # pred_3d2 = fkt(pred[:,:,144:].contiguous(), skeleton)\n",
    "        # # pred_3d = fkt(pred, skeleton).cpu().data.numpy()\n",
    "        # pred_3d1 = pred_3d1.reshape((pred_3d1.shape[0], pred_3d1.shape[1], 24,-1)).cpu().data.numpy()\n",
    "        # pred_3d2 = pred_3d2.reshape((pred_3d2.shape[0], pred_3d2.shape[1], 24,-1)).cpu().data.numpy()\n",
    "        # root_pred = root_pred.cpu().data.numpy()\n",
    "        # pred_3d2 = pred_3d2 + root_pred\n",
    "\n",
    "        return alpha, betas, root_pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot(alpha, betas, root,ind, single=False, save_gif=False, save_name='example'):\n",
    "    '''\n",
    "        This function takes 2 persons as input and plots them\n",
    "        skeleton_motion1: Skeleton of person 1\n",
    "        skeleton_motion2: Skeleton of person 2\n",
    "        single: True if single person class else false\n",
    "        save_fig: Save gif file if True\n",
    "    '''\n",
    "    # print(betas.device, alpha.device, root.device)\n",
    "    output1 = smpl(betas=betas[ind,:,0,:],body_pose=alpha[ind,:,0,1:], global_orient=alpha[ind,:,0,0].unsqueeze(1), pose2rot=False)\n",
    "    output2 = smpl(betas=betas[ind,:,1,:],body_pose=alpha[ind,:,1,1:], global_orient=alpha[ind,:,1,0].unsqueeze(1), pose2rot=False)\n",
    "\n",
    "    v1 = output1.vertices.data.cpu().numpy()\n",
    "    skeleton_motion1 = output1.joints.data.cpu().numpy()\n",
    "\n",
    "    v2 = output2.vertices.data.cpu().numpy() + root[ind].data.cpu().numpy()\n",
    "    skeleton_motion2 = output2.joints.data.cpu().numpy() + root[ind].data.cpu().numpy()\n",
    "\n",
    "    v1[:,:,1] *= -1\n",
    "    skeleton_motion1[:,:,1] *= -1\n",
    "    v2[:,:,1] *= -1\n",
    "    skeleton_motion2[:,:,1] *= -1\n",
    "\n",
    "    v1[:,:,2] *= -1\n",
    "    skeleton_motion1[:,:,2] *= -1\n",
    "    v2[:,:,2] *= -1\n",
    "    skeleton_motion2[:,:,2] *= -1\n",
    "\n",
    "    fig = ipv.figure(height=800,width=800)\n",
    "\n",
    "    plot_mesh1 = ipv.plot_trisurf(x=v1[:,:,0],y=v1[:,:,1],z=v1[:,:,2],triangles=smpl_data['f'],color=[1,1,.5,1])\n",
    "    plot_mesh1.material.transparent = True\n",
    "    plot_mesh1.material.side = \"FrontSide\"\n",
    "    \n",
    "    a = np.arange(-np.max(skeleton_motion2[:,:,:])-1,np.max(skeleton_motion1[:,:,:])+1)\n",
    "    if not single:\n",
    "        a = np.arange(np.min(skeleton_motion1[:,:,:])-1,np.max(skeleton_motion1[:,:,:])+1)\n",
    "        plot_mesh2 = ipv.plot_trisurf(x=v2[:,:,0],y=v2[:,:,1],z=v2[:,:,2],triangles=smpl_data['f'],color=[1,0,1,1])\n",
    "        plot_mesh2.material.transparent = True\n",
    "        plot_mesh2.material.side = \"FrontSide\"\n",
    "    x,z = np.meshgrid(a,a)\n",
    "    y = np.ones_like(x)*np.min(skeleton_motion1[:,:,1])\n",
    "    s1 = ipv.scatter(skeleton_motion1[:,:,0],skeleton_motion1[:,:,1],skeleton_motion1[:,:,2],size=1,color='indigo',marker='sphere')\n",
    "    if not single:\n",
    "        s2 = ipv.scatter(skeleton_motion2[:,:,0],skeleton_motion2[:,:,1],skeleton_motion2[:,:,2],size=1,color='indigo',marker='sphere')\n",
    "\n",
    "    if single:\n",
    "        anim_list = [plot_mesh1,s1]\n",
    "    else:\n",
    "        anim_list = [plot_mesh1,s1,plot_mesh2,s2]\n",
    "    for i,p in enumerate(parent_array): # Run loop for each bone\n",
    "            if p == -1:\n",
    "                continue\n",
    "            b1 = ipv.plot(np.array([skeleton_motion1[:,i,0],skeleton_motion1[:,p,0]]).T,np.array([skeleton_motion1[:,i,1],skeleton_motion1[:,p,1]]).T,np.array([skeleton_motion1[:,i,2],skeleton_motion1[:,p,2]]).T ,size=10, color='darkviolet')\n",
    "            \n",
    "            if not single:\n",
    "                b2 = ipv.plot(np.array([skeleton_motion2[:,i,0],skeleton_motion2[:,p,0]]).T,np.array([skeleton_motion2[:,i,1],skeleton_motion2[:,p,1]]).T,np.array([skeleton_motion2[:,i,2],skeleton_motion2[:,p,2]]).T ,size=10, color='darkviolet')\n",
    "            anim_list.append(b1)\n",
    "            if not single:\n",
    "                anim_list.append(b2)\n",
    "    \n",
    "    ipv.plot_surface(x,y,z, color='lightgray')\n",
    "    ipv.plot_wireframe(x,y,z,color='black')\n",
    "    ipv.animation_control(anim_list)\n",
    "    ipv.style.background_color('dark')\n",
    "    ipv.style.box_off()\n",
    "    ipv.style.axes_off()\n",
    "    ipv.xyzlim(min(np.min(v1),np.min(v2)),max(np.max(v1),np.max(v2)))\n",
    "    ipv.show()\n",
    "    \n",
    "    if save_gif:\n",
    "        def slide(figure, framenr, fraction):\n",
    "            for a in anim_list:\n",
    "                if a.sequence_index == skeleton_motion1.shape[0]:\n",
    "                    a.sequence_index = 0\n",
    "                a.sequence_index += 1        \n",
    "        ipv.movie(save_name + '.gif', slide, fps=5, frames=skeleton_motion1.shape[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model(latent_dim).to(device)\n",
    "model.load_state_dict(torch.load('./checkpoints/' + 'model_199.pt', map_location=torch.device('cpu')))\n",
    "print('model loaded..')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cls_lbl = 58\n",
    "alpha, betas, root_pred = infer(model, cls_lbl, v=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = 6\n",
    "plot(alpha, betas, root_pred, idx, single=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ENV",
   "language": "python",
   "name": "env"
  },
  "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}