traffic-AI/skeleton.py

import os
import cv2 
import argparse
import time
import random
import multiprocessing as mp
import threading as th
from queue import Queue

inputQueue = mp.Queue()
vehicleDetectionQueue = mp.Queue()
outputQueue = mp.Queue()
class ReadFrame(th.Thread):
    def __init__(self,source,name='Input thread',custom_id=1) -> None:
        super().__init__()
        self.frameId = 1
        self.stopped = False
        self.grabbed = True
        self.name = f'{name} {custom_id}'
        

        self.videoCaptureObject = cv2.VideoCapture(source)
        print(f'Reading from source = {source}')

    global inputQueue
    
    def run(self):
        while (self.grabbed):
            (self.grabbed, self.frame) = self.videoCaptureObject.read()
            inputQueue.put((self.frame,self.frameId))
            print(f"{self.name}frame added with id {self.frameId}\n")
            self.frameId+=1
        return
            

class VehicleDetection(mp.Process):
    def __init__(self,name='Vehicle Detection Process',custom_id=1):
        super(VehicleDetection,self).__init__()
        self.name = f'{name} {custom_id}'
        
    global inputQueue
    def run(self):
        while (True):
            (frame,frameId) = inputQueue.get()
            #inputQueue.task_done()
            print(f"{self.name}Got frame with ID {frameId} qsize = {inputQueue.qsize()}\n")
            #do some processing here.
            vehicleDetectionQueue.put((frame,frameId))
            if(inputQueue.qsize() < 1):
                return
            
            
    
class NumberPlateOcr(mp.Process):
    def __init__(self,name='Number plate OCR Process',custom_id=1):
        super(NumberPlateOcr,self).__init__()
        self.name=f'{name} {custom_id}'


    global inputQueue
    global numberPlateOcrQueue
    def run(self):
        while True:
            (frame,frameId) = vehicleDetectionQueue.get()
            #inputQueue.task_done()
            print(f"{self.name} Got frame with ID {frameId}\n")
            #do some processing here.
            outputQueue.put((frame,frameId))

class outputframe(th.Thread):
    def __init__(self,name='output thread',custom_id=1):
        super().__init__()
        self.name = f'{name} {custom_id}'
    
    def run(self):
        while True:
            (frame,frameId) = outputQueue.get()
            print(f'{self.name} got frame {frameId}\n')



      
if __name__ == '__main__':
    import cProfile

    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')
      
    args = parser.parse_args()
    #enable profiler here.
    app_profiler.enable()

    readFramesThread = ReadFrame(args.video)
    vehicleDetectionProcess = VehicleDetection()
    numberPlateOcrProcess = NumberPlateOcr()
    readFramesThread.start()
    vehicleDetectionProcess.start()
    numberPlateOcrProcess.start()
        
    

    #disable profiler here.
    app_profiler.disable()
    
    profile_name = str('temp.prof'.format(os.path.basename(args.video)[0:-4]))
    print("------------------------\nEnd of execution, dumping profile stats\n-------------------------")
    app_profiler.dump_stats(profile_name)