263 lines
8.4 KiB
Python
263 lines
8.4 KiB
Python
####################################################################################
|
|
# Load Libraries
|
|
|
|
from ultralytics import YOLO
|
|
import cvzone
|
|
import cv2
|
|
import threading
|
|
import queue
|
|
import paho.mqtt.client as mqtt
|
|
import time
|
|
from filterpy.kalman import UnscentedKalmanFilter as UKF
|
|
from filterpy.kalman import MerweScaledSigmaPoints
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
####################################################################################
|
|
# Initialize Variables
|
|
|
|
# Load YOLOv11 model
|
|
device = 'cuda'
|
|
model = YOLO("yolo11n-pecan.pt") # Update with correct model path
|
|
|
|
cy1 = 475 # Count Threshold (Center Y-coordinate)
|
|
offset = 60 # Offset for Center Y-coordinate
|
|
ids = set()
|
|
pecanCount = 0
|
|
samplePeriod = 0.1 # seconds
|
|
refThroughput = 0 # Count / Second
|
|
switchState = 0 # 0 = Off, 1 = On
|
|
prevSwitchState = 0 # 0 = Off, 1 = On
|
|
|
|
####################################################################################
|
|
# Initialize Camera
|
|
|
|
cap = cv2.VideoCapture(0) # Ensure that you have the correct camera index
|
|
|
|
# Set the resolution and other properties for each camera
|
|
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
|
|
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
|
|
cap.set(cv2.CAP_PROP_FPS, 120)
|
|
|
|
####################################################################################
|
|
# Unscented Kalman Filter
|
|
|
|
# State transition function (identity - no control input)
|
|
def fx(x, dt):
|
|
""" State transition function (no external control) """
|
|
return x # No change in state without an input
|
|
|
|
# Measurement function (identity)
|
|
def hx(x):
|
|
""" Measurement function (direct observation) """
|
|
return x # We measure the state directly
|
|
|
|
points = MerweScaledSigmaPoints(n=1, alpha=0.1, beta=2, kappa=0)# Define sigma points
|
|
|
|
ukf = UKF(dim_x=1, dim_z=1, fx=fx, hx=hx, points=points, dt=samplePeriod) # Initial State Estimate
|
|
ukf.x = np.array([refThroughput]) # Initial state estimate
|
|
ukf.Q = np.array([[0.02]]) # Process noise covariance (Q) - controls how much the state changes naturally
|
|
ukf.R = np.array([[0.5]]) # Measurement noise covariance (R) - how noisy the measurements are
|
|
ukf.P = np.eye(1) * 0.1 # Initial state covariance (P) - initial uncertainty
|
|
|
|
####################################################################################
|
|
# PI Controller
|
|
|
|
class PIController:
|
|
def __init__(self, Kp, Ki, Ts):
|
|
self.Kp = Kp
|
|
self.Ki = Ki
|
|
self.Ts = Ts
|
|
self.prevError = 0
|
|
self.prevOutput = 55
|
|
self.error = 0
|
|
|
|
|
|
def compute(self, setpoint, measurement):
|
|
"""Compute PI control output."""
|
|
self.error = setpoint - measurement
|
|
|
|
output = self.prevOutput + self.Kp * self.error + (self.Ki * self.Ts - self.Kp) * self.prevError # Z domain
|
|
|
|
self.prevError = self.error
|
|
|
|
# Apply saturation limits (0 to 90)
|
|
if output > 90:
|
|
output = 90
|
|
elif output < 0:
|
|
output = 0
|
|
|
|
self.prevOutput = output
|
|
|
|
return output
|
|
|
|
# PI Initialization
|
|
controller = PIController(0.063, 0.5, samplePeriod) # Kp, Ki, Ts
|
|
|
|
####################################################################################
|
|
# MQTT
|
|
|
|
# Mqtt Configuration
|
|
MQTT_BROKER = "192.168.1.110"
|
|
MQTT_PORT = 1883
|
|
MQTT_REF_TOPIC = "/jc/feedrate/"
|
|
MQTT_COUNT_TOPIC = "/jc/feedrate/count/"
|
|
MQTT_CONTROL_TOPIC = "jc/status/light:0"
|
|
MQTT_CONTROL_VISUAL_TOPIC = "/pi_controller/output"
|
|
|
|
# MQTT Handling
|
|
def on_connect(client, userdata, flags, rc, properties=None):
|
|
print("Connected with result code " + str(rc))
|
|
client.subscribe(MQTT_REF_TOPIC)
|
|
|
|
def on_message(client, userdata, message):
|
|
global refThroughput, switchState, prevSwitchState, controller
|
|
|
|
with refThroughput_lock:
|
|
refThroughput = float(message.payload.decode())
|
|
|
|
|
|
with ukf_lock:
|
|
ukf.x = np.array([refThroughput])
|
|
|
|
with switchState_lock:
|
|
if refThroughput == 0:
|
|
switchState = 0
|
|
elif prevSwitchState == 0 and refThroughput != 0:
|
|
switchState = 1
|
|
with controller_lock:
|
|
controller.prevError = 0
|
|
controller.prevOutput = 60
|
|
else:
|
|
switchState = 1
|
|
|
|
prevSwitchState = switchState
|
|
|
|
# Initialize MQTT client
|
|
mqtt_client = mqtt.Client()
|
|
mqtt_client.on_connect = on_connect
|
|
mqtt_client.on_message = on_message
|
|
mqtt_client.connect(MQTT_BROKER, MQTT_PORT, 60)
|
|
mqtt_client.loop_start() # Starts the loop in the background
|
|
|
|
####################################################################################
|
|
# Thread Initialization
|
|
|
|
# Queue to hold frames captured by the capture thread
|
|
frame_queue = queue.Queue(maxsize=1) # Limit to 1 frames in the queue
|
|
|
|
# Locks for thread synchronization
|
|
refThroughput_lock = threading.Lock()
|
|
ukf_lock = threading.Lock()
|
|
switchState_lock = threading.Lock()
|
|
controller_lock = threading.Lock()
|
|
|
|
####################################################################################
|
|
# Camera Thread
|
|
|
|
def capture_thread():
|
|
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
# Define the black box position (adjust as needed)
|
|
top_left = (0, 600)
|
|
bottom_right = (1280, 720)
|
|
|
|
# Draw a black rectangle (filled)
|
|
cv2.rectangle(frame, top_left, bottom_right, (0, 0, 0), thickness=-1)
|
|
|
|
# Flip Frame
|
|
frame = cv2.flip(frame, 0)
|
|
|
|
# Put the frame into the queue
|
|
if not frame_queue.full():
|
|
frame_queue.put(frame)
|
|
|
|
####################################################################################
|
|
# Processing Thread
|
|
|
|
def processing_thread():
|
|
global pecanCount, ids
|
|
sampleStart = time.time()
|
|
|
|
while True:
|
|
# Check if there are any frames in the queue
|
|
if not frame_queue.empty():
|
|
frame = frame_queue.get()
|
|
|
|
results = model.track(frame, persist=True, classes=0, device=device, stream=True, tracker='bytetrack_custom.yaml', iou = 0.6)
|
|
|
|
for result in results:
|
|
# Check if there are any boxes in the result
|
|
if result.boxes is not None and result.boxes.id is not None:
|
|
# Get the boxes (x, y, w, h), class IDs, track IDs, and confidences
|
|
boxes = result.boxes.xyxy.int().tolist() # Bounding boxes
|
|
track_ids = result.boxes.id.int().tolist() # Track IDs
|
|
|
|
# Iterate over each detected object
|
|
for box, track_id in zip(boxes, track_ids):
|
|
x1, y1, x2, y2 = box
|
|
cy = int((y1 + y2) // 2) # Center Y-coordinate
|
|
|
|
if cy < (cy1 + offset) and cy > (cy1 - offset) and track_id not in ids:
|
|
pecanCount += 1
|
|
ids.add(track_id)
|
|
|
|
sampleEnd = time.time()
|
|
if (sampleEnd - sampleStart) > samplePeriod:
|
|
measuredCount = pecanCount / samplePeriod
|
|
|
|
with ukf_lock:
|
|
ukf.predict()
|
|
ukf.update(np.array([measuredCount]))
|
|
|
|
filteredCount = ukf.x[0]
|
|
|
|
if switchState == 1:
|
|
with controller_lock:
|
|
controllerOutput = (controller.compute(refThroughput, filteredCount))
|
|
mqtt_client.publish(MQTT_CONTROL_TOPIC, str(int((controllerOutput) / 90 * 129)))
|
|
mqtt_client.publish(MQTT_CONTROL_VISUAL_TOPIC, str(controllerOutput))
|
|
|
|
mqtt_client.publish(MQTT_COUNT_TOPIC, str(filteredCount))
|
|
|
|
pecanCount = 0 # Reset count for next sample period
|
|
sampleStart = time.time() # Reset start time for next sample period
|
|
|
|
####################################################################################
|
|
# Threading
|
|
|
|
# Start capture and processing threads
|
|
capture_thread = threading.Thread(target=capture_thread)
|
|
capture_thread.daemon = True # Ensures the thread exits when the main program exits
|
|
capture_thread.start()
|
|
|
|
processing_thread = threading.Thread(target=processing_thread)
|
|
processing_thread.daemon = True # Ensures the thread exits when the main program exits
|
|
processing_thread.start()
|
|
|
|
####################################################################################
|
|
# Clean Up
|
|
|
|
try:
|
|
while True:
|
|
# Keep the main loop alive
|
|
if cv2.waitKey(1) & 0xFF == ord('q'): # Press 'q' to quit
|
|
break
|
|
except KeyboardInterrupt:
|
|
print("Process interrupted")
|
|
|
|
# Release resources when exiting
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
mqtt_client.loop_stop()
|
|
mqtt_client.disconnect()
|
|
print("Capture Released")
|
|
|
|
####################################################################################
|
|
|
|
|