improve

dayglo_detector/dayglo_detector.py

@@ -1,4 +1,5 @@
 import os
+import time
 import paho.mqtt.client as mqtt
 import numpy as np
 import cv2
@@ -22,13 +23,16 @@ LOWER_COLOR = np.array([20, 100, 100])
 UPPER_COLOR = np.array([40, 255, 255])
 
 def on_connect(client, userdata, flags, rc):
-    print("Connected with result code", rc)
+    if rc == 0:
-    client.subscribe(MQTT_SNAPSHOT_TOPIC)
+        print("Connected successfully to MQTT broker")
-    publish_discovery_configurations()
+        client.subscribe(MQTT_SNAPSHOT_TOPIC)
-    # Publish initial rating of 0
+        publish_discovery_configurations()
-    publish_rating(0)
+        # Publish initial rating of 0
-    # Attempt to process the latest snapshot upon connection
+        publish_rating(0)
-    process_latest_snapshot()
+        # Attempt to process the latest snapshot upon connection
+        process_latest_snapshot()
+    else:
+        print(f"Failed to connect, return code {rc}")
 
 def publish_discovery_configurations():
     rating_config = {
@@ -79,10 +83,20 @@ def calculate_dayglo_rating(image):
     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.connect(MQTT_BROKER, MQTT_PORT, 60)
+
+connect_mqtt()
 client.loop_forever()