Handle fetching latest snapshot from MQTT topic upon connection

- Added `get_latest_snapshot()` function to subscribe to the MQTT topic for the latest snapshot image.
- Updated `process_latest_snapshot()` to utilize the `get_latest_snapshot()` function.
- Improved the logging to provide better information about snapshot subscription status.
This commit is contained in:
2024-10-20 16:31:57 +11:00
parent fe25895b38
commit be63a43613
+118
View File
@@ -110,3 +110,121 @@ client.on_message = on_message
connect_mqtt() connect_mqtt()
client.loop_forever() client.loop_forever()
import os
import time
import paho.mqtt.client as mqtt
import numpy as np
import cv2
import base64
import json
# Configuration
MQTT_BROKER = os.environ.get('MQTT_BROKER', '10.59.221.172')
MQTT_PORT = int(os.environ.get('MQTT_PORT', '1883'))
MQTT_USERNAME = os.getenv('MQTT_USERNAME', 'your_username')
MQTT_PASSWORD = os.getenv('MQTT_PASSWORD', 'your_password')
MQTT_SNAPSHOT_TOPIC = "/frigate/patiocam/person/snapshot"
MQTT_TOPIC_PUBLISH = "homeassistant/sensor/dayglo_rating/state"
DISCOVERY_PREFIX = "homeassistant"
# Default rating
last_rating = 0
# Default color thresholds for dayglo detection
LOWER_COLOR = np.array([20, 100, 100])
UPPER_COLOR = np.array([40, 255, 255])
def on_connect(client, userdata, flags, rc):
if rc == 0:
print("Connected successfully to MQTT broker")
client.subscribe(MQTT_SNAPSHOT_TOPIC)
publish_discovery_configurations()
# Publish initial rating of 0
publish_rating(0)
else:
print(f"Failed to connect, return code {rc}")
def publish_discovery_configurations():
rating_config = {
"name": "Dayglo Rating",
"state_topic": MQTT_TOPIC_PUBLISH,
"unit_of_measurement": "%",
"value_template": "{{ value_json.rating }}",
"icon": "mdi:brush",
"unique_id": "mqtt_dayglo_rating"
}
client.publish(f"{DISCOVERY_PREFIX}/sensor/dayglo_rating/config", json.dumps(rating_config), retain=True)
def publish_rating(rating):
global last_rating
last_rating = rating
client.publish(MQTT_TOPIC_PUBLISH, json.dumps({"rating": rating}))
def on_message(client, userdata, msg):
if msg.topic == MQTT_SNAPSHOT_TOPIC:
print("Snapshot received")
process_snapshot(msg.payload)
def process_latest_snapshot():
print("Attempting to process the latest snapshot...")
latest_snapshot_payload = base64.b64encode(get_latest_snapshot()).decode('utf-8')
if latest_snapshot_payload:
process_snapshot(latest_snapshot_payload.encode('utf-8'))
else:
print("No latest snapshot available to process.")
def get_latest_snapshot():
# Fetch the latest snapshot from MQTT topic
result = client.subscribe(MQTT_SNAPSHOT_TOPIC)
if result[0] == mqtt.MQTT_ERR_SUCCESS:
print("Subscribed to snapshot topic for latest image.")
else:
print("Failed to subscribe to snapshot topic.")
return b''
def process_snapshot(payload):
if not payload:
print("Empty payload received, skipping processing.")
return
print("Processing snapshot...")
try:
image_data = base64.b64decode(payload)
nparr = np.frombuffer(image_data, np.uint8)
image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
if image is not None:
rating = calculate_dayglo_rating(image)
print("Dayglo Rating calculated:", rating)
publish_rating(rating)
else:
print("Invalid image received")
except Exception as e:
print(f"Error processing snapshot: {e}")
def calculate_dayglo_rating(image):
print("Calculating dayglo rating...")
hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
mask = cv2.inRange(hsv_image, LOWER_COLOR, UPPER_COLOR)
dayglo_pixels = cv2.countNonZero(mask)
total_pixels = image.shape[0] * image.shape[1]
rating = (dayglo_pixels / total_pixels) * 100
return rating
def connect_mqtt():
while True:
try:
client.connect(MQTT_BROKER, MQTT_PORT, 60)
break
except Exception as e:
print(f"Connection failed: {e}. Retrying in 5 seconds...")
time.sleep(5)
client = mqtt.Client(protocol=mqtt.MQTTv311)
client.username_pw_set(MQTT_USERNAME, MQTT_PASSWORD)
client.on_connect = on_connect
client.on_message = on_message
connect_mqtt()
client.loop_forever()