import time
import multiprocessing
from multiprocessing import Process
#from torch.multiprocessing import Pool, Process, set_start_method
from yolov5_trt_lib import YoloV5TRT
import cv2
#print(torch.cuda.is_available())
engine_file_path = "build/yolov5s.engine"
yolo_v5_wrapper = YoloV5TRT(engine_file_path, conf_thresh= 0.5, iou_thresh= 0.4, )

def detectObject(video,name):
    cap = cv2.VideoCapture(video)
    while cap.isOpened():
        pTime = time.time()
        ret, img = cap.read()
        if img is None:
            break
        else:
            result_boxes, result_scores, result_classid = yolo_v5_wrapper.infer(img)
            print(str(result_classid), str(name))
            cTime = time.time()
            print("fps", str(1/(cTime - pTime)))

        #cv2.imshow(name, img)
        #cv2.waitKey(1)
        
    cap.release()

try:
     #set_start_method('spawn', force=True)
    multiprocessing.set_start_method('spawn')
except RuntimeError:
    pass

Videos = ['traffic1.mp4', 'traffic3.mp4']
if __name__ =='__main__':

    for i in Videos:
        process = Process(target=detectObject, args=(i, str(i)))
        process.start()



