detector_YOLOv5.py

import pathlib
pathlib.PosixPath = pathlib.WindowsPath

MODEL = 0 # {0, 1}
MODELS = ["yolov5s.pt", "yolov5s_package_detection.pt"]

import cv2
import torch.hub as hub
import numpy as np

yolo = hub.load(
	"ultralytics/yolov5", 
	"custom", 
	path = f"YOLOv5/{MODELS[MODEL]}",
)

if MODEL == 1:
	yolo.conf = 0.8

cap = cv2.VideoCapture(1)

while True:
	ret, frame = cap.read()
	img = np.array(frame)
	results = yolo(img)
	cv2.imshow(f"YOLOv5 - {MODELS[MODEL]}", np.squeeze(results.render()))
	
	if cv2.waitKey(1) & 0xFF == ord("q"):
		break

cap.release()
cv2.destroyAllWindows()