178 lines
6.4 KiB
Python
178 lines
6.4 KiB
Python
import cv2
|
|
import numpy as np
|
|
from ultralytics import YOLO
|
|
import paho.mqtt.client as mqtt
|
|
from filterpy.kalman import UnscentedKalmanFilter as UKF
|
|
from filterpy.kalman import MerweScaledSigmaPoints
|
|
import time
|
|
|
|
############################################################################################################
|
|
# Initialize Variables
|
|
|
|
cy1 = 550
|
|
offset = 90
|
|
ids = set()
|
|
pecanCount = 0
|
|
avgPecanWeight = 0.0270453125 # lbs
|
|
refThroughput = 0 # lbs / 15 seconds
|
|
samplePeriod = 0.25 # seconds
|
|
width = 1280
|
|
height = 720
|
|
totalPecanCount = 0
|
|
|
|
# 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('rtsp://192.168.1.10:8554/stream')
|
|
|
|
############################################################################################################
|
|
# 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 / 15 / avgPecanWeight * samplePeriod]) # 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/"
|
|
MQTT_COUNT_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)
|
|
############################################################################################################
|
|
|
|
sampleStart = time.time()
|
|
frameTime = time.time()
|
|
frameCount = 0
|
|
|
|
while True:
|
|
ret,frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
# Define the black box position (adjust as needed)
|
|
top_left = (0, 0)
|
|
bottom_right = (1280, 220)
|
|
|
|
# Draw a black rectangle (filled)
|
|
cv2.rectangle(frame, top_left, bottom_right, (0, 0, 0), thickness=-1)
|
|
|
|
frameCount += 1
|
|
|
|
# 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)
|
|
|
|
# Check if there are any boxes in the results
|
|
if results[0].boxes is not None and results[0].boxes.id is not None:
|
|
# Get the boxes (x, y, w, h), class IDs, track IDs, and confidences
|
|
boxes = results[0].boxes.xyxy.int().cpu().tolist() # Bounding boxes
|
|
class_ids = results[0].boxes.cls.int().cpu().tolist() # Class IDs
|
|
track_ids = results[0].boxes.id.int().cpu().tolist() # Track IDs
|
|
confidences = results[0].boxes.conf.cpu().tolist() # Confidence score
|
|
|
|
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 ids:
|
|
pecanCount += 1
|
|
totalPecanCount += 1
|
|
print(f'New Count: {pecanCount}')
|
|
ids.add(track_id)
|
|
|
|
# filtered_count = ema_filter.update(pecanCount, refThroughput) # Applies exponential moving average filter
|
|
|
|
sampleEnd = time.time()
|
|
print(f'Pecan Count: {pecanCount}')
|
|
print(f'Total Count: {totalPecanCount}')
|
|
if (sampleEnd - sampleStart) > samplePeriod:
|
|
ukf.predict()
|
|
print(f'Predicted State: {ukf.x[0]}')
|
|
ukf.update(np.array([pecanCount]))
|
|
print(f'Updated State: {ukf.x[0]}')
|
|
filtered_count = ukf.x[0]
|
|
print(filtered_count)
|
|
measuredThroughput = (filtered_count * avgPecanWeight) / (samplePeriod) * 15 # lbs / 15 seconds
|
|
print(f'Published Throughput: {measuredThroughput}')
|
|
client.publish(MQTT_COUNT_TOPIC, str(measuredThroughput))
|
|
pecanCount = 0
|
|
sampleStart = time.time()
|
|
print(f'Publish Time: {sampleStart-sampleEnd}')
|
|
|
|
# sampleEnd = time.time() # End Sample Timer
|
|
|
|
# samplePeriod = sampleEnd - sampleStart
|
|
|
|
# measuredThroughput = (filtered_count * avgPecanWeight) / (samplePeriod) * 15 # lbs / 15 seconds
|
|
|
|
#if (new_time := time.time()) > last_time + 0.25:
|
|
#client.publish(MQTT_COUNT_TOPIC, str(measuredThroughput))
|
|
#print(samplePeriod)
|
|
#print(str(last_time) + " " + str(new_time))
|
|
#last_time = new_time
|
|
|
|
if (time.time()-frameTime) > 10:
|
|
fps = frameCount / (time.time() - frameTime)
|
|
print(fps)
|
|
break
|
|
|
|
# Release the video capture object and close the display window
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|