docker_get_points_from_image.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import cv2\n",
    "import os\n",
    "import pyopenpose as op\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = (20,10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_dir = \"../rachel_wong/example/Videos\"\n",
    "output_dir = \"../rachel_wong/example/Videos_Pose\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "pose_dir = os.path.join(output_dir, \"Pose\")\n",
    "black_pose_dir = os.path.join(output_dir, \"Black_Pose\")\n",
    "coords_dir = os.path.join(output_dir, \"Coords\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#os.walk recursively goes through all the files in our input_dir\n",
    "input_paths = []\n",
    "for directory, folders, files in os.walk(input_dir):\n",
    "    for file in files:\n",
    "        input_paths.append(os.path.join(directory, file))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "found 26 input files\n",
      "example: ../rachel_wong/example/Videos/Hair Pulling/IMG_8455.MOV\n"
     ]
    }
   ],
   "source": [
    "print(\"found\", len(input_paths), \"input files\")\n",
    "print(\"example:\", input_paths[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Custom Params (refer to include/openpose/flags.hpp for more parameters)\n",
    "params = dict()\n",
    "params[\"model_folder\"] = \"/openpose/models/\"\n",
    "params[\"face\"] = True\n",
    "params[\"hand\"] = True\n",
    "params[\"render_pose\"] = 0 #we will manually draw pose, so we turn this off"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Starting OpenPose\n",
    "opWrapper = op.WrapperPython()\n",
    "opWrapper.configure(params)\n",
    "opWrapper.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def do_openpose(image):\n",
    "    datum = op.Datum()\n",
    "    datum.cvInputData = image\n",
    "    opWrapper.emplaceAndPop([datum])\n",
    "    return datum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_empty(keypoint):\n",
    "    return keypoint.sum() == 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_point(image, x, y, color):\n",
    "    radius = int(min(image.shape[0], image.shape[1]) * .005)\n",
    "    cv2.circle(image, (x, y), radius, color, thickness=radius, lineType=8, shift=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_line(image, x1, y1, x2, y2, color):\n",
    "    thickness = max(2, int(min(image.shape[0], image.shape[1]) / 200.0))\n",
    "    cv2.line(image, (x1, y1), (x2, y2), color, thickness)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "CocoColors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0],\n",
    "              [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255],\n",
    "              [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]\n",
    "\n",
    "Body25Pairs = [\n",
    "    (1,8), (1,2), (1,5), (2,3), (3,4), (5,6), (6,7), (8,9), (9,10), (10,11), (8,12), \n",
    "    (12,13), (13,14), (1,0), (0,15), (15,17), (0,16), (16,18), (14,19), (19,20), \n",
    "    (14,21), (11,22), (22,23), (11,24)\n",
    "]\n",
    "\n",
    "HandPairs = [\n",
    "    (0,1), (1,2), (2,3), (3,4), (0,5), (5,6), (6,7), (7,8), (0,9), (9,10), (10,11), \n",
    "    (11,12), (0,13), (13,14), (14,15), (15,16), (0,17), (17,18), (18,19), (19,20)\n",
    "]\n",
    "\n",
    "FacePairs = [\n",
    "    (0,1), (1,2), (2,3), (3,4), (4,5), (5,6), (6,7), (7,8), (8,9), (9,10), (10,11),\n",
    "    (11,12), (12,13), (13,14), (14,15), (15,16), (17,18), (18,19), (19,20), (20,21), \n",
    "    (22,23), (23,24), (24,25), (25,26), (27,28), (28,29), (29,30), (31,32), (32,33), \n",
    "    (33,34), (34,35), (36,37), (37,38), (38,39), (39,40), (40,41), (41,36), (42,43), \n",
    "    (43,44), (44,45), (45,46), (46,47), (47,42), (48,49), (49,50), (50,51), (51,52), \n",
    "    (52,53), (53,54), (54,55), (55,56), (56,57), (57,58), (58,59), (59,48), (60,61), \n",
    "    (61,62), (62,63), (63,64), (64,65), (65,66), (66,67), (67,60)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_color(i, mode):\n",
    "    if mode == \"face\":\n",
    "        return [255, 255, 255] #white\n",
    "    \n",
    "    return CocoColors[i % len(CocoColors)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "pairs_dict = {\n",
    "    \"body\": Body25Pairs,\n",
    "    \"hand\": HandPairs,\n",
    "    \"face\": FacePairs\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "def array_empty(array):\n",
    "    return len(array.shape) == 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_keypoints(image, keypoints, mode=\"body\"):\n",
    "    pairs = pairs_dict[mode]\n",
    "    \n",
    "    if array_empty(keypoints):\n",
    "        return image\n",
    "    \n",
    "    for human in keypoints:\n",
    "        for i, keypoint in enumerate(human):\n",
    "            if not is_empty(keypoint):\n",
    "                point_color = get_color(i, mode)\n",
    "                draw_point(image, int(keypoint[0].round()), int(keypoint[1].round()), point_color)\n",
    "\n",
    "        for i, pair in enumerate(pairs):\n",
    "            keypoint0 = human[pair[0]]\n",
    "            keypoint1 = human[pair[1]]\n",
    "            if not is_empty(keypoint0) and not is_empty(keypoint1):\n",
    "                line_color = get_color(i, mode)\n",
    "                draw_line(image, int(keypoint0[0].round()), int(keypoint0[1].round()), int(keypoint1[0].round()), int(keypoint1[1].round()), line_color)\n",
    "\n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_keypoints_from_result(image, result):\n",
    "    image_pose = draw_keypoints(image, result.poseKeypoints)\n",
    "    image_pose = draw_keypoints(image_pose, result.faceKeypoints, mode=\"face\")\n",
    "    image_pose = draw_keypoints(image_pose, result.handKeypoints[0], mode=\"hand\")\n",
    "    return draw_keypoints(image_pose, result.handKeypoints[1], mode=\"hand\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_black_image(image):\n",
    "    image = np.copy(image)\n",
    "    image[:,:,:] = 0\n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# '\\r' is a \"carriage return\"\n",
    "# it moves the cursor back to the biginning of the line\n",
    "# the next line will then overwrite the previus one\n",
    "def log_video_progress(video):\n",
    "    frame_idx = video.get(cv2.CAP_PROP_POS_FRAMES)\n",
    "    sys.stdout.write(str(int(frame_idx)) + \" frames\\r\")\n",
    "    sys.stdout.flush()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_image(image):\n",
    "    plt.imshow(image.swapaxes(0, 1)[:,:,[2,1,0]])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def decode_fourcc(codec):\n",
    "    return ''.join([\n",
    "        chr(codec & 255),\n",
    "        chr((codec >> 8) & 255),\n",
    "        chr((codec >> 16) & 255),\n",
    "        chr((codec >> 24) & 255)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def move_path(path, folder_from, folder_to):\n",
    "    return path.replace(folder_from, folder_to, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dirs(path):\n",
    "    directory = os.path.dirname(path)\n",
    "    if not os.path.exists(directory):\n",
    "        os.makedirs(directory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_coords(path, result):\n",
    "    pose = np.empty(shape=(0,25,3)) if array_empty(result.poseKeypoints) else result.poseKeypoints\n",
    "    face = np.empty(shape=(0,70,3)) if array_empty(result.faceKeypoints) else result.faceKeypoints\n",
    "    handl = np.empty(shape=(0,21,3)) if array_empty(result.handKeypoints[0]) else result.handKeypoints[0]\n",
    "    handr = np.empty(shape=(0,21,3)) if array_empty(result.handKeypoints[1]) else result.handKeypoints[1]\n",
    "    \n",
    "    np.savez(path, pose=pose, face=face, handl=handl, handr=handr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1255 frames\r"
     ]
    }
   ],
   "source": [
    "for input_path in input_paths:\n",
    "    video = cv2.VideoCapture(input_path)\n",
    "    \n",
    "    framerate = video.get(cv2.CAP_PROP_FPS)\n",
    "    codec = int(video.get(cv2.CAP_PROP_FOURCC))\n",
    "    width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "    height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "    \n",
    "    pose_path = move_path(input_path, input_dir, pose_dir)\n",
    "    create_dirs(pose_path)\n",
    "    video_pose = cv2.VideoWriter(pose_path, cv2.VideoWriter_fourcc(*'MPEG'), framerate, (width, height))\n",
    "    \n",
    "    black_pose_path = move_path(input_path, input_dir, black_pose_dir)\n",
    "    create_dirs(black_pose_path)\n",
    "    video_black_pose = cv2.VideoWriter(black_pose_path, cv2.VideoWriter_fourcc(*'MPEG'), framerate, (width, height))\n",
    "    \n",
    "    coords_path = move_path(input_path, input_dir, coords_dir)\n",
    "    coords_path, _ext = os.path.splitext(coords_path)\n",
    "    create_dirs(os.path.join(coords_path, \"test\"))\n",
    "    \n",
    "    frames_remaining, frame = video.read()\n",
    "    frame_idx = 0\n",
    "    while frames_remaining:\n",
    "        result = do_openpose(frame)\n",
    "\n",
    "        image_pose = draw_keypoints_from_result(frame, result)\n",
    "        video_pose.write(image_pose)\n",
    "\n",
    "        black = get_black_image(frame)\n",
    "        black_pose = draw_keypoints_from_result(black, result)\n",
    "        video_black_pose.write(black_pose)\n",
    "\n",
    "        coords_frame_path = os.path.join(coords_path, str(frame_idx) + \".npz\")\n",
    "        save_coords(coords_frame_path, result)\n",
    "\n",
    "        log_video_progress(video)\n",
    "        frames_remaining, frame = video.read()\n",
    "        frame_idx += 1\n",
    "        \n",
    "    print(input_path)\n",
    "    video.release()\n",
    "    video_pose.release()\n",
    "    video_black_pose.release()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Extra Code Below"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![(no touchy!)](notouchy.jpg \"(No Touchy!)\")"
   ]
  },
  {
   "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": []
  },
  {
   "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": []
  },
  {
   "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": "markdown",
   "metadata": {},
   "source": [
    "![(I said no touchy!)](notouchy2.jpg \"(I Said No Touchy!)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "../rachel_wong/example/Videos/Hair Pulling/IMG_8455.MOV\n",
      "None\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "a bytes-like object is required, not 'NoneType'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0mTraceback (most recent call last)",
      "\u001b[0;32m<ipython-input-55-57de2f21ddd3>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0mfile\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcv2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mimread\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_path\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfile\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 7\u001b[0;31m     \u001b[0mdatum\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcvInputData\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfile\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      8\u001b[0m     \u001b[0mopWrapper\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0memplaceAndPop\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mdatum\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdatum\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: a bytes-like object is required, not 'NoneType'"
     ]
    }
   ],
   "source": [
    "# Process Files\n",
    "for input_path in input_paths:\n",
    "    datum = op.Datum()\n",
    "    print(input_path)\n",
    "    file = cv2.imread(input_path)\n",
    "    print(file)\n",
    "    datum.cvInputData = file\n",
    "    opWrapper.emplaceAndPop([datum])\n",
    "    print(datum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Body keypoints: \n",
      "8.9500005e-37\n",
      "Face keypoints: \n",
      "8.9500005e-37\n",
      "Left hand keypoints: \n",
      "8.9500005e-37\n",
      "Right hand keypoints: \n",
      "1.0406712e-36\n",
      "1.0406712e-36\n"
     ]
    }
   ],
   "source": [
    "# Display Image\n",
    "print(\"Body keypoints: \\n\" + str(datum.poseKeypoints))\n",
    "print(\"Face keypoints: \\n\" + str(datum.faceKeypoints))\n",
    "print(\"Left hand keypoints: \\n\" + str(datum.handKeypoints[0]))\n",
    "print(\"Right hand keypoints: \\n\" + str(datum.handKeypoints[1]))\n",
    "print(datum.outputData)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pims"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = pims.Video(\"../rachel_wong/example/Videos/Hair Pulling/IMG_8455.MOV\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sliced PyAVReaderTimed. Original repr:\n",
       "    <Frames>\n",
       "    Format: HEVC (High Efficiency Video Coding)\n",
       "    Source: ../rachel_wong/example/Videos/Hair Pulling/IMG_8455.MOV\n",
       "    Duration: 50.278 seconds\n",
       "    Frame rate: 29.973 fps\n",
       "    Length: 1507 frames\n",
       "    Frame Shape: (1080, 1920, 3)\n",
       "    "
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v[2:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "relative_path = os.path.relpath(full_path, input_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plt.imshow(np.flip(np.flip(datum.cvOutputData.swapaxes(0, 1), axis=2), axis=1))\n",
    "        "
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}