diff --git a/dayglo_detector/dayglo_detector.py b/dayglo_detector/dayglo_detector.py index 0daeef4..c9f8160 100644 --- a/dayglo_detector/dayglo_detector.py +++ b/dayglo_detector/dayglo_detector.py @@ -99,3 +99,110 @@ 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) + print(f"Subscribed to topic: {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_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 +client.enable_logger() + +connect_mqtt() +client.loop_start() + +# Keep the script running +while True: + time.sleep(1) +