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.
2024-10-20 16:31:57 +11:00
parent fe25895b38
commit be63a43613
1 changed files with 118 additions and 0 deletions
@@ -110,3 +110,121 @@ client.on_message = on_message
 connect_mqtt()
 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()
+