Updated Count Script

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)