
import time
import multiprocessing
from multiprocessing import Process
#from torch.multiprocessing import Pool, Process, set_start_method
from yolov7_trt import YoLov7TRT
import cv2
#print(torch.cuda.is_available())
engine_file_path = "build/yolov7.engine"
yolo_v7_wrapper = YoLov7TRT(engine_file_path)

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

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

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

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

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




