traffic-AI/src/trafficApp.py

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)