atmadeep.a
/
traffic-AI


			
				
					
						
						
							
							{
 "metadata": {
  "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
  },
  "orig_nbformat": 2
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "\n",
    "import cv2 as cv\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "os.environ['DISPLAY'] = ':0'\n",
    "\n",
    "from config.config import PARAMS\n",
    "from src.numberPlateRoiDetection import NumberPlateROIDetection\n",
    "from src.objectDetection import ObjectDetection\n",
    "from src.ocrNumberPlate import get_number_plate_ocr_from_rois\n",
    "from src.parkingDetection import ParkingDetection\n",
    "from src.trackingManager import TrackerManager\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrafficApp(object):\n",
    "    def __init__(self,args):\n",
    "        self.args = args\n",
    "\n",
    "        #get Object Detection Up\n",
    "        self.objectDetection = ObjectDetection(debug=args.debug,target=args.target)\n",
    "        self.numberPlateDetection = NumberPlateROIDetection(args= args,algoType='NumberPlate')\n",
    "        self.parkingDetection = None #intilize later when we will have height/width\n",
    "        np.random.seed(41)\n",
    "\n",
    "        #fix color\n",
    "        self.colorToDisplay = {'numberplate':(0,255,255),'car':(0,255,0),'bus':(128,255,0),'truck':(0,0,255),'moterbike':(255,0,255),'ocr':(0,140,240)}\n",
    "        if self.args.video is not None:\n",
    "            self.vid_writer = None\n",
    "            self.runVideoFlow()\n",
    "\n",
    "    def runVideoFlow(self):\n",
    "        frame_count = 0\n",
    "        if args.video is not None:\n",
    "            try:\n",
    "                videoObj = cv.VideoCapture(args.video)\n",
    "                imgH, imgW = None, None\n",
    "                writer = None\n",
    "            except:\n",
    "                raise Exception('Video cannot be loaded! Please check the path provided!')\n",
    "\n",
    "            finally:\n",
    "                try:\n",
    "                    totalFrames = videoObj.get(cv.cv.CV_CAP_PROP_FRAME_COUNT)\n",
    "                except:\n",
    "                    totalFrames = -1\n",
    "\n",
    "                try:\n",
    "                    totalFrames = videoObj.get(cv.CAP_PROP_FRAME_COUNT)\n",
    "                except:\n",
    "                    totalFrames = -1\n",
    "                try:\n",
    "                    imgH = int(videoObj.get(cv.CAP_PROP_FRAME_HEIGHT))\n",
    "                    imgW = int(videoObj.get(cv.CAP_PROP_FRAME_WIDTH))\n",
    "                    TrackerManager.FrameHeight = imgH\n",
    "                    TrackerManager.FrameWidth = imgW\n",
    "                    print('Height, Width',imgH,imgW)\n",
    "                    if PARAMS._ALGO_MODE_PARKING:\n",
    "                        self.parkingDetection = ParkingDetection(imgW=imgW,imgH=imgH)\n",
    "                        self.parkingDetection.getParkingRegionMask()\n",
    "                    #videoObj.set(cv.CAP_PROP_POS_FRAMES, 225)\n",
    "                except:\n",
    "                    imgH = -1\n",
    "                    imgW = -1\n",
    "                    raise ValueError('Issue with video')\n",
    "                if self.args.debug:\n",
    "                    print('Frames-{},Height-{}, Width-{}'.format(totalFrames,imgH,imgW))\n",
    "\n",
    "                if self.args.saveoutput and (imgH > 0 and imgW > 0):\n",
    "                    self.vid_writer = cv.VideoWriter(self.args.outputfile,\n",
    "                                                      cv.VideoWriter_fourcc(*\"MJPG\"), 30,\n",
    "                                                      (round(imgW),round(imgH)))\n",
    "                \n",
    "                progress_bar=tqdm(total = totalFrames)\n",
    "                # start reading frame\n",
    "                while True:\n",
    "                    grabbed, frame = videoObj.read()\n",
    "                    #frame[:,450:,:] = 0\n",
    "                    # end of frame\n",
    "                    if not grabbed:\n",
    "                        break\n",
    "                    frame_count +=1\n",
    "\n",
    "                    #print('Frame_count-',frame_count)\n",
    "                    #Use jump argument to skip frames.\n",
    "                    if (frame_count % self.args.jump == 0):\n",
    "                        \n",
    "                        # get object detection on this frame\n",
    "                        img_objectMarking, boxes, confidences, classids, idxs,status = self.objectDetection.run_object_detection(frame.copy(),imageH=imgH,imageW=imgW)\n",
    "                        '''Assign Trcakers'''\n",
    "                        object_detect_info = [boxes, confidences, classids, idxs, status]\n",
    "                        bbox_labels_tracking = self.parseObjDetectInfo(object_detect_info)\n",
    "                        TrackerManager.FrameCount = frame_count\n",
    "                        TrackerManager.manageTracker(bbox_labels_tracking)\n",
    "\n",
    "                        ''' Get Parking Status'''\n",
    "                        if PARAMS._ALGO_MODE_PARKING:\n",
    "                            self.parkingDetection.getParkingStatus(TrackerManager.TrackerList)\n",
    "\n",
    "                        '''Filter ROIs for Number Plate Detection'''\n",
    "                        tentative_numberplate_rios = self.objectDetection.filterRoiforNumberPlate(boxes, classids, idxs)\n",
    "\n",
    "\n",
    "                        ''' Get Number Plate ROI'''\n",
    "                        detected_np_info = self.numberPlateDetection.run_number_plate_detection_rois(image=frame.copy(),rois=tentative_numberplate_rios)\n",
    "\n",
    "\n",
    "                        ''' Get Number plate OCR '''\n",
    "                        number_plate_ocr_dict = get_number_plate_ocr_from_rois(frame.copy(),detected_np_info, False)\n",
    "\n",
    "                        #Display frame\n",
    "                        displayFrame = self.displayFrame(frame.copy(),detected_np_info,number_plate_ocr_dict,object_detect_info)\n",
    "\n",
    "                        winName = 'YOLOV3 Object Detection'\n",
    "                        cv.namedWindow(winName, cv.WINDOW_NORMAL)\n",
    "                        cv.imshow(winName, displayFrame)\n",
    "                        #cv.resizeWindow('objectDetection',680,420)\n",
    "                        if self.vid_writer:\n",
    "                            self.vid_writer.write(displayFrame.astype(np.uint8))\n",
    "                        c = cv.waitKey(1)\n",
    "                        if c & 0xFF == ord('q'):\n",
    "                            self.vid_writer.release()\n",
    "                            videoObj.release()\n",
    "                            break\n",
    "                        progress_bar.close()\n",
    "    def parseObjDetectInfo(self,object_roi_info):\n",
    "        boxes, confidences, classids, idxs, status = object_roi_info\n",
    "        #[[list of bbox ][list of conf and labels]]\n",
    "        bboxList =[]\n",
    "        confidence_labels = []\n",
    "        if len(idxs) > 0 and status:\n",
    "            for i in idxs.flatten():\n",
    "                # Get the bounding box coordinates\n",
    "                if self.objectDetection.labels[classids[i]] not in PARAMS._TRACKER_OBJECT_LIST +\\\n",
    "                        PARAMS._YOLOV3_OD_NUMBER_PLATE_OBJECT_LIST:\n",
    "                    continue\n",
    "                x, y = boxes[i][0], boxes[i][1]\n",
    "                w, h = boxes[i][2], boxes[i][3]\n",
    "                bboxList.append ([x,y,w,h])\n",
    "                confidence_labels.append([confidences[i],self.objectDetection.labels[classids[i]]])\n",
    "        return [bboxList,confidence_labels]\n",
    "\n",
    "    def displayFrame(self,displayFrame,numberplate_roi,number_plate_ocr_dict,object_roi_info):\n",
    "        debug = self.args.debug\n",
    "        if PARAMS._ALGO_MODE_NUMBER_PLATE:\n",
    "            #for nuber plate\n",
    "            for idx,roiinfo in enumerate(numberplate_roi):\n",
    "                conf, classID, roi = roiinfo\n",
    "                x, y, w, h = roi\n",
    "                cv.rectangle(displayFrame, (x, y), (x + w, y + h), self.colorToDisplay['numberplate'], 2)\n",
    "                text = \"{}: {:.3f}\".format(self.numberPlateDetection.labels[classID], conf)\n",
    "                #cv.putText(displayFrame, text, (x, y - 10), cv.FONT_HERSHEY_SIMPLEX, 0.5, self.colorToDisplay['numberplate'], 1)\n",
    "\n",
    "                #add Number plate OCR\n",
    "                if number_plate_ocr_dict[idx]:\n",
    "                    thickness = 4\n",
    "                    font_face = cv.FONT_HERSHEY_SIMPLEX\n",
    "                    font_scale = 1\n",
    "                    cv.putText(displayFrame, number_plate_ocr_dict[idx], (x, y-5), font_face, font_scale,\\\n",
    "                                self.colorToDisplay['ocr'], thickness)\n",
    "\n",
    "        if False:\n",
    "            #for objects\n",
    "            boxes, confidences, classids, idxs, status = object_roi_info\n",
    "            if len(idxs) > 0 and status:\n",
    "                for i in idxs.flatten():\n",
    "                    # Get the bounding box coordinates\n",
    "                    x, y = boxes[i][0], boxes[i][1]\n",
    "                    w, h = boxes[i][2], boxes[i][3]\n",
    "\n",
    "                    # Get the unique color for this class\n",
    "                    if self.objectDetection.labels[classids[i]] in self.colorToDisplay:\n",
    "                        color = self.colorToDisplay[self.objectDetection.labels[classids[i]]]\n",
    "                    else:\n",
    "                        color = [int(c) for c in self.objectDetection.colors[classids[i]]]\n",
    "                    #color = (255,255,255)\n",
    "                    # Draw the bounding box rectangle and label on the image\n",
    "                    cv.rectangle(displayFrame, (x, y), (x + w, y + h), color, 2)\n",
    "                    text = \"{}: {:.3f}\".format(self.objectDetection.labels[classids[i]], confidences[i])\n",
    "                    cv.putText(displayFrame, text, (x, y - 5), cv.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)\n",
    "\n",
    "\n",
    "        if True:\n",
    "            if len(TrackerManager.DetectionWithNoTracker)>0:\n",
    "                color = (0,0,0)\n",
    "                for item in TrackerManager.DetectionWithNoTracker:\n",
    "                    bbox,(conf,label) = item\n",
    "                    x,y,w,h = bbox\n",
    "                    # Draw the bounding box rectangle and label on the image\n",
    "                    cv.rectangle(displayFrame, (x, y), (x + w, y + h), color, 2)\n",
    "                    if debug:\n",
    "                        text = \"NotTrack-{}: {:.3f}\".format(label,conf)\n",
    "                        cv.putText(displayFrame, text, (x, y - 5), cv.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)\n",
    "\n",
    "        if PARAMS._ALGO_MODE_PARKING:\n",
    "            cv.line(displayFrame,PARAMS._NO_PARAKING_LINE_POINT_1_XY,PARAMS._NO_PARAKING_LINE_POINT_2_XY,\\\n",
    "                     (0,0,255),3,2)\n",
    "\n",
    "        if PARAMS._ALGO_MODE_KALMAN_TRCAKING:\n",
    "            if len(TrackerManager.TrackerList) > 0:\n",
    "                color = (0,255,0)\n",
    "                for tracker in TrackerManager.TrackerList:\n",
    "                    bbox = tracker.curr_frame_predict_bbox\n",
    "                    x,y,w,h = np.int32(bbox)\n",
    "                    missframe = tracker.objectInfo.ObjectTrackerMissedFrame\n",
    "                    direction = 'XX' if tracker.objectInfo.ObjectDirection is None else tracker.objectInfo.ObjectDirection\n",
    "                    objectType = tracker.objectInfo.ObjectType\n",
    "                    objectID = tracker.objectID\n",
    "                    if not tracker.objectInfo.ObjectParkingStatus:\n",
    "                        cv.rectangle(displayFrame, (x, y), (x + w, y + h), color, 2)\n",
    "                    else:\n",
    "                        cv.rectangle(displayFrame, (x, y), (x + w, y + h), (0,0,0), 3)\n",
    "\n",
    "                    #update curr box by which it was updated\n",
    "                    if False:\n",
    "                        bbox_detect = tracker.curr_frame_update_bbox\n",
    "                        xp,yp,wp,hp = bbox_detect\n",
    "                        cv.rectangle(displayFrame, (xp, yp), (xp + wp, yp + hp), (0,255,255), 2)\n",
    "                    if debug:\n",
    "                        text = \"{}-f{}-{}\".format(objectID,missframe,direction)\n",
    "                    else:\n",
    "                        text = \"{}\".format(direction)\n",
    "\n",
    "                    if tracker.objectInfo.ObjectParkingStatus and PARAMS._ALGO_MODE_PARKING:\n",
    "                        if tracker.objectInfo.ObjectType in PARAMS._YOLOV3_OD_NUMBER_PLATE_OBJECT_LIST:\n",
    "                            text = \"{}\".format(PARAMS._PARKING_STRING)\n",
    "                    font_scale = 1.5\n",
    "                    font = cv.FONT_HERSHEY_SIMPLEX #PLAIN #cv.FONT_HERSHEY_SIMPLEX\n",
    "                    # set the rect bg - BLACK\n",
    "                    rect_bgr = (0,0,0)\n",
    "                    # get the width and height of the text box\n",
    "                    (text_width, text_height) = np.int32(cv.getTextSize(text, font, fontScale=font_scale, thickness=2)[0])\n",
    "                    # make the coords of the box with a small padding of two pixels\n",
    "                    box_coords = ((x, y), (x + text_width + 5, y - text_height - 5))\n",
    "                    cv.rectangle(displayFrame, box_coords[0], box_coords[1], rect_bgr, cv.FILLED)\n",
    "                    cv.putText(displayFrame, text, (x, y), font, fontScale=font_scale, color=(0, 0, 255),thickness=2)\n",
    "        if True:\n",
    "            imglogo = cv.imread(PARAMS.LOGO_FILE_PATH)\n",
    "            logo = cv.resize(imglogo,dsize=(300,100),interpolation=cv.INTER_LINEAR)\n",
    "            h,w,c = logo.shape\n",
    "            H,W,C = displayFrame.shape\n",
    "            displayFrame[0:h,W-w-10:W-10,:] = logo\n",
    "\n",
    "        return displayFrame\n"
   ]
  },
  {
   "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": []
  }
 ]
}