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": [
      "Overwriting 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",
    "        self.w_image = 0\n",
    "        self.h_image = 0\n",
    "\n",
    "        try:\n",
    "            try:\n",
    "                self.core=IECore()\n",
    "                self.model=self.core.read_network(model=self.model_structure, weights=self.model_weights)\n",
    "            except AttributeError:\n",
    "                self.model=IENetwork(model=self.model_structure, weights=self.model_weights)\n",
    "        except Exception:\n",
    "            raise ValueError(\"Could not Initialise the network. Have you entered the correct model path?\")\n",
    "\n",
    "        self.input_name=next(iter(self.model.inputs))\n",
    "        # Gets the input shape of the network\n",
    "        self.input_shape=self.model.inputs[self.input_name].shape\n",
    "        self.output_name=next(iter(self.model.outputs))\n",
    "        # Gets the output shape of the network\n",
    "        self.output_shape=self.model.outputs[self.output_name].shape\n",
    "\n",
    "    def load_model(self):\n",
    "        '''\n",
    "        Loads the network \n",
    "        '''\n",
    "        num_req=1\n",
    "        \n",
    "        self.net=self.core.load_network(network=self.model, device_name=self.device, num_requests=num_req)\n",
    "    \n",
    "        return self.net\n",
    "\n",
    "        \n",
    "    def predict(self, image):\n",
    "        '''\n",
    "        Image prepreprocessing, postprocessing and drawing bounding boxes\n",
    "        '''\n",
    "        request_id = 0\n",
    "        \n",
    "        preprocess_image = self.preprocess_input(image)\n",
    "        input_dict = {self.input_name:preprocess_image}\n",
    "        \n",
    "        # Makes an asynchronous inference request, given an input image.\n",
    "        self.net.start_async(request_id=request_id, inputs=input_dict)\n",
    "        if self.net.requests[request_id].wait(-1) ==0:\n",
    "            # Parse detection results to the current request\n",
    "            out = self.net.requests[request_id].outputs[self.output_name]\n",
    "\n",
    "            coords = self.preprocess_outputs(out)\n",
    "            box, image = self.draw_outputs(coords, image)\n",
    "            \n",
    "        return box, image\n",
    "              \n",
    "    \n",
    "    def draw_outputs(self, coords, image):\n",
    "        '''\n",
    "        Find and draw the bounding box\n",
    "        '''\n",
    "        print(\"draw_outputs\")\n",
    "        h_image = image.shape[0]\n",
    "        w_image = image.shape[1]\n",
    "        img_coords = []\n",
    "        color = (11, 156, 254)\n",
    "        for coord in coords:\n",
    "            # Draw objects\n",
    "                xmin = int(coord[3] * w_image)\n",
    "                ymin = int(coord[4] * h_image)\n",
    "                xmax = int(coord[5] * w_image)\n",
    "                ymax = int(coord[6] * h_image)\n",
    "                # Draw the bounding box\n",
    "                cv2.rectangle(image, (xmin, ymin), (xmax, ymax), color, 2)\n",
    "                img_coords.append([xmin, ymin, xmax, ymax])\n",
    "        \n",
    "        return img_coords, image\n",
    "    \n",
    "\n",
    "    def preprocess_outputs(self, outputs):\n",
    "        '''\n",
    "        # Draw objects only when probability is more than specified threshold\n",
    "        '''\n",
    "        thres_coord = []\n",
    "        for coord in outputs[0][0]:\n",
    "            if coord[2] > self.threshold:\n",
    "                thres_coord.append(coord)\n",
    "        return thres_coord \n",
    "\n",
    "    def preprocess_input(self, image):\n",
    "        '''\n",
    "        Change image-data layout from HWC to CHW \n",
    "        '''\n",
    "        h_img = self.input_shape[2]\n",
    "        w_img = self.input_shape[3]\n",
    "        \n",
    "        img = cv2.resize(image, (w_img, h_img))\n",
    "        img = img.transpose((2, 0, 1))\n",
    "        img = img.reshape(1, *img.shape)\n",
    "        \n",
    "        return 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",
    "            \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
}