Merge pull request #68 from z65536/main

Fix: Corrupted JPEG

Author: SummerSigh

Diff for: BabbleApp/camera.py

@@ -12,6 +12,8 @@
 from lang_manager import LocaleStringManager as lang
 from colorama import Fore
 from config import BabbleConfig, BabbleSettingsConfig
+from PIL import Image
+from io import BytesIO
 from utils.misc_utils import get_camera_index_by_name, list_camera_names, os_type
 
 from vivefacialtracker.vivetracker import ViveTracker
@@ -245,20 +247,21 @@ def get_next_packet_bounds(self):
         beg = -1
         while beg == -1:
             self.buffer += self.serial_connection.read(2048)
-            beg = self.buffer.find(ETVR_HEADER)
+            beg = self.buffer.find(b"\xff\xd8\xff")
         # Discard any data before the frame header.
         if beg > 0:
             self.buffer = self.buffer[beg:]
             beg = 0
-        # We know exactly how long the jpeg packet is
-        end = int.from_bytes(self.buffer[4:6], signed=False, byteorder="little")
-        self.buffer += self.serial_connection.read(end - len(self.buffer))
+        end = -1
+        while end == -1:
+            self.buffer += self.serial_connection.read(128)
+            end = self.buffer.find(b"\xff\xd9")
         return beg, end
 
     def get_next_jpeg_frame(self):
         beg, end = self.get_next_packet_bounds()
-        jpeg = self.buffer[beg + ETVR_HEADER_LEN : end + ETVR_HEADER_LEN]
-        self.buffer = self.buffer[end + ETVR_HEADER_LEN :]
+        jpeg = self.buffer[beg: end + 2]
+        self.buffer = self.buffer[end + 2 :]
         return jpeg
 
     def get_serial_camera_picture(self, should_push):
@@ -271,10 +274,9 @@ def get_serial_camera_picture(self, should_push):
                 jpeg = self.get_next_jpeg_frame()
                 if jpeg:
                     # Create jpeg frame from byte string
-                    image = cv2.imdecode(
-                        np.fromstring(jpeg, dtype=np.uint8), cv2.IMREAD_UNCHANGED
-                    )
-                    if image is None:
+                    try:
+                        image = np.array(Image.open(BytesIO(jpeg)))
+                    except Exception:
                         print(
                             f'{Fore.YELLOW}[{lang._instance.get_string("log.warn")}] {lang._instance.get_string("warn.frameDrop")}{Fore.RESET}'
                         )