From 05c9b1f713aba4f1bbdccc3e20ea9d23351b01c8 Mon Sep 17 00:00:00 2001
From: Marcus Brown <marcus@freedoh.net>
Date: Sun, 20 Oct 2024 16:25:27 +1100
Subject: [PATCH] improve

---
 dayglo_detector/dayglo_detector.py | 30 ++++++++++++++++++++++--------
 1 file changed, 22 insertions(+), 8 deletions(-)

diff --git a/dayglo_detector/dayglo_detector.py b/dayglo_detector/dayglo_detector.py
index 3c485ac..32d1c36 100644
--- a/dayglo_detector/dayglo_detector.py
+++ b/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)
-    client.subscribe(MQTT_SNAPSHOT_TOPIC)
-    publish_discovery_configurations()
-    # Publish initial rating of 0
-    publish_rating(0)
-    # Attempt to process the latest snapshot upon connection
-    process_latest_snapshot()
+    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)
+        # 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()