Save current changes before reorganizing drivers and cleaning workspace

test_stream.py

@@ -80,6 +80,28 @@ def generate_animation_frame(frame_idx, mode_name):
     
     return img
 
+def connect_socket(esp32_ip, port, protocol):
+    """Establishes connection to the ESP32 stream server, retrying on failure for TCP."""
+    while True:
+        try:
+            if protocol == "tcp":
+                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+                sock.connect((esp32_ip, port))
+                print(f"Connected to stream server at {esp32_ip}:{port}")
+                return sock
+            else: # udp
+                sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+                if esp32_ip.endswith(".255") or esp32_ip == "255.255.255.255":
+                    sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
+                return sock
+        except (socket.error, Exception) as e:
+            if protocol == "tcp":
+                print(f"Failed to connect to stream server: {e}. Retrying in 2 seconds...")
+                time.sleep(2.0)
+            else:
+                # UDP doesn't establish connection, so any error here is fatal
+                raise e
+
 def stream_camera(esp32_ip, port, protocol):
     """Streams camera capture frames using OpenCV (requires opencv-python)."""
     try:
@@ -94,16 +116,8 @@ def stream_camera(esp32_ip, port, protocol):
         print("Error: Could not open camera.")
         sys.exit(1)
         
-    sock = None
-    if protocol == "tcp":
-        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        sock.connect((esp32_ip, port))
-    else: # udp
-        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-        # Enable broadcast if destination is broadcast
-        if esp32_ip.endswith(".255") or esp32_ip == "255.255.255.255":
-            sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
-            
+    sock = connect_socket(esp32_ip, port, protocol)
+        
     print(f"Streaming webcam to {esp32_ip}:{port} via {protocol.upper()}...")
     frame_count = 0
     start_time = time.time()
@@ -128,11 +142,19 @@ def stream_camera(esp32_ip, port, protocol):
             # Map pixels to hardware buffer
             raw_bytes = map_to_rlcd_hw_buffer(bg)
             
-            # Transmit
-            if protocol == "tcp":
-                sock.sendall(raw_bytes)
-            else:
-                send_udp_frame_chunked(sock, esp32_ip, port, frame_count, raw_bytes)
+            # Transmit with reconnect logic
+            try:
+                if protocol == "tcp":
+                    sock.sendall(raw_bytes)
+                else:
+                    send_udp_frame_chunked(sock, esp32_ip, port, frame_count, raw_bytes)
+            except (socket.error, Exception) as se:
+                print(f"Transmission failed: {se}. Reconnecting...")
+                try:
+                    sock.close()
+                except:
+                    pass
+                sock = connect_socket(esp32_ip, port, protocol)
                 
             frame_count += 1
             elapsed = time.time() - start_time
@@ -151,14 +173,7 @@ def stream_camera(esp32_ip, port, protocol):
 
 def stream_animation(esp32_ip, port, protocol):
     """Streams a generated Pillow vector animation."""
-    sock = None
-    if protocol == "tcp":
-        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        sock.connect((esp32_ip, port))
-    else: # udp
-        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-        if esp32_ip.endswith(".255") or esp32_ip == "255.255.255.255":
-            sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
+    sock = connect_socket(esp32_ip, port, protocol)
             
     print(f"Streaming vector animation to {esp32_ip}:{port} via {protocol.upper()}...")
     frame_idx = 0
@@ -172,11 +187,19 @@ def stream_animation(esp32_ip, port, protocol):
             # Map pixels to hardware buffer
             raw_bytes = map_to_rlcd_hw_buffer(pil_img)
             
-            # Transmit
-            if protocol == "tcp":
-                sock.sendall(raw_bytes)
-            else:
-                send_udp_frame_chunked(sock, esp32_ip, port, frame_idx, raw_bytes)
+            # Transmit with reconnect logic
+            try:
+                if protocol == "tcp":
+                    sock.sendall(raw_bytes)
+                else:
+                    send_udp_frame_chunked(sock, esp32_ip, port, frame_idx, raw_bytes)
+            except (socket.error, Exception) as se:
+                print(f"Transmission failed: {se}. Reconnecting...")
+                try:
+                    sock.close()
+                except:
+                    pass
+                sock = connect_socket(esp32_ip, port, protocol)
                 
             frame_idx += 1