diff --git a/etc/yolo_vision.py b/etc/yolo_vision.py index 39a8dac..ecaba72 100644 --- a/etc/yolo_vision.py +++ b/etc/yolo_vision.py @@ -8,13 +8,11 @@ import torch from PIL import Image import numpy as np import time -from picamera2 import Picamera2 import warnings import sys import datetime -# Suppress FutureWarnings from libraries -warnings.filterwarnings("ignore", category=FutureWarning) +PI_CAM = 0 # 1 for Raspberry Pi Camera, 0 for USB webcam # Load YOLOv5 model, other options include 'yolov5m', 'yolov5l', 'yolov5x' model = torch.hub.load("ultralytics/yolov5", "yolov5s") @@ -27,14 +25,31 @@ IGNORE_STATIONARY = True # Whether to ignore stationary objects in output ALERT_FUSE_COUNT = 5 # Number of consecutive detections before alerting ALERT_FILE_PATH = "alert.txt" # e.g., "/opt/meshing-around/alert.txt" or None for no file output -picam2 = Picamera2() -if LOW_RES_MODE: - picam2.preview_configuration.main.size = (320, 240) +if PI_CAM: + from picamera2 import Picamera2 else: - picam2.preview_configuration.main.size = (640, 480) -picam2.preview_configuration.main.format = "RGB888" -picam2.configure("preview") -picam2.start() + import cv2 + +# Suppress FutureWarnings from libraries +warnings.filterwarnings("ignore", category=FutureWarning) + +if PI_CAM: + picam2 = Picamera2() + if LOW_RES_MODE: + picam2.preview_configuration.main.size = (320, 240) + else: + picam2.preview_configuration.main.size = (640, 480) + picam2.preview_configuration.main.format = "RGB888" + picam2.configure("preview") + picam2.start() +else: + if LOW_RES_MODE: + cam_res = (320, 240) + else: + cam_res = (640, 480) + cap = cv2.VideoCapture(0) + cap.set(cv2.CAP_PROP_FRAME_WIDTH, cam_res[0]) + cap.set(cv2.CAP_PROP_FRAME_HEIGHT, cam_res[1]) print("="*40) print(" Sentinal Vision 3000 Booting Up!") @@ -54,7 +69,14 @@ try: system_normal_printed = False # system nominal flag, if true disables printing while True: i += 1 - frame = picam2.capture_array() + if PI_CAM: + frame = picam2.capture_array() + else: + ret, frame = cap.read() + if not ret: + print("Failed to grab frame from webcam.") + break + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) img = Image.fromarray(frame) results = model(img) @@ -133,5 +155,9 @@ except KeyboardInterrupt: except Exception as e: print(f"\nAn error occurred: {e}", file=sys.stderr) finally: - picam2.close() - print("Camera closed. Goodbye!") \ No newline at end of file + if PI_CAM: + picam2.close() + print("Camera closed. Goodbye!") + else: + cap.release() + print("Webcam released. Goodbye!") \ No newline at end of file