diff --git a/Count.py b/Count.py index 9e0e5b0..5e27a58 100644 --- a/Count.py +++ b/Count.py @@ -1,27 +1,106 @@ import cv2 import numpy as np from ultralytics import YOLO -import cvzone +import paho.mqtt.client as mqtt +from filterpy.kalman import UnscentedKalmanFilter as UKF +from filterpy.kalman import MerweScaledSigmaPoints +import time -device = 'cuda' +############################################################################################################ +# Initialize Variables -# Load the YOLO11 model +cy1=550 +offset=30 +ids=set() +pecanCount = 0 +avgPecanWeight = 0.0270453125 # lbs +refThroughput = 0 # lbs / 15 seconds +avgSampleTime = 0.5 # seconds + +# Load the YOLO11 model and set device model = YOLO("yolo11m-pecan.pt") +device = 'cuda' # Open the video file (use video file or webcam, here using webcam) cap = cv2.VideoCapture(0) -cy1=550 -offset=30 -idDict={} -pecanCount = 0 + +############################################################################################################ +# 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) # Initialize UKF +ukf.x = np.array([refThroughput / 15 / avgPecanWeight * avgSampleTime ]) # Initial state estimate +ukf.Q = np.array([[0.02]]) # Process noise covariance (Q) - controls how much the state changes naturally +ukf.R = np.array([[1]]) # Measurement noise covariance (R) - how noisy the measurements are +ukf.P = np.eye(1) * 0.1 # Initial state covariance (P) - initial uncertainty + +############################################################################################################ +# MQTT Configuration + +MQTT_BROKER = "192.168.1.110" +MQTT_PORT = 1883 +MQTT_TOPIC = "/jc/feedrate/count" + +def on_connect(client, userdata, flags, rc, properties=None): + print("Connected with result code " + str(rc)) + client.subscribe(MQTT_TOPIC) + +def on_message(client, userdata, message): + global refThroughput # lbs of pecans / 15 seconds + prevRefThroughput = refThroughput + refThroughput = float(message.payload.decode()) + if refThroughput != prevRefThroughput: + ukf.x = np.array([refThroughput]) + +client = mqtt.Client() +client.on_connect = on_connect +client.on_message = on_message + +client.connect(MQTT_BROKER, MQTT_PORT, 60) +client.loop_start() # Starts MQTT loop in the background +############################################################################################################ +# Exponential Moving Average + +# class FastEMA: # Exponential Moving Average Function +# def __init__(self, alpha): +# self.alpha = alpha +# self.ema = None # Initialize without a value +# self.input = None + +# def update(self, measure, input): +# if self.ema is None or self.input != input: +# self.ema = measure # Set first value +# self.input = input +# else: +# self.ema = self.alpha * measure + (1 - self.alpha) * self.ema +# return self.ema + +# # Initialize the EMA filter +# ema_filter = FastEMA(alpha=0.1) # Adjust alpha (lower = smoother) +############################################################################################################ while True: ret,frame = cap.read() if not ret: break + sampleStart = time.time() # Sample Start Time + + pecanCount = 0 # Reset count for new sample period + # Run YOLO11 tracking on the frame, persisting tracks between frames - results = model.track(frame, persist=True,classes=0,device = device) + results = model.track(frame, persist=True, classes=0, device = device) # Check if there are any boxes in the results if results[0].boxes is not None and results[0].boxes.id is not None: @@ -34,12 +113,26 @@ while True: for box, class_id, track_id, conf in zip(boxes, class_ids, track_ids, confidences): x1, y1, x2, y2 = box cy = int(y1+y2)//2 - if cy<(cy1+offset) and cy>(cy1-offset) and track_id not in idDict.keys(): + if cy<(cy1+offset) and cy>(cy1-offset) and track_id not in ids: pecanCount += 1 - idDict[track_id] = pecanCount + ids.add(track_id) + + # filtered_count = ema_filter.update(pecanCount, refThroughput) # Applies exponential moving average filter + + ukf.predict() + ukf.update(np.array([pecanCount])) + filtered_count = ukf.x[0] + + sampleEnd = time.time() # End Sample Timer + + samplePeriod = sampleEnd - sampleStart + print(samplePeriod) + + measuredThroughput = (filtered_count * avgPecanWeight) / (samplePeriod) * 15 # lbs / 15 seconds + + client.publish(MQTT_TOPIC, str(measuredThroughput)) + # Release the video capture object and close the display window cap.release() cv2.destroyAllWindows() - -print(pecanCount) \ No newline at end of file