import time
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):
    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(result_classid)


        
    cap.release()

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

Videos = ['traffic1.mp4']
if __name__ =='__main__':
    detectObject(video="traffic1.mp4")
    # for i in Videos:
    #     process = Process(target=detectObject, args=(i, str(i)))
    #     process.start()

