02_preprocess_mocap_result.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import config\n",
    "import utils\n",
    "\n",
    "import json\n",
    "import pickle\n",
    "import os\n",
    "import numpy as np\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make directory for saving"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(config.skeleton_pkl_dir, exist_ok=True)\n",
    "os.makedirs(config.skeleton_npz_dir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "76cde7a9045f44d98f14b2b2a8d3e83b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/69 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "video_idx = [v_idx for v_idx in os.listdir(config.mocap_dir) if v_idx[0]!='.']\n",
    "for v_idx in video_idx:\n",
    "    result_dir = os.path.join(config.mocap_dir, v_idx)\n",
    "    bbox_dir = os.path.join(result_dir, 'bbox')\n",
    "    mocap_dir = os.path.join(result_dir, 'mocap')\n",
    "    \n",
    "    pkl_dir = os.path.join(config.skeleton_pkl_dir, v_idx)\n",
    "    npz_dir = os.path.join(config.skeleton_npz_dir, v_idx)\n",
    "    \n",
    "    os.makedirs(pkl_dir, exist_ok=True)\n",
    "    os.makedirs(npz_dir, exist_ok=True)\n",
    "    \n",
    "    for mocap_fp in tqdm(sorted(os.listdir(mocap_dir))):\n",
    "        scene_name = mocap_fp.split('/')[-1][:14]\n",
    "        \n",
    "        bbox_fp = os.path.join(bbox_dir, scene_name+'_bbox.json')\n",
    "        mocap_fp = os.path.join(mocap_dir, scene_name+'_prediction_result.pkl')\n",
    "        \n",
    "        ### load info\n",
    "        bbox = json.load(open(bbox_fp, 'rb'))\n",
    "        mocap = pickle.load(open(mocap_fp, 'rb'))\n",
    "        \n",
    "        ### get bbox size\n",
    "        bbox_np = np.array(bbox['body_bbox_list'])\n",
    "        bbox_size = [bb[2]*bb[3] for bb in bbox_np]\n",
    "\n",
    "        ### only choose the biggest bbox\n",
    "        max_idx = np.argmax(bbox_size)\n",
    "        curr_body_bbox = bbox_np[max_idx]\n",
    "        curr_mocap = mocap['pred_output_list'][max_idx]\n",
    "\n",
    "        ### extract body and hand pose\n",
    "        curr_body_pose = curr_mocap['pred_joints_img']\n",
    "        curr_open_pose = curr_body_pose[:25]\n",
    "        curr_aux_pose = curr_body_pose[25:]\n",
    "\n",
    "        curr_right_hand = curr_mocap['right_hand_joints_img_coord']\n",
    "        curr_left_hand = curr_mocap['left_hand_joints_img_coord']\n",
    "        \n",
    "        # extract hand bbox\n",
    "        if None not in mocap['hand_bbox_list']:\n",
    "            curr_left_hand_bbox = np.asarray(mocap['hand_bbox_list'][max_idx]['left_hand'])\n",
    "            curr_right_hand_bbox = np.asarray(mocap['hand_bbox_list'][max_idx]['right_hand'])\n",
    "            if curr_left_hand_bbox.all() == None:\n",
    "                curr_left_hand_bbox = np.array([0, 0, 0, 0])\n",
    "            if curr_right_hand_bbox.all() == None:\n",
    "                curr_right_hand_bbox = np.array([0, 0, 0, 0])\n",
    "\n",
    "            \n",
    "        ### extract simplified pose\n",
    "        curr_simple_pose = curr_open_pose[:15] # 15 by 3\n",
    "\n",
    "        ### process for .npz file.\n",
    "        left_hand_info, left_hand_array = utils.get_hand_info(curr_left_hand, 'left')\n",
    "        right_hand_info, right_hand_array = utils.get_hand_info(curr_right_hand, 'right')\n",
    "\n",
    "        left_heel_info, right_heel_info, heel_array = utils.get_heel_info(curr_open_pose)\n",
    "\n",
    "        head_info, head_array = utils.get_head_info(curr_open_pose)\n",
    "\n",
    "        ### prepare dictionary\n",
    "        pkl_saver = dict()\n",
    "        pkl_saver['open_pose'] = curr_open_pose\n",
    "        pkl_saver['aux_pose'] = curr_aux_pose\n",
    "        pkl_saver['simple_pose'] = curr_simple_pose\n",
    "        pkl_saver['left_hand'] = curr_left_hand\n",
    "        pkl_saver['right_hand'] = curr_right_hand\n",
    "        pkl_saver['hand_info'] = {'left':left_hand_info, 'right':right_hand_info}\n",
    "        pkl_saver['heel_info'] = {'left':left_heel_info, 'right':right_heel_info}\n",
    "        pkl_saver['head_info'] = head_info\n",
    "        pkl_saver['body_bbox'] = curr_body_bbox\n",
    "        pkl_saver['left_hand_bbox'] = curr_left_hand_bbox\n",
    "        pkl_saver['right_hand_bbox'] = curr_right_hand_bbox\n",
    "\n",
    "        ### prepare the array\n",
    "        npz_saver = [curr_simple_pose.flatten(), head_array, left_hand_array, right_hand_array, heel_array, \n",
    "                     curr_body_bbox, curr_left_hand_bbox, curr_right_hand_bbox]\n",
    "        npz_saver = np.concatenate(npz_saver)\n",
    "        # 0  ~ 45 : pose that was 15 by 3\n",
    "        # 45 ~ 51 : head pos(3), vec(3)\n",
    "        # 51 ~ 63 : left hand thumb vector(3), index vector(3), palm vector(3), wrist pos(3)\n",
    "        # 63 ~ 75 : right hand thumb vector(3), index vector(3), palm vector(3), wrist pos(3)\n",
    "        # 75 ~ 87 : left heel pos (3), left foot vec (3), right heel pos (3), right foot vec (3)\n",
    "        # 87 ~ 91 : body bbox (4) left hand bbox (4) right hand bbox(4)\n",
    "\n",
    "        ### save the data\n",
    "        pkl_fp = os.path.join(pkl_dir, scene_name+'.pkl')\n",
    "        npz_fp = os.path.join(npz_dir, scene_name+'.npz')\n",
    "        pickle.dump(pkl_saver, open(pkl_fp, 'wb'))\n",
    "        np.savez(npz_fp, npz_saver)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv_frankmocap",
   "language": "python",
   "name": "venv_frankmocap"
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}