From fbf8f8de59b7a28cd4fd65724cab2b9c6cb02e94 Mon Sep 17 00:00:00 2001
From: Marcus <marcus.c.eliasson@gmail.com>
Date: Tue, 16 Mar 2021 15:18:02 +0100
Subject: [PATCH] Improved WiFi streamer example (#40)

---
 .../main/ai-deck-jpeg-streamer-demo.c         |  4 +--
 NINA/viewer.py                                | 31 +++++++++++--------
 2 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/NINA/firmware/main/ai-deck-jpeg-streamer-demo.c b/NINA/firmware/main/ai-deck-jpeg-streamer-demo.c
index d48c2f5e..f0ff8055 100644
--- a/NINA/firmware/main/ai-deck-jpeg-streamer-demo.c
+++ b/NINA/firmware/main/ai-deck-jpeg-streamer-demo.c
@@ -132,12 +132,12 @@ static void send_imagedata_to_client(uint32_t type, uint32_t size, uint32_t info
   if (wifi_is_socket_connected())
   {
 
+    wifi_send_packet( (const char*) buffer, size);
+
     if (new_frame) {
       wifi_send_packet( (const char*) &jpeg_footer, sizeof(jpeg_footer) );
     }
 
-    wifi_send_packet( (const char*) buffer, size);
-
     if (new_frame) {
       wifi_send_packet( (const char*) &jpeg_header, sizeof(jpeg_header) );
     }
diff --git a/NINA/viewer.py b/NINA/viewer.py
index 1bf5188e..82fe9f48 100644
--- a/NINA/viewer.py
+++ b/NINA/viewer.py
@@ -60,26 +60,31 @@ def run(self):
         print("Socket connected")
 
         imgdata = None
+        data_buffer = bytearray()
 
         while(1):
-            strng = client_socket.recv(512)
+            # Reveive image data from the AI-deck
+            data_buffer.extend(client_socket.recv(512))
 
             # Look for start-of-frame and end-of-frame
-            start_idx = strng.find(b"\xff\xd8")
-            end_idx = strng.find(b"\xff\xd9")
-
-            # Concatenate image data, once finished send it to the UI
-            if start_idx >= 0:
-                imgdata = strng[start_idx:len(strng)] # Does not support start and end in same package
-            elif end_idx >= 0 and imgdata:
-                imgdata += strng[0:end_idx]
+            start_idx = data_buffer.find(b"\xff\xd8")
+            end_idx = data_buffer.find(b"\xff\xd9")
+
+            # At startup we might get an end before we get the first start, if
+            # that is the case then throw away the data before start
+            if end_idx > -1 and end_idx < start_idx:
+                data_buffer = data_buffer[start_idx:]
+
+            # We have a start and an end of the image in the buffer now
+            if start_idx > -1 and end_idx > -1 and end_idx > start_idx:
+                # Pick out the image to render ...
+                imgdata = data_buffer[start_idx:end_idx + 2]
+                # .. and remove it from the buffer
+                data_buffer = data_buffer[end_idx + 2 :]
                 try:
                     self._callback(imgdata)
                 except gi.repository.GLib.Error:
-                    pass
-                imgdata = strng[end_idx:len(strng)]
-            elif imgdata:
-                imgdata += strng
+                    print("Error rendering image")
 
 # UI for showing frames from AI-deck example
 class FrameViewer(Gtk.Window):