quail_egg_detection / code /detect_send_mqtt_v2.py
jan-martens0124's picture
Update code/detect_send_mqtt_v2.py
8818220 verified
import time
import cv2
from ultralytics import YOLO
from paho.mqtt import client as mqtt
import json
import random
import threading
# bufferless VideoCapture
class VideoCapture:
def __init__(self, name):
self.cap = cv2.VideoCapture(name)
self.lock = threading.Lock()
self.t = threading.Thread(target=self._reader)
self.t.daemon = True
self.t.start()
# grab frames as soon as they are available
def _reader(self):
while True:
with self.lock:
ret = self.cap.grab()
if not ret:
break
# retrieve latest frame
def read(self):
print("read");
with self.lock:
_, frame = self.cap.retrieve()
return frame
# Desired frame rate (in seconds)
frame_rate = 1 # 1 / fps
# MQTT Settings
mqtt_username = "mqtt-user"
mqtt_password = "mqtt-password"
broker = 'ip'
port = 1883
topic = "your_topic/subject"
client_id = f'publish-{random.randint(0, 1000)}'
retain = False
aantal_eieren_detected = 0
# Load the YOLOv8 model
model = YOLO("./best.pt")
#video_path = "your_path"
#cap = cv2.VideoCapture(video_path)
# Alternatively: use webcam as input
# cap = cv2.VideoCapture(0)
# Alternatively: use rtsp-stream as input
url = "rtsp://admin:password@ip:8554/Streaming/Channels/101"
cap = cv2.VideoCapture(url)
def open_stream(url):
cap = cv2.VideoCapture(url)
return cap
def restart_stream_input():
cap.release()
cap = cv2.VideoCapture(url)
mqtt_client = mqtt.Client()
mqtt_client.username_pw_set(mqtt_username, mqtt_password)
mqtt_client.connect(broker, port)
mqtt_client.loop_start()
config_payload_json = {
"object_id": client_id,
"name": "kwartel_eieren",
"aantal_eieren": aantal_eieren_detected
}
def make_payload(aantal_eieren_detected):
config_payload_json = {
"object_id": client_id,
"name": "kwartel_eieren",
"aantal_eieren": aantal_eieren_detected
}
config_payload = json.dumps(config_payload_json);
return config_payload;
# Eerste keer publishen van aantal eieren (0)
mqtt_client.publish(topic, json.dumps(config_payload_json), retain=retain, qos=1)
print("Published", aantal_eieren_detected, "eieren detected")
# Variables to track time
start_time = time.time()
prev_frame_time = start_time
# Start infinite loop for when the stream stops / is too slow so end is reached and we need to restart the stream
while True:
cap = open_stream(url);
# Loop through the video frames
while cap.isOpened():
# Get current time
current_time = time.time()
# Calculate time elapsed since last frame
elapsed_time = current_time - prev_frame_time
# If elapsed time is less than the desired frame rate, continue to the next iteration
if elapsed_time < 1.0 / frame_rate:
ret, _ = cap.read()
continue
# Update previous frame time
prev_frame_time = current_time
# Read a frame from the video
success, frame = cap.read()
if success:
# Run YOLOv8 inference on the frame;
results = model(frame);
# Count the number of detections in the frame
aantal_eieren_detected = len(results[0]);
# Print the number of detections in the frame
print(f"Number of detections in frame: {aantal_eieren_detected}");
payload = make_payload(aantal_eieren_detected);
mqtt_client.publish(topic, payload, retain=retain, qos=1);
# SHOW BOUNDING BOXES (optional)
# Visualize the results on the frame
# annotated_frame = results[0].plot()
# Display the annotated frame
# cv2.imshow("YOLOv8 Inference", annotated_frame)
# Break the loop if 'q' is pressed
if cv2.waitKey(1) & 0xFF == ord("q"):
break
else:
# Break the loop if the end of the video is reached
print("End of video reached")
break
cap.release()
cv2.destroyAllWindows()
time.sleep(.2);
# Release the video capture object and close the display window
cap.release()
cv2.destroyAllWindows()