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": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting person_detect.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile person_detect.py\n",
    "\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",
    "    Performs basic operations for queues like adding to a queue, getting the queues \n",
    "    and checking the coordinates for queues.\n",
    "    \n",
    "    Attributes:\n",
    "        queues: A list containing the queues data\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self):\n",
    "        self.queues=[]\n",
    "\n",
    "    def add_queue(self, points):\n",
    "        \"\"\"\n",
    "        Add points to the queue.\n",
    "\n",
    "        Args:\n",
    "            points: A list of points to be added.\n",
    "\n",
    "        Raises:\n",
    "            TypeError: points is None.\n",
    "        \"\"\"\n",
    "        \n",
    "        self.queues.append(points)\n",
    "\n",
    "    def get_queues(self, image):\n",
    "        \"\"\"\n",
    "        Get queues from images.\n",
    "\n",
    "        Args:\n",
    "            image: A list of the image.\n",
    "\n",
    "        Yields:\n",
    "            A list containing each frame.\n",
    "        \"\"\"\n",
    "            \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, init_w, init_h):\n",
    "        \"\"\"\n",
    "        Check coordinates for queues.\n",
    "\n",
    "        Args:\n",
    "            coords: A list of the coordinates.\n",
    "        \"\"\"\n",
    "        \n",
    "        d={k+1:0 for k in range(len(self.queues))}\n",
    "        \n",
    "        dum = ['0', '1' , '2', '3']\n",
    "        \n",
    "        for coord in coords:\n",
    "            xmin = int(coord[3] * init_w)\n",
    "            ymin = int(coord[4] * init_h)\n",
    "            xmax = int(coord[5] * init_w)\n",
    "            ymax = int(coord[6] * init_h)\n",
    "            \n",
    "            dum[0] = xmin\n",
    "            dum[1] = ymin\n",
    "            dum[2] = xmax\n",
    "            dum[3] = ymax\n",
    "            \n",
    "            for i, q in enumerate(self.queues):\n",
    "                if dum[0]>q[0] and dum[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",
    "    Performs person detection and preprocessing.\n",
    "    \n",
    "    Attributes:\n",
    "        model_weights: A string containing model weights path.\n",
    "        model_structure: A string conatining model structure path.\n",
    "        device: A string conatining device name.\n",
    "        threshold: A floating point number containing threshold value.\n",
    "        input_name: A list of input names.\n",
    "        input_shape: A tuple of the input shape.\n",
    "        output_name: A list of output names.\n",
    "        output_shape: A tuple of the output shape.\n",
    "        core: IECore object.\n",
    "        net: Loaded net object.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, model_name, device, threshold=0.60):\n",
    "        \"\"\"\n",
    "        Inits PersonDetect class with model_name, device, threshold.\n",
    "        \"\"\"\n",
    "        \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(model = self.model_structure,weights = self.model_weights)\n",
    "            \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",
    "        print('Model creation...')\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",
    "        \"\"\"\n",
    "        Loads the model.\n",
    "        \"\"\"\n",
    "        self.network = self.core.load_network(network=self.model, device_name=self.device, num_requests=1)\n",
    "        print('Network Loading...')\n",
    "        \n",
    "    def predict(self, image):\n",
    "        \"\"\"\n",
    "        Make asynchronous predictions from images.\n",
    "\n",
    "        Args:\n",
    "            image: List of the image data.\n",
    "\n",
    "        Returns:\n",
    "            The outputs and the image.\n",
    "        \"\"\"\n",
    "        \n",
    "        input_name = self.input_name\n",
    "        input_img = self.preprocess_input(image)      \n",
    "        input_dict={input_name: input_img}  \n",
    "        \n",
    "        # Start asynchronous inference for specified request.\n",
    "\n",
    "        infer_request = self.network.start_async(request_id=0, inputs=input_dict)\n",
    "        infer_stat = infer_request.wait()\n",
    "        if infer_stat == 0:\n",
    "            outputs = infer_request.outputs[self.output_name]\n",
    "            \n",
    "        return outputs, image\n",
    "    \n",
    "    def draw_outputs(self, coords, frame, init_w, init_h):\n",
    "        \"\"\"\n",
    "        Draws outputs or predictions on image.\n",
    "\n",
    "        Args:\n",
    "            coords: The coordinates of predictions.\n",
    "            image: The image on which boxes need to be drawn.\n",
    "\n",
    "        Returns:\n",
    "            the frame\n",
    "            the count of people\n",
    "            bounding boxes above threshold\n",
    "        \"\"\"\n",
    "        \n",
    "        cur_count = 0\n",
    "        det = []\n",
    "        \n",
    "        for ob in coords[0][0]:\n",
    "            \n",
    "            # Draw bounding box for the detected object when it's probability \n",
    "            # is more than the specified threshold\n",
    "            if ob[2] > self.threshold:\n",
    "                start = (int(ob[3] * init_w),int(ob[4] * init_h))\n",
    "                end = (int(ob[5] * init_w),int(ob[6] * init_h))\n",
    "                cv2.rectangle(frame, start, end, (0, 255, 0), 1)\n",
    "                cur_count = cur_count + 1\n",
    "                \n",
    "                det.append(ob)\n",
    "                \n",
    "        return frame, cur_count, det\n",
    "\n",
    "    def preprocess_outputs(self, outputs):\n",
    "        \"\"\"\n",
    "        Preprocess the outputs.\n",
    "\n",
    "        Args:\n",
    "            outputs: The output from predictions.\n",
    "\n",
    "        Returns:\n",
    "            Preprocessed dictionary.\n",
    "        \"\"\"\n",
    "        #out_dict = {self.output_name:out for out in outputs}\n",
    "        out_dict = {}\n",
    "        for out in outputs:\n",
    "            out_name = self.output_name\n",
    "            out_img = out\n",
    "            out_dict[output_name] = out_img\n",
    "        \n",
    "        return out_dict\n",
    "    \n",
    "        return out\n",
    "        \n",
    "\n",
    "    def preprocess_input(self, image):\n",
    "      \n",
    "        in_img = image\n",
    "        \n",
    "        # Preprocessing input\n",
    "        n, c, h, w = self.input_shape\n",
    "        \n",
    "        in_img=cv2.resize(in_img, (w, h), interpolation = cv2.INTER_AREA)\n",
    "        \n",
    "        # Change image from HWC to CHW\n",
    "        in_img = in_img.transpose((2, 0, 1))\n",
    "        in_img = in_img.reshape((n, c, h, w))\n",
    "\n",
    "        return in_img\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",
    "            coords, image= pd.predict(frame)\n",
    "            frame, current_count, coords = pd.draw_outputs(coords, image, initial_w, initial_h)\n",
    "        \n",
    "            num_people = queue.check_coords(coords, initial_w, initial_h)\n",
    "            print(f\"Total People in frame = {len(coords)}\")\n",
    "            print(f\"Number of people in queue = {num_people}\")\n",
    "            \n",
    "            out_text=\"\"\n",
    "            y_pixel=25\n",
    "            \n",
    "            for k, v in num_people.items():\n",
    "                print(k, v)\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",
    "\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",
    "        \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
}