marcus
e9e5ffb931
- Save snapshot as a temporary file before decoding, allowing better handling of JFIF or non-standard JPEG formats. - Improved error handling for invalid or corrupted image data.
120 lines
3.7 KiB
Python
120 lines
3.7 KiB
Python
import os
|
|
import time
|
|
import paho.mqtt.client as mqtt
|
|
import numpy as np
|
|
import cv2
|
|
import base64
|
|
import json
|
|
import tempfile
|
|
|
|
# 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])
|
|
|
|
# Track if the initial snapshot has been processed
|
|
initial_snapshot_processed = False
|
|
|
|
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)
|
|
# Request the retained snapshot to process the initial image
|
|
client.publish(MQTT_SNAPSHOT_TOPIC, "", qos=1)
|
|
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):
|
|
global initial_snapshot_processed
|
|
if msg.topic == MQTT_SNAPSHOT_TOPIC:
|
|
print("Snapshot received")
|
|
process_snapshot(msg.payload)
|
|
initial_snapshot_processed = True
|
|
|
|
def process_snapshot(payload):
|
|
if not payload:
|
|
print("Empty payload received, skipping processing.")
|
|
return
|
|
|
|
print("Processing snapshot...")
|
|
try:
|
|
image_data = base64.b64decode(payload)
|
|
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_image_file:
|
|
temp_image_file.write(image_data)
|
|
temp_image_path = temp_image_file.name
|
|
|
|
# Attempt to read the saved image with OpenCV
|
|
image = cv2.imread(temp_image_path)
|
|
|
|
if image is not None:
|
|
rating = calculate_dayglo_rating(image)
|
|
print("Dayglo Rating calculated:", rating)
|
|
publish_rating(rating)
|
|
else:
|
|
print("Invalid image format or corrupted 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)
|
|
|