pipeline_demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import pickle\n",
    "import glob\n",
    "import cv2\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import clear_output\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "video_dir = '/mnt/disks/datastorage/videos/'\n",
    "weights_dir = '/mnt/disks/datastorage/weights/'\n",
    "\n",
    "# Video stream\n",
    "video_file = 'cold_day/demos/shooting_ballroom_light6.mov'\n",
    "video_path = os.path.join(video_dir, 'fens', video_file)\n",
    "\n",
    "# For live stream\n",
    "# video_path = './stream.sdp'\n",
    "\n",
    "# Parameters\n",
    "nb_classes = 3\n",
    "skip_num = 1\n",
    "svm_model = os.path.join(weights_dir, 'svm_demo.pkl')\n",
    "prediction_path = '/mnt/disks/datastorage/predictions/'\n",
    "flownet_model = 'FlowNet2CSS'\n",
    "\n",
    "# Models\n",
    "flownet_weights = os.path.join(weights_dir, 'FlowNet2-CSS.pth.tar')\n",
    "spatial_weights = os.path.join(weights_dir, 'spatial_weights_demo.pth.tar')\n",
    "motion_weights = os.path.join(weights_dir, 'motion_weights_demo.pth.tar')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run pipeline and ignore warnings\n",
    "!python3 -W ignore pipeline.py --stream $video_path \\\n",
    "                               --nb_classes $nb_classes \\\n",
    "                               --skip_frames $skip_num \\\n",
    "                               --svm $svm_model \\\n",
    "                               --save $prediction_path \\\n",
    "                               --model $flownet_model \\\n",
    "                               --optical_weights $flownet_weights \\\n",
    "                               --spatial_weights $spatial_weights \\\n",
    "                               --motion_weights $motion_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction_files = sorted(glob.glob(os.path.join(prediction_path, 'predictions_*.pkl')), reverse=True)\n",
    "\n",
    "with open(prediction_files[0], 'rb') as pf:\n",
    "    predictions = pickle.load(pf)\n",
    "    predictions = np.array(predictions).squeeze()\n",
    "    actions = np.argmax(predictions, axis=1)\n",
    "    \n",
    "predictions_argmax = predictions.argmax(1)\n",
    "for pred in list(zip(predictions, predictions_argmax)):\n",
    "    print(pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Interpolate predictions for skipped frames\n",
    "predictions_extended = np.zeros((skip_num + 1) * len(predictions_argmax))\n",
    "for i in range(len(predictions_extended)):\n",
    "    predictions_extended[i] = int(predictions_argmax[i // (skip_num + 1)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "action_type = 'demo'\n",
    "action_dict = {\n",
    "    'fens': {\n",
    "        0: 'Theft',\n",
    "        1: 'Assault',\n",
    "        2: 'Shooting',\n",
    "        3: 'No Action',\n",
    "    },\n",
    "    'demo': {\n",
    "        0: 'Assault',\n",
    "        1: 'Shooting',\n",
    "        2: 'No Action',\n",
    "    }\n",
    "}\n",
    "\n",
    "cap = cv2.VideoCapture(video_path)\n",
    "\n",
    "frame_time = 0\n",
    "frame_idx = 0\n",
    "display = None\n",
    "i = 0\n",
    "while True:\n",
    "    start_time = time.time()\n",
    "    \n",
    "    ret, frame = cap.read()\n",
    "    if not ret:\n",
    "        break\n",
    "    \n",
    "    try:\n",
    "        label = action_dict[action_type][actions[frame_idx]]\n",
    "    except IndexError:\n",
    "        break\n",
    "        \n",
    "    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
    "    cv2.putText(frame, label, (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), thickness=3)\n",
    "    \n",
    "    plt.axis('off')\n",
    "    plt.imshow(frame)\n",
    "    plt.show()\n",
    "    clear_output(wait=True)\n",
    "\n",
    "    frame_time += time.time() - start_time\n",
    "    frame_idx += 1\n",
    "    \n",
    "cap.release()\n",
    "\n",
    "frame_time /= frame_idx\n",
    "print('Avg. frame processing time: {} s'.format(round(frame_time, 2)))"
   ]
  }
 ],
 "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.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}