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