raw_pose_processing_beat2_smplx.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract Poses from Amass Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib notebook\n",
    "%matplotlib inline\n",
    "\n",
    "import sys, os\n",
    "import zipfile\n",
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "\n",
    "from human_body_prior.tools.omni_tools import copy2cpu as c2c\n",
    "\n",
    "os.environ['PYOPENGL_PLATFORM'] = 'egl'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose the device to run the body model on.\n",
    "comp_device = torch.device(\"cuda:1\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from human_body_prior.body_model.body_model import BodyModel\n",
    "\n",
    "neutral_bm_path = './body_models/smplx/SMPLX_NEUTRAL_2020.npz'\n",
    "neutral_bm = BodyModel(bm_fname=neutral_bm_path, num_betas=300, num_expressions=100).to(comp_device)\n",
    "faces = c2c(neutral_bm.f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num files: 2048\n"
     ]
    }
   ],
   "source": [
    "import pathlib\n",
    "\n",
    "beat2_d = pathlib.Path('beat2_data')\n",
    "motion_fs = [pathlib.Path(f'{rt}/{f}')\n",
    "             for rt,ds,fs in os.walk(beat2_d)\n",
    "             for f in fs\n",
    "             if f.endswith('.npz')]\n",
    "motion_fs.sort()\n",
    "print('num files:',len(motion_fs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def beat2_to_pose(motion_f, joints_f, i_motion):\n",
    "    bdata = np.load(motion_f, allow_pickle=True)\n",
    "    assert bdata['mocap_frame_rate'].item() == 30\n",
    "    assert bdata['model'].item() == 'smplx2020'\n",
    "    assert bdata['gender'].item() == 'neutral'\n",
    "    bm = neutral_bm\n",
    "    T = bdata['poses'].shape[0]\n",
    "    B = 1024 # batch size\n",
    "    results = list()\n",
    "    for i_beg in range(0,T,B):\n",
    "        bdata_poses = bdata['poses'][i_beg:i_beg+B]\n",
    "        bdata_trans = bdata['trans'][i_beg:i_beg+B]\n",
    "        bdata_expression = bdata['expressions'][i_beg:i_beg+B]\n",
    "        body_parms = {\n",
    "                'root_orient': torch.Tensor(bdata_poses[:,:3]).to(comp_device),\n",
    "                'pose_body'  : torch.Tensor(bdata_poses[:,3:66]).to(comp_device),\n",
    "                'pose_hand'  : torch.Tensor(bdata_poses[:,75:]).to(comp_device),\n",
    "                'trans'      : torch.Tensor(bdata_trans).to(comp_device),\n",
    "                'betas'      : torch.Tensor(np.repeat(bdata['betas'][np.newaxis], repeats=len(bdata_trans), axis=0)).to(comp_device),\n",
    "                'expression' : torch.Tensor(bdata_expression).to(comp_device)\n",
    "            }    \n",
    "        with torch.no_grad():\n",
    "            body = bm(**body_parms)\n",
    "        pose_seq_np = body.Jtr.detach().cpu().numpy()\n",
    "        # Unlike AMASS datasets, here the XZ plane is already the ground plane.\n",
    "        # So no need for `trans_matrix`.\n",
    "        results.append(pose_seq_np)\n",
    "    results = np.concatenate(results,axis=0)\n",
    "    assert results.shape[0] == T\n",
    "    np.save(joints_f,results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "beat2-to-pose:   0%|▏                                                                                              | 4/2048 [00:11<2:12:40,  3.89s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/vision/vision_data_2/VGGSound_shards_fixed/shrinidhi/meshes_beat2_pose_data/beat_chinese_v2.0.0/smplxflame_30/12_zhao_2_103_103\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "beat2-to-pose:   1%|▌                                                                                             | 12/2048 [00:20<1:15:40,  2.23s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/vision/vision_data_2/VGGSound_shards_fixed/shrinidhi/meshes_beat2_pose_data/beat_chinese_v2.0.0/smplxflame_30/12_zhao_2_111_111\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "beat2-to-pose:   1%|█                                                                                             | 23/2048 [00:31<1:37:12,  2.88s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/vision/vision_data_2/VGGSound_shards_fixed/shrinidhi/meshes_beat2_pose_data/beat_chinese_v2.0.0/smplxflame_30/12_zhao_2_1_1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "beat2-to-pose:   1%|█▍                                                                                            | 30/2048 [00:41<1:13:57,  2.20s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/vision/vision_data_2/VGGSound_shards_fixed/shrinidhi/meshes_beat2_pose_data/beat_chinese_v2.0.0/smplxflame_30/12_zhao_2_2_2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "beat2-to-pose: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 2048/2048 [09:51<00:00,  3.46it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "bad zip files: 0\n"
     ]
    }
   ],
   "source": [
    "pose_data_d = pathlib.Path('pose_data_beat2')\n",
    "bad_zip_files = list()\n",
    "for i_motion,motion_f in enumerate(tqdm(motion_fs,desc='beat2-to-pose',ncols=150)):\n",
    "    joints_f = pose_data_d / motion_f.relative_to(beat2_d).with_suffix('.npy')\n",
    "    if joints_f.is_file():\n",
    "        joints_f.unlink()\n",
    "    joints_f.parent.mkdir(parents=True,exist_ok=True)\n",
    "    try:\n",
    "        beat2_to_pose(motion_f,joints_f,i_motion)\n",
    "    except zipfile.BadZipFile:\n",
    "        bad_zip_files.append(motion_f)\n",
    "print('bad zip files:',len(bad_zip_files))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Segment, Mirror and Relocate Motions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "import csv\n",
    "import os\n",
    "import pathlib\n",
    "\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "\n",
    "bm_params_f = pathlib.Path('body_models/smplx/SMPLX_NEUTRAL_2020.npz')\n",
    "index_f = pathlib.Path('index.csv')\n",
    "pose_data_d  = pathlib.Path('pose_data_beat2')\n",
    "joints_d = pathlib.Path('joints_beat2')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Find the corresponding left/right joints from model npy file. We will mirror left/right joints to augment data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num joints to swap: 24\n",
      "left joints: [30, 23, 13, 34, 35, 36, 39, 38, 37, 27, 26, 25, 10, 4, 18, 31, 32, 33, 28, 29, 20, 16, 1, 7]\n",
      "right joints: [45, 24, 14, 49, 50, 51, 54, 53, 52, 42, 41, 40, 11, 5, 19, 46, 47, 48, 43, 44, 21, 17, 2, 8]\n",
      "joints to drop: [22, 23, 24]\n",
      "L_Hip      ( 1) <--> ( 2) R_Hip\n",
      "L_Knee     ( 4) <--> ( 5) R_Knee\n",
      "L_Ankle    ( 7) <--> ( 8) R_Ankle\n",
      "L_Foot     (10) <--> (11) R_Foot\n",
      "L_Collar   (13) <--> (14) R_Collar\n",
      "L_Shoulder (16) <--> (17) R_Shoulder\n",
      "L_Elbow    (18) <--> (19) R_Elbow\n",
      "L_Wrist    (20) <--> (21) R_Wrist\n",
      "L_Eye      (23) <--> (24) R_Eye\n",
      "L_Index1   (25) <--> (40) R_Index1\n",
      "L_Index2   (26) <--> (41) R_Index2\n",
      "L_Index3   (27) <--> (42) R_Index3\n",
      "L_Middle1  (28) <--> (43) R_Middle1\n",
      "L_Middle2  (29) <--> (44) R_Middle2\n",
      "L_Middle3  (30) <--> (45) R_Middle3\n",
      "L_Pinky1   (31) <--> (46) R_Pinky1\n",
      "L_Pinky2   (32) <--> (47) R_Pinky2\n",
      "L_Pinky3   (33) <--> (48) R_Pinky3\n",
      "L_Ring1    (34) <--> (49) R_Ring1\n",
      "L_Ring2    (35) <--> (50) R_Ring2\n",
      "L_Ring3    (36) <--> (51) R_Ring3\n",
      "L_Thumb1   (37) <--> (52) R_Thumb1\n",
      "L_Thumb2   (38) <--> (53) R_Thumb2\n",
      "L_Thumb3   (39) <--> (54) R_Thumb3\n"
     ]
    }
   ],
   "source": [
    "bm_params = np.load(bm_params_f,allow_pickle=True)\n",
    "joint2ind = bm_params['joint2num'].item()\n",
    "ind2joint = {v:k\n",
    "             for k,v in joint2ind.items()}\n",
    "l_joints,r_joints = list(),list()\n",
    "for j in joint2ind:\n",
    "    if j.startswith('L_'):\n",
    "        l_j = j\n",
    "        r_j = j.replace('L_','R_')\n",
    "        l_joints.append(joint2ind[l_j])\n",
    "        r_joints.append(joint2ind[r_j])\n",
    "joints_to_drop = [joint2ind['Jaw'],\n",
    "                  joint2ind['L_Eye'],\n",
    "                  joint2ind['R_Eye']]\n",
    "\n",
    "print('num joints to swap:',len(l_joints))\n",
    "print('left joints:',l_joints)\n",
    "print('right joints:',r_joints)\n",
    "print('joints to drop:',joints_to_drop)\n",
    "for l,r in sorted(zip(l_joints,r_joints)):\n",
    "    print(f'{ind2joint[l]:10} ({l:2}) <--> ({r:2}) {ind2joint[r]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Mirror each file and split to max sequence length of 200 each because MotionGPT uses these limits.  \n",
    "Following code is explained in [this issue](https://github.com/EricGuo5513/HumanML3D/issues/20).\n",
    "```\n",
    "data_m[...,0] *= -1\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "mirror & split files:   0%|                                                                                       | 0/2048 [00:00<?, ?it/s, samples=6]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "mirror & split files: 100%|████████████████████████████████████████████████████████████████████████| 2048/2048 [08:08<00:00,  4.19it/s, samples=73548]\n"
     ]
    }
   ],
   "source": [
    "data_fs = [pathlib.Path(f'{rt}/{f}')\n",
    "           for rt,ds,fs in os.walk(pose_data_d)\n",
    "           for f in fs\n",
    "           if f.endswith('.npy')]\n",
    "data_fs.sort()\n",
    "max_seq_len = 200\n",
    "pose_data_to_joints_map = ['file,id']\n",
    "pbar = tqdm(data_fs,desc='mirror & split files',ncols=150)\n",
    "for i,f in enumerate(pbar):\n",
    "    data = np.load(f)\n",
    "    data_m = data.copy()\n",
    "    data_m[:,l_joints] = data[:,r_joints]\n",
    "    data_m[:,r_joints] = data[:,l_joints]\n",
    "    # Unlike AMASS datasets, we have flip the mirrored version.\n",
    "    data_m[...,0] *= -1\n",
    "    data = np.delete(data,joints_to_drop,axis=1)\n",
    "    data_m = np.delete(data_m,joints_to_drop,axis=1)\n",
    "    for j,beg in enumerate(range(0,data.shape[0],max_seq_len)):\n",
    "        id = f'{i:06}_{j:03}'\n",
    "        id_m = f'M{i:06}_{j:03}'\n",
    "        out_f = joints_d / f'{id}.npy'\n",
    "        out_m_f = joints_d / f'{id_m}.npy'\n",
    "        if out_f.is_file():\n",
    "            out_f.unlink()\n",
    "        if out_m_f.is_file():\n",
    "            out_m_f.unlink()\n",
    "        np.save(out_f,data[beg:beg+max_seq_len])\n",
    "        np.save(out_m_f,data_m[beg:beg+max_seq_len])\n",
    "        pose_data_to_joints_map.append(f'{f},{id}')\n",
    "        pbar.set_postfix({'samples':2*len(pose_data_to_joints_map)})\n",
    "_ = open('pose_data_to_joints_map_beat2.txt','w').write('\\n'.join(pose_data_to_joints_map) + '\\n')\n",
    "pbar.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Write fake text files for each example. Copy a single text file from AMASS dataset to all the examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "EnvironmentNameNotFound: Could not find conda environment: avjoint\n",
      "You can list all discoverable environments with `conda info --envs`.\n",
      "\n",
      "\n",
      "\n",
      "EnvironmentNameNotFound: Could not find conda environment: avjoint\n",
      "You can list all discoverable environments with `conda info --envs`.\n",
      "\n",
      "\n",
      "\n",
      "EnvironmentNameNotFound: Could not find conda environment: avjoint\n",
      "You can list all discoverable environments with `conda info --envs`.\n",
      "\n",
      "\n",
      "\n",
      "EnvironmentNameNotFound: Could not find conda environment: avjoint\n",
      "You can list all discoverable environments with `conda info --envs`.\n",
      "\n",
      "\n",
      "\n",
      "EnvironmentNameNotFound: Could not find conda environment: avjoint\n",
      "You can list all discoverable environments with `conda info --envs`.\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!git checkout main -- HumanML3D/texts.zip\n",
    "!rm -rf HumanML3D/texts\n",
    "!unzip -q HumanML3D/texts.zip -d HumanML3D\n",
    "!rm HumanML3D/texts.zip\n",
    "!mv HumanML3D/texts HumanML3D/texts_orig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "texts: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████| 73546/73546 [04:03<00:00, 302.38it/s]\n"
     ]
    }
   ],
   "source": [
    "import shutil\n",
    "\n",
    "src_f = pathlib.Path('HumanML3D/texts_orig/000000.txt')\n",
    "texts_d = pathlib.Path('HumanML3D_beat2/texts')\n",
    "texts_d.mkdir(parents=True,exist_ok=True)\n",
    "data_fs = list(joints_d.iterdir())\n",
    "data_fs.sort()\n",
    "\n",
    "for f in tqdm(data_fs,'texts',ncols=150):\n",
    "    txt_f = f.with_suffix('.txt').name\n",
    "    txt_f = texts_d / txt_f\n",
    "    shutil.copyfile(src_f,txt_f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "EnvironmentNameNotFound: Could not find conda environment: avjoint\n",
      "You can list all discoverable environments with `conda info --envs`.\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!rm -r HumanML3D/texts_orig"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "shrik_mgpt",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}