{ "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": [] } ] }