Create_Python_Script.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step1: Create the Python Script\n",
    "\n",
    "In the cell below, you will need to complete the Python script and run the cell to generate the file using the magic `%%writefile` command. Your main task is to complete the following methods for the `PersonDetect` class:\n",
    "* `load_model`\n",
    "* `predict`\n",
    "* `draw_outputs`\n",
    "* `preprocess_outputs`\n",
    "* `preprocess_inputs`\n",
    "\n",
    "For your reference, here are all the arguments used for the argument parser in the command line:\n",
    "* `--model`:  The file path of the pre-trained IR model, which has been pre-processed using the model optimizer. There is automated support built in this argument to support both FP32 and FP16 models targeting different hardware.\n",
    "* `--device`: The type of hardware you want to load the model on (CPU, GPU, MYRIAD, HETERO:FPGA,CPU)\n",
    "* `--video`: The file path of the input video.\n",
    "* `--output_path`: The location where the output stats and video file with inference needs to be stored (results/[device]).\n",
    "* `--max_people`: The max number of people in queue before directing a person to another queue.\n",
    "* `--threshold`: The probability threshold value for the person detection. Optional arg; default value is 0.60."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing person_detect.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile person_detect.py\n",
    "\n",
    "import numpy as np\n",
    "import time\n",
    "from openvino.inference_engine import IENetwork, IECore\n",
    "import os\n",
    "import cv2\n",
    "import argparse\n",
    "import sys\n",
    "\n",
    "\n",
    "class Queue:\n",
    "    '''\n",
    "    Class for dealing with queues\n",
    "    '''\n",
    "    def __init__(self):\n",
    "        self.queues=[]\n",
    "\n",
    "    def add_queue(self, points):\n",
    "        self.queues.append(points)\n",
    "\n",
    "    def get_queues(self, image):\n",
    "        for q in self.queues:\n",
    "            x_min, y_min, x_max, y_max=q\n",
    "            frame=image[y_min:y_max, x_min:x_max]\n",
    "            yield frame\n",
    "    \n",
    "    def check_coords(self, coords):\n",
    "        d={k+1:0 for k in range(len(self.queues))}\n",
    "        for coord in coords:\n",
    "            for i, q in enumerate(self.queues):\n",
    "                if coord[0]>q[0] and coord[2]<q[2]:\n",
    "                    d[i+1]+=1\n",
    "        return d\n",
    "\n",
    "\n",
    "class PersonDetect:\n",
    "    '''\n",
    "    Class for the Person Detection Model.\n",
    "    '''\n",
    "\n",
    "    def __init__(self, model_name, device, threshold=0.60):\n",
    "        self.model_weights=model_name+'.bin'\n",
    "        self.model_structure=model_name+'.xml'\n",
    "        self.device=device\n",
    "        self.threshold=threshold\n",
    "\n",
    "        try:\n",
    "            self.core = IECore()\n",
    "            self.model = self.core.read_network(self.model_structure,self.model_weights)\n",
    "#             self.model=IENetwork(self.model_structure, self.model_weights)\n",
    "        except Exception as e:\n",
    "            raise ValueError(\"Could not Initialise the network. Have you enterred the correct model path?\")\n",
    "\n",
    "        self.input_name=next(iter(self.model.inputs))\n",
    "        self.input_shape=self.model.inputs[self.input_name].shape\n",
    "        self.output_name=next(iter(self.model.outputs))\n",
    "        self.output_shape=self.model.outputs[self.output_name].shape\n",
    "\n",
    "    def load_model(self):\n",
    "        return self.core.load_network(self.model, self.device)\n",
    "        \n",
    "    def predict(self, image):\n",
    "        processed_frame = preprocess_input(frame)\n",
    "        self.core.exec_net(processed_frame)\n",
    "        if self.core.wait() == 0:\n",
    "            result = {}\n",
    "            for i in list(self.model.outputs.keys()):\n",
    "                result[i] = self.core.requests[0].outputs[i]\n",
    "            \n",
    "            outputs = preprocess_outputs(result)\n",
    "            image = draw_outputs(outputs, image)\n",
    "            \n",
    "            return outputs, image\n",
    "\n",
    "        raise NotImplementedError\n",
    "    \n",
    "    def draw_outputs(self, coords, image):\n",
    "        for detection in outputs:\n",
    "            image_id, label, conf, x_min, y_min, x_max, y_max = detection\n",
    "            if conf > self.threshold:\n",
    "                x_min = int(x_min * width)\n",
    "                x_max = int(x_max * width)\n",
    "                y_min = int(y_min * height)\n",
    "                y_max = int(y_max * height)\n",
    "                \n",
    "                cv2.rectangle(image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)\n",
    "                \n",
    "        return image\n",
    "\n",
    "    def preprocess_outputs(self, outputs):\n",
    "        output = result['DetectionOutput']\n",
    "        return output[0][0]\n",
    "\n",
    "    def preprocess_input(self, image):\n",
    "        p_frame = cv2.resize(image, (self.input_shape[3], self.input_shape[2]))\n",
    "        p_frame = p_frame.transpose(2, 0, 1)\n",
    "        p_frame = p_frame.reshape(1, *p_frame.shape)\n",
    "        return p_frame\n",
    "\n",
    "\n",
    "def main(args):\n",
    "    model=args.model\n",
    "    device=args.device\n",
    "    video_file=args.video\n",
    "    max_people=args.max_people\n",
    "    threshold=args.threshold\n",
    "    output_path=args.output_path\n",
    "\n",
    "    start_model_load_time=time.time()\n",
    "    pd= PersonDetect(model, device, threshold)\n",
    "    pd.load_model()\n",
    "    total_model_load_time = time.time() - start_model_load_time\n",
    "\n",
    "    queue=Queue()\n",
    "    \n",
    "    try:\n",
    "        queue_param=np.load(args.queue_param)\n",
    "        for q in queue_param:\n",
    "            queue.add_queue(q)\n",
    "    except:\n",
    "        print(\"error loading queue param file\")\n",
    "\n",
    "    try:\n",
    "        cap=cv2.VideoCapture(video_file)\n",
    "    except FileNotFoundError:\n",
    "        print(\"Cannot locate video file: \"+ video_file)\n",
    "    except Exception as e:\n",
    "        print(\"Something else went wrong with the video file: \", e)\n",
    "    \n",
    "    initial_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "    initial_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "    video_len = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "    fps = int(cap.get(cv2.CAP_PROP_FPS))\n",
    "    out_video = cv2.VideoWriter(os.path.join(output_path, 'output_video.mp4'), cv2.VideoWriter_fourcc(*'avc1'), fps, (initial_w, initial_h), True)\n",
    "    \n",
    "    counter=0\n",
    "    start_inference_time=time.time()\n",
    "\n",
    "    try:\n",
    "        while cap.isOpened():\n",
    "            ret, frame=cap.read()\n",
    "            if not ret:\n",
    "                break\n",
    "            counter+=1\n",
    "            \n",
    "            coords, image= pd.predict(frame)\n",
    "            num_people= queue.check_coords(coords)\n",
    "            print(f\"Total People in frame = {len(coords)}\")\n",
    "            print(f\"Number of people in queue = {num_people}\")\n",
    "            out_text=\"\"\n",
    "            y_pixel=25\n",
    "            \n",
    "            for k, v in num_people.items():\n",
    "                out_text += f\"No. of People in Queue {k} is {v} \"\n",
    "                if v >= int(max_people):\n",
    "                    out_text += f\" Queue full; Please move to next Queue \"\n",
    "                cv2.putText(image, out_text, (15, y_pixel), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0), 2)\n",
    "                out_text=\"\"\n",
    "                y_pixel+=40\n",
    "            out_video.write(image)\n",
    "            \n",
    "        total_time=time.time()-start_inference_time\n",
    "        total_inference_time=round(total_time, 1)\n",
    "        fps=counter/total_inference_time\n",
    "\n",
    "        with open(os.path.join(output_path, 'stats.txt'), 'w') as f:\n",
    "            f.write(str(total_inference_time)+'\\n')\n",
    "            f.write(str(fps)+'\\n')\n",
    "            f.write(str(total_model_load_time)+'\\n')\n",
    "\n",
    "        cap.release()\n",
    "        cv2.destroyAllWindows()\n",
    "    except Exception as e:\n",
    "        print(\"Could not run Inference: \", e)\n",
    "\n",
    "if __name__=='__main__':\n",
    "    parser=argparse.ArgumentParser()\n",
    "    parser.add_argument('--model', required=True)\n",
    "    parser.add_argument('--device', default='CPU')\n",
    "    parser.add_argument('--video', default=None)\n",
    "    parser.add_argument('--queue_param', default=None)\n",
    "    parser.add_argument('--output_path', default='/results')\n",
    "    parser.add_argument('--max_people', default=2)\n",
    "    parser.add_argument('--threshold', default=0.60)\n",
    "    \n",
    "    args=parser.parse_args()\n",
    "\n",
    "    main(args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Next Step\n",
    "\n",
    "Now that you've run the above cell and created your Python script, you will create your job submission shell script in the next workspace.\n",
    "\n",
    "**Note**: As a reminder, if you need to make any changes to the Python script, you can come back to this workspace to edit and run the above cell to overwrite the file with your changes."
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}