From 90a306622d67959673348831f8a9b5c1c61e84aa Mon Sep 17 00:00:00 2001
From: Benjamin Ruesink <benjamin.ruesink@gmail.com>
Date: Tue, 8 Oct 2024 15:55:26 +0200
Subject: [PATCH] Only create WS message once, instead of recreating it for
 every client, and put both images into 1 message

---
 samples/AiMaskStreamApp/video_stream.h | 53 ++++++++++--------
 streamWebApp/index.html                | 76 ++++++++++++++++++++------
 2 files changed, 89 insertions(+), 40 deletions(-)

diff --git a/samples/AiMaskStreamApp/video_stream.h b/samples/AiMaskStreamApp/video_stream.h
index be28b74..04d5124 100644
--- a/samples/AiMaskStreamApp/video_stream.h
+++ b/samples/AiMaskStreamApp/video_stream.h
@@ -86,32 +86,40 @@ struct PerSocketData {
 std::mutex wsMutex;
 std::vector<uWS::WebSocket<false, true, PerSocketData>*> activeWebSockets;
 
-
 void SendFramesToAllClients() {
     std::lock_guard<std::mutex> lock(wsMutex);
-    for (auto ws : activeWebSockets) {
-        for (int type = 0; type < ST_MAX; ++type) {
-            cv::Mat frame;
-            {
-                std::lock_guard<std::mutex> lock(frameQueues[type].frameMutex);
-                if (!frameQueues[type].frameQueue.empty()) {
-                    frame = frameQueues[type].frameQueue.front();
-                    frameQueues[type].frameQueue.pop();
-                }
-            }
+    std::string message;
+    cv::Mat originalFrame, maskedFrame;
+    {
+        std::lock_guard<std::mutex> lock(frameQueues[ST_ORIGINAL].frameMutex);
+        if (!frameQueues[ST_ORIGINAL].frameQueue.empty()) {
+            originalFrame = frameQueues[ST_ORIGINAL].frameQueue.front();
+            frameQueues[ST_ORIGINAL].frameQueue.pop();
+        }
+    }
+    {
+        std::lock_guard<std::mutex> lock(frameQueues[ST_MASKED].frameMutex);
+        if (!frameQueues[ST_MASKED].frameQueue.empty()) {
+            maskedFrame = frameQueues[ST_MASKED].frameQueue.front();
+            frameQueues[ST_MASKED].frameQueue.pop();
+        }
+    }
 
-            if (!frame.empty()) {
-                std::vector<uchar> jpeg = matToJpeg(frame);
-                uint32_t size = jpeg.size();
+    if (!originalFrame.empty() && !maskedFrame.empty()) {
+        std::vector<uchar> originalJpeg = matToJpeg(originalFrame);
+        std::vector<uchar> maskedJpeg = matToJpeg(maskedFrame);
 
-                // Prepare the message: type (1 byte) + size (4 bytes) + image data
-                std::string message(1, static_cast<char>(type));
-                message.append(reinterpret_cast<char*>(&size), 4);
-                message.append(jpeg.begin(), jpeg.end());
+        uint32_t originalSize = originalJpeg.size();
+        uint32_t maskedSize = maskedJpeg.size();
 
-                ws->send(message, uWS::OpCode::BINARY);
-            }
-        }
+        // Prepare the message: originalSize (4 bytes) + maskedSize (4 bytes) + original image data + masked image data
+        message.append(reinterpret_cast<char*>(&originalSize), 4);
+        message.append(reinterpret_cast<char*>(&maskedSize), 4);
+        message.append(originalJpeg.begin(), originalJpeg.end());
+        message.append(maskedJpeg.begin(), maskedJpeg.end());
+    }
+    for (auto ws : activeWebSockets) {
+        ws->send(message, uWS::OpCode::BINARY);
     }
 }
 
@@ -133,9 +141,10 @@ void RunWebSocketServer(int websocketPort) {
         .resetIdleTimeoutOnSend = false,
         .sendPingsAutomatically = true,
         .open = [](auto* ws) {
-		    Logger("WebSocket connection opened\n");
+		    Logger("WebSocket connection opened, current connections: %d\n", activeWebSockets.size() + 1);
             std::lock_guard<std::mutex> lock(wsMutex);
             activeWebSockets.push_back(ws);
+            Logger("Client address: %s\n", ws->getRemoteAddressAsText().data());
         },
         .message = [](auto* ws, std::string_view message, uWS::OpCode opCode) {
             // Handle incoming messages if needed
diff --git a/streamWebApp/index.html b/streamWebApp/index.html
index 8b574e2..4581696 100644
--- a/streamWebApp/index.html
+++ b/streamWebApp/index.html
@@ -12,11 +12,20 @@
     .streams {
       display: flex;
       justify-content: center;
-      background: #ddd;
+      position: relative;
+      width: 100%;
     }
     .streams img {
-      margin: 50px;
       outline: 1px solid #000;
+      position: absolute;
+      top: 0;
+      left: 50%;
+      transform: translateX(-50%);
+      height: 450px;
+      background: #ddd;
+    }
+    .stream-2 {
+      mix-blend-mode: darken;
     }
     .stats {
       display: flex;
@@ -24,21 +33,30 @@
       width: 100%;
       margin-top: 20px;
     }
+    .connect-info {
+      display: flex;
+      justify-content: center;
+      align-items: center;
+      margin-top: 20px;
+    }
   </style>
 </head>
 <body>
   <h1>Stream Web App</h1>
-  <button onclick="connect()">Connect</button>
-  <input type="text" placeholder="Enter the WebSocket URL" />
-  <div class="streams">
-    <img class="stream-1" src="" alt="Original Stream">
-    <img class="stream-2" src="" alt="Masked Stream">
+  <div class="connect-info">
+    <input type="text" placeholder="Enter the WebSocket URL" />
+    <button onclick="connect()">Connect</button>
+    <button onclick="disconnect()">Disconnect</button>
   </div>
   <div class="stats">
     <h4>Messages received: <label class="message-count">0</label></h4>
     <h4>FPS: <label class="fps">0</label></h4>
     <h4>Total MB received: <label class="total-mb">0</label></h4>
   </div>
+  <div class="streams">
+    <img class="stream-1" src="" alt="Stream">
+    <img class="stream-2" src="" alt="">
+  </div>
   <script>
     let websocketUrl = 'ws://localhost:9001';
     let ws;
@@ -51,7 +69,16 @@
     let lastFpsUpdateTime = 0;
     let totalBytesReceived = 0;
 
+    const disconnect = () => {
+      if (ws) {
+        ws.close();
+        ws = null;
+      }
+    };
+
     const connect = () => {
+      disconnect();
+
       websocketUrl = input.value;
       console.log('Connecting to the server...', websocketUrl);
       ws = new WebSocket(websocketUrl);
@@ -74,21 +101,34 @@
       try {
         const arrayBuffer = await blob.arrayBuffer();
         const dataView = new DataView(arrayBuffer);
-        const type = dataView.getUint8(0);
-        const size = dataView.getUint32(1, true);
-        const imageData = arrayBuffer.slice(5, 5 + size);
-        const imageBlob = new Blob([imageData], { type: 'image/jpeg' });
-        const imageUrl = URL.createObjectURL(imageBlob);
-        const streamElement = document.querySelector(type === 0 ? '.stream-1' : '.stream-2');
-        streamElement.onload = () => {
-          URL.revokeObjectURL(streamElement.src); // Clean up the old object URL
+        const originalSize = dataView.getUint32(0, true);
+        const maskedSize = dataView.getUint32(4, true);
+        
+        const originalImageData = arrayBuffer.slice(8, 8 + originalSize);
+        const maskedImageData = arrayBuffer.slice(8 + originalSize, 8 + originalSize + maskedSize);
+        
+        const originalImageBlob = new Blob([originalImageData], { type: 'image/jpeg' });
+        const maskedImageBlob = new Blob([maskedImageData], { type: 'image/jpeg' });
+        
+        const originalImageUrl = URL.createObjectURL(originalImageBlob);
+        const maskedImageUrl = URL.createObjectURL(maskedImageBlob);
+        
+        const originalStreamElement = document.querySelector('.stream-1');
+        const maskedStreamElement = document.querySelector('.stream-2');
+        
+        originalStreamElement.onload = () => {
+          URL.revokeObjectURL(originalStreamElement.src);
         };
-        streamElement.src = imageUrl;
-
+        maskedStreamElement.onload = () => {
+          URL.revokeObjectURL(maskedStreamElement.src);
+        };
+        
+        originalStreamElement.src = originalImageUrl;
+        maskedStreamElement.src = maskedImageUrl;
+        
         // Update total MB received
         totalBytesReceived += arrayBuffer.byteLength;
         totalMBElement.textContent = (totalBytesReceived / (1024 * 1024)).toFixed(2);
-
         updateFPS();
       } catch (error) {
         console.error('Error processing message:', error);