_PY01-OptiTrack_Preprocessing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import math\n",
    "import datetime\n",
    "import time\n",
    "import os\n",
    "import glob\n",
    "import multiprocessing\n",
    "import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run py/constants.py\n",
    "%run py/take_processing_methods.py\n",
    "%run py/marker_processing_methods.py\n",
    "%run py/heatmap_methods.py\n",
    "%run py/grasp_property_methods.py\n",
    "%run py/rasterization_methods.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def get_transformed_marker(frame_index, marker_id, df, rot_matrices):\n",
    "    rot_matrix = rot_matrices[df.index.get_loc(frame_index)].T\n",
    "\n",
    "    rb_pivot =[df[df.columns[6]].ix[frame_index], df[df.columns[7]].ix[frame_index], df[df.columns[8]].ix[frame_index]]\n",
    "\n",
    "    marker = [df[marker_id + \"_X\"].ix[frame_index], df[marker_id + \"_Y\"].ix[frame_index], df[marker_id + \"_Z\"].ix[frame_index]]  \n",
    "    marker = inv_trans(marker, rb_pivot, rot_matrix)\n",
    "\n",
    "    return marker\n",
    "\n",
    "def quaternion_matrix(row, rigid_body_name):\n",
    "    \"\"\"\n",
    "    Converts a quaternion to a rotation matrix.\n",
    "    Source: http://fabiensanglard.net/doom3_documentation/37726-293748.pdf\n",
    "    RETURN: Rotation matrix as numpy array.\n",
    "    \"\"\"\n",
    "    x = row[rigid_body_name + '_X_Rotation']\n",
    "    y = row[rigid_body_name + '_Y_Rotation']\n",
    "    z = row[rigid_body_name + '_Z_Rotation']\n",
    "    w = row[rigid_body_name + '_W_Rotation']\n",
    "    \n",
    "    rot_matrix = [\n",
    "        1.0-2*y*y-2*z*z, 2.0*x*y+2*w*z, 2.0*x*z - 2*w*y,\n",
    "        2.0*x*y - 2*w*z, 1.0-2*x*x-2*z*z, 2.0*y*z+2*w*x,\n",
    "        2.0*x*z+2*w*y, 2.0*y*z-2*w*x, 1.0-2*x*x-2*y*y]\n",
    "    \n",
    "    return rot_matrix\n",
    "\n",
    "\n",
    "def inv_trans(p, pos, rot):\n",
    "    \"\"\"\n",
    "    Transform a marker into the rigid body coordinate system\n",
    "    \"\"\"\n",
    "    p = np.subtract(p, pos)\n",
    "    p = np.dot(p, rot.T)\n",
    "    p = [p[0], p[2], p[1]]\n",
    "    return p\n",
    "\n",
    "def trans(prow):\n",
    "    markers = [\"Wrist\", \"R_Shape_1\",\"R_Shape_2\",\"R_Shape_3\",\"R_Shape_4\",\n",
    "               \"Thumb_Fn\", \"Thumb_DIP\", \"Thumb_PIP\" ,\"Thumb_MCP\", \"Index_Fn\", \"Index_DIP\", \"Index_PIP\",\"Index_MCP\", \n",
    "               \"Middle_Fn\", \"Middle_DIP\", \"Middle_PIP\", \"Middle_MCP\", \"Ring_Fn\", \"Ring_DIP\", \"Ring_PIP\",\"Ring_MCP\", \n",
    "               \"Little_Fn\", \"Little_DIP\", \"Little_PIP\", \"Little_MCP\"]\n",
    "    row = prow.copy(deep=True)\n",
    "    \n",
    "    # iterate over all rows\n",
    "    for m in markers:\n",
    "        \n",
    "        x = row[m + \"_X\"]\n",
    "        y = row[m + \"_Y\"]\n",
    "        z = row[m + \"_Z\"]\n",
    "        \n",
    "        pivot = row['Pivot']\n",
    "        rot_matrix = row['RotationMatrix']\n",
    "        \n",
    "        if not math.isnan(x):\n",
    "            marker = [x, y, z]\n",
    "            marker = inv_trans(marker, pivot, rot_matrix)\n",
    "            \n",
    "            row[m + \"_X\"] = marker[0]\n",
    "            row[m + \"_Y\"] = marker[1]\n",
    "            row[m + \"_Z\"] = marker[2]\n",
    "            \n",
    "    return row\n",
    "\n",
    "def log(s):\n",
    "    with open(\"./out/status_PY01_preprocessing.txt\", \"a\") as myfile:\n",
    "        myfile.write(\"[\" + str(datetime.datetime.now()) + \"] \" + s + \"\\n\")\n",
    "    #print(\"[\" + str(datetime.datetime.now()) + \"] \" + s)\n",
    "\n",
    "\n",
    "test_participants = [1,2,9,15,19]\n",
    "\n",
    "def doJob(filename):\n",
    "    if filename.endswith(\".csv\"):\n",
    "        current_participant = filename.split(\".\")[0]\n",
    "        cp = int(current_participant.replace(\"P\",\"\").split(\"_\")[0])\n",
    "        if cp not in test_participants:\n",
    "            rigid_body_name = filename.split(\"_\")[1]\n",
    "\n",
    "            if os.path.isfile(\"Pickles/dfTransformed_\" + current_participant + \".pkl\") :\n",
    "                log(current_participant + \" is already available. Skipped.\")\n",
    "            else:\n",
    "                log(\"Reading \" + current_participant)\n",
    "                df, t_adjusted = read_take(\"raw-data/\" + filename)\n",
    "\n",
    "                df.Time = df.Time + t_adjusted\n",
    "                df.Time = df.Time * 1000 # sec to ms\n",
    "                df.Time = df.Time.astype(np.int64)\n",
    "                # Cells for quaternions\n",
    "                x = df[df.columns[2]]\n",
    "                y = df[df.columns[3]]\n",
    "                z = df[df.columns[4]]\n",
    "                w = df[df.columns[5]]\n",
    "\n",
    "                rot_matrix = np.array([\n",
    "                    [1-2*y*y-2*z*z, 2*x*y+2*w*z, 2*x*z - 2*w*y],\n",
    "                    [2*x*y - 2*w*z, 1-2*x*x-2*z*z, 2*y*z+2*w*x],\n",
    "                    [2*x*z+2*w*y, 2*y*z-2*w*x, 1-2*x*x-2*y*y]])\n",
    "\n",
    "                angle = np.degrees(np.arccos(np.dot((rot_matrix[:,1,:].T), [0, 1, 0])))\n",
    "                #adf = pd.DataFrame(angle)\n",
    "                #df2 = adf[np.logical_not(adf[0].isnull())]\n",
    "                #df3 = df2[df2[0] < df2[0].mean() + df2[0].std()*2]\n",
    "                dfFinal = df#.iloc[df3.index]\n",
    "\n",
    "                log(\"Creating RotationMatrix\")\n",
    "                dfFinal['RotationMatrix'] = dfFinal.apply(lambda row : quaternion_matrix(row, rigid_body_name), axis=1)\n",
    "                dfFinal.RotationMatrix = dfFinal.RotationMatrix.apply(lambda x : np.array(x).reshape(3,3))\n",
    "                dfFinal['Pivot'] = dfFinal.apply(lambda row : [row[rigid_body_name + '_X'], row[rigid_body_name + '_Y'], row[rigid_body_name + '_Z']], axis=1)\n",
    "                dfFinal['Pivot'] = dfFinal.Pivot.apply(lambda x : np.array(x))\n",
    "\n",
    "                log(\"Applying transformation\")\n",
    "                dfTransformed = dfFinal.apply(lambda x : trans(x), axis=1)\n",
    "                log(\"Saving to File\")\n",
    "                dfTransformed.to_pickle(\"Pickles/dfTransformed_\" + current_participant + \".pkl\")        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HINT: This may take some hours. Detailed progress see output file: out/status_PY01_preprocessing.txt\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 130/130 [25:10<00:00, 11.62s/it] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 5.4 s, sys: 5.61 s, total: 11 s\n",
      "Wall time: 25min 11s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print(\"HINT: This may take some hours. Detailed progress see output file: out/status_PY01_preprocessing.txt\")\n",
    "pool = multiprocessing.Pool(max(multiprocessing.cpu_count()-1, 1))\n",
    "files = list(os.listdir(\"raw-data/\"))\n",
    "for _ in tqdm.tqdm(pool.imap_unordered(doJob, files), total=len(files)):\n",
    "    pass"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}