data_preprocessing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "import json\n",
    "import math\n",
    "import os\n",
    "import pickle\n",
    "import random\n",
    "\n",
    "import joblib\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from utils import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read Data from CSV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Video Name  Label\n",
      "0         v1      1\n",
      "1         v2      1\n",
      "2         v3      3\n",
      "3         v4      0\n",
      "4         v5      0\n",
      "5         v6      1\n",
      "6         v7      0\n",
      "7         v8      3\n",
      "8         v9      0\n",
      "9        v10      2\n"
     ]
    }
   ],
   "source": [
    "# Set path to joint data, label_file, and save path\n",
    "joints_path = '/home/ubuntu/EPD/data/all_gait_clips/' # location of 3d joints \n",
    "exam_df = pd.read_csv('data/gait_labels.csv') # label file with example format\n",
    "save_dir = '/home/ubuntu/EPD/splits/default/'\n",
    "print(exam_df.head(10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# list of valid video names that we can use for gait\n",
    "video_names = exam_df['Video Name'].tolist()\n",
    "labels = exam_df['Label'].tolist()\n",
    "    \n",
    "# generate a dictionary mapping video name to label\n",
    "labels_dict = {}\n",
    "for i in range(len(video_names)):\n",
    "    labels_dict[video_names[i]] = labels[i]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Track subject and read poses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "processing video name: v1\n",
      "processing video name: v2\n",
      "processing video name: v3\n",
      "processing video name: v4\n",
      "processing video name: v5\n",
      "processing video name: v6\n",
      "processing video name: v7\n",
      "processing video name: v8\n",
      "processing video name: v9\n",
      "processing video name: v10\n",
      "processing video name: v11\n",
      "processing video name: v12\n",
      "processing video name: v13\n",
      "processing video name: v14\n",
      "processing video name: v15\n",
      "processing video name: v16\n",
      "processing video name: v17\n",
      "processing video name: v18\n",
      "processing video name: v19\n",
      "processing video name: v20\n"
     ]
    }
   ],
   "source": [
    "def read_keypoints(keypoint_path):\n",
    "    \"\"\"\n",
    "    Read json files in given directory into arrays of pose keypoints. \n",
    "    Remove confidence scores and format pose keypoints as a list of tuples, as the preprocessor expects. \n",
    "    :param keypoint_path: path to directory of keypoints\n",
    "    :return: dictionary with <key=video name, value=keypoints>\n",
    "    \"\"\"\n",
    "    pose_dict = {}\n",
    "    for video_name in video_names:\n",
    "        vibe_output = joblib.load(keypoint_path + video_name + '/vibe_output.pkl') \n",
    "        print('processing video name:', video_name)\n",
    "\n",
    "        # Choose the subject with pose detected for the most frames. Modify as necessary. \n",
    "        max_key = list(vibe_output)[0]\n",
    "        max_frames = len(vibe_output[max_key]['joints3d'])\n",
    "        for key in vibe_output:\n",
    "            num_frames = len(vibe_output[key]['joints3d'])\n",
    "            if num_frames > max_frames:\n",
    "                max_frames = num_frames\n",
    "                max_key = key\n",
    "        joints3d = vibe_output[max_key]['joints3d']\n",
    "\n",
    "        # normalize each point by the a single joint\n",
    "        for i in range(len(joints3d)):\n",
    "            pelvis = 48\n",
    "            joints3d[i] = joints3d[i] - joints3d[i][pelvis]\n",
    "        pose_dict[video_name] = np.stack(joints3d)\n",
    "    return pose_dict\n",
    "\n",
    "\n",
    "\n",
    "def normalize_poses(pose_dict):\n",
    "    \"\"\"\n",
    "    Normalize each pose along each axis by video. Divide by the largest value in each direction\n",
    "    and center around the origin.\n",
    "    :param pose_dict: dictionary of poses\n",
    "    :return: dictionary of normalized poses\n",
    "    \"\"\"\n",
    "    normalized_pose_dict = {}\n",
    "    for video_name in pose_dict:\n",
    "        poses = pose_dict[video_name].copy()\n",
    "\n",
    "        maxes = [-1, -1, -1]\n",
    "        mins = [1, 1, 1]\n",
    "\n",
    "        for i in range(len(poses)):\n",
    "            pose = poses[i]\n",
    "            for j in range(49):\n",
    "                [x, y, z] = pose[j]\n",
    "                maxes[0] = max(maxes[0], x)\n",
    "                maxes[1] = max(maxes[1], y)\n",
    "                maxes[2] = max(maxes[2], z)\n",
    "                mins[0] = min(mins[0], x)\n",
    "                mins[1] = min(mins[1], y)\n",
    "                mins[2] = min(mins[2], z)\n",
    "        for i in range(len(poses)):\n",
    "            pose = poses[i]\n",
    "            for j in range(49):\n",
    "                [x, y, z] = pose[j]\n",
    "                poses[i][j][0] = x / (maxes[0] - mins[0])\n",
    "                poses[i][j][1] = y / (maxes[1] - mins[1])\n",
    "                poses[i][j][2] = z / (maxes[2] - mins[2])   \n",
    "        normalized_pose_dict[video_name] = poses\n",
    "    return normalized_pose_dict\n",
    "\n",
    "pose_dict = read_keypoints(joints_path)\n",
    "normalized_pose_dict = normalize_poses(pose_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Partition gait videos into clips of specified length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_clips(video, length, offset):\n",
    "    \"\"\"\n",
    "    Returns a list of partitioned gait segments of given length in frames and offset to next clip\n",
    "    :param video: input video\n",
    "    :param length: length of clip\n",
    "    :param offset: offset between clips\n",
    "    :return: partition of gait segments into clips\n",
    "    \"\"\"\n",
    "    clips = []\n",
    "    n = len(video)\n",
    "    num_clips = math.ceil((n - length+1) / offset)\n",
    "\n",
    "    i = n\n",
    "    while i > length:\n",
    "        temp_end = i\n",
    "        temp_start = temp_end - length\n",
    "        clip = np.array(video[temp_start:temp_end])\n",
    "        clips.append(clip)\n",
    "        i -= offset\n",
    "    return clips\n",
    "\n",
    "def partition_videos(video_names, pose_dict, length=200, offset=50):\n",
    "    \"\"\"\n",
    "    Partition poses from each video into clips.\n",
    "    :param video_names: names of videos to partition\n",
    "    :param pose_dict: dictionary of poses for each video\n",
    "    :param length: length of clip\n",
    "    :param offset: offset between clips\n",
    "    :return: dictionary of clips for each video\n",
    "    \"\"\"\n",
    "    clip_dict = {}\n",
    "    for name in video_names:\n",
    "        print(name)\n",
    "        clips = get_clips(pose_dict[name], length, offset)\n",
    "        clip_dict[name] = clips\n",
    "    return clip_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "v1\n",
      "v2\n",
      "v3\n",
      "v4\n",
      "v5\n",
      "v6\n",
      "v7\n",
      "v8\n",
      "v9\n",
      "v10\n",
      "v11\n",
      "v12\n",
      "v13\n",
      "v14\n",
      "v15\n",
      "v16\n",
      "v17\n",
      "v18\n",
      "v19\n",
      "v20\n"
     ]
    }
   ],
   "source": [
    "clip_dict = partition_videos(normalized_pose_dict.keys(), normalized_pose_dict, length=100, offset=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate leave-one-out cross validation splits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['v1', 'v2', 'v3', 'v4', 'v5', 'v6', 'v7', 'v8', 'v9', 'v10', 'v11', 'v12', 'v13', 'v14', 'v15', 'v16', 'v17', 'v18', 'v19', 'v20']\n",
      "len train 329\n",
      "test: ['v1']\n",
      "train: 329\n",
      "labels: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]\n",
      "len train 348\n",
      "test: ['v2']\n",
      "train: 348\n",
      "labels: [1, 1, 1, 1]\n",
      "len train 334\n",
      "test: ['v3']\n",
      "train: 334\n",
      "labels: [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]\n",
      "len train 345\n",
      "test: ['v4']\n",
      "train: 345\n",
      "labels: [0, 0, 0, 0, 0, 0, 0]\n",
      "len train 331\n",
      "test: ['v5']\n",
      "train: 331\n",
      "labels: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "len train 344\n",
      "test: ['v6']\n",
      "train: 344\n",
      "labels: [1, 1, 1, 1, 1, 1, 1, 1]\n",
      "len train 336\n",
      "test: ['v7']\n",
      "train: 336\n",
      "labels: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "len train 330\n",
      "test: ['v8']\n",
      "train: 330\n",
      "labels: [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]\n",
      "len train 334\n",
      "test: ['v9']\n",
      "train: 334\n",
      "labels: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "len train 324\n",
      "test: ['v10']\n",
      "train: 324\n",
      "labels: [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]\n",
      "len train 331\n",
      "test: ['v11']\n",
      "train: 331\n",
      "labels: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]\n",
      "len train 335\n",
      "test: ['v12']\n",
      "train: 335\n",
      "labels: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "len train 330\n",
      "test: ['v13']\n",
      "train: 330\n",
      "labels: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "len train 346\n",
      "test: ['v14']\n",
      "train: 346\n",
      "labels: [1, 1, 1, 1, 1, 1]\n",
      "len train 350\n",
      "test: ['v15']\n",
      "train: 350\n",
      "labels: [1, 1]\n",
      "len train 319\n",
      "test: ['v16']\n",
      "train: 319\n",
      "labels: [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]\n",
      "len train 327\n",
      "test: ['v17']\n",
      "train: 327\n",
      "labels: [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]\n",
      "len train 331\n",
      "test: ['v18']\n",
      "train: 331\n",
      "labels: [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]\n",
      "len train 332\n",
      "test: ['v19']\n",
      "train: 332\n",
      "labels: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]\n",
      "len train 332\n",
      "test: ['v20']\n",
      "train: 332\n",
      "labels: [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]\n"
     ]
    }
   ],
   "source": [
    "def generate_pose_label(clip_dict, test_clip_dict, train_list, val_list, test_list):\n",
    "    train = {}\n",
    "    train['pose'] = []\n",
    "    train['label'] = []\n",
    "    val = {}\n",
    "    val['pose'] = []\n",
    "    val['label'] = []\n",
    "    test = {}\n",
    "    test['pose'] = []\n",
    "    test['label'] = []\n",
    "    complete_list = train_list + val_list + test_list\n",
    "    \n",
    "    # Place each clip in the correct split\n",
    "    for video_name in train_list:\n",
    "        clips = test_clip_dict[video_name]\n",
    "        for clip in clips:\n",
    "            train['label'].append(labels_dict[video_name])\n",
    "            train['pose'].append(clip)\n",
    "    for video_name in val_list + test_list:\n",
    "        clips = test_clip_dict[video_name]\n",
    "        for clip in clips:\n",
    "            if video_name in val_list:\n",
    "                val['label'].append(labels_dict[video_name])\n",
    "                val['pose'].append(clip)\n",
    "            elif video_name in test_list:\n",
    "                test['label'].append(labels_dict[video_name])\n",
    "                test['pose'].append(clip)\n",
    "    print(\"len train\", len(train['label']))\n",
    "    return train, val, test\n",
    "\n",
    "def generate_leave_one_out_folds(clip_dict, test_clip_dict, save_dir, seed=None):\n",
    "    \"\"\"\n",
    "    Generate folds for leave-one-out CV.\n",
    "    :param clip_dict: dictionary of clips for each video\n",
    "    :param test_clip_dict: dictionary of poses for each test video\n",
    "    :param save_dir: save directory for folds\n",
    "    \"\"\"\n",
    "    if not os.path.exists(save_dir):\n",
    "        os.mkdir(save_dir)\n",
    "        \n",
    "    video_names_list = []\n",
    "    k = 0\n",
    "    for video_name in clip_dict:\n",
    "        video_names_list.append(video_name)\n",
    "        k += 1\n",
    "    \n",
    "    print(video_names_list)\n",
    "    for j in range(len(video_names)):\n",
    "        i = j + 1\n",
    "        train_list = video_names_list[:]\n",
    "        train_list.remove(video_names[j])\n",
    "        val_list = []\n",
    "        test_list = [video_names[j]]\n",
    "                    \n",
    "        train, _, test = generate_pose_label(clip_dict, test_clip_dict, train_list, val_list, test_list)\n",
    "        print(\"test:\", test_list)\n",
    "        print(\"train:\", len(train['label']))\n",
    "        print(\"labels:\", test['label'])\n",
    "        pickle.dump(train_list, open(save_dir+\"EPG_train_list_\"+str(i)+\".pkl\", \"wb\"))\n",
    "        pickle.dump(test_list, open(save_dir+\"EPG_test_list_\"+str(i)+\".pkl\", \"wb\"))\n",
    "        pickle.dump(train, open(save_dir+\"EPG_train_\"+str(i)+\".pkl\", \"wb\"))\n",
    "        pickle.dump(test, open(save_dir+\"EPG_test_\"+str(i)+\".pkl\", \"wb\"))\n",
    "    pickle.dump(labels_dict, open(save_dir+\"EPG_labels.pkl\", \"wb\"))\n",
    "    \n",
    "generate_leave_one_out_folds(clip_dict, clip_dict, save_dir, seed=4096)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}