42 lines
1.1 KiB
Python
42 lines
1.1 KiB
Python
from ultralytics import YOLO
|
|
from paths import * # make sure model_path is defined here
|
|
import cv2
|
|
|
|
if __name__ == "__main__":
|
|
|
|
print("Initializing model...")
|
|
model = YOLO(model_path)
|
|
|
|
print("Initializing camera...")
|
|
cap = cv2.VideoCapture(0)
|
|
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
|
|
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
|
|
cap.set(cv2.CAP_PROP_FPS, 30)
|
|
|
|
print("Initialized")
|
|
if not cap.isOpened():
|
|
print("Error: Could not open camera")
|
|
exit()
|
|
|
|
cv2.namedWindow("Predictions", cv2.WINDOW_NORMAL)
|
|
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
print("Error: Failed to grab frame")
|
|
break
|
|
|
|
# Optional: resize frame to improve YOLO performance
|
|
# frame = cv2.resize(frame, (416, 416))
|
|
|
|
results = model.predict(source=frame, conf=0.5)
|
|
|
|
annotated_frame = results[0].plot() # annotated frame as NumPy array
|
|
|
|
cv2.imshow("Predictions", annotated_frame)
|
|
cv2.resizeWindow("Predictions", 640, 640)
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
break
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|