nvidia-vfx-prototyping/samples/AiMaskStreamApp/video_stream.h

#pragma once

#include <vector>
#include <queue>
#include <mutex>
#include <thread>

#define UWS_NO_ZLIB
#include "App.h"

void Logger(const char* format, ...) {
    time_t now = time(0);
    tm* ltm = localtime(&now);
    char buffer[32];
    strftime(buffer, 32, "[%H:%M:%S] ", ltm);
    printf("%s", buffer);

    va_list args;
    va_start(args, format);
    vprintf(format, args);
    va_end(args);
}

enum StreamType {
    ST_ORIGINAL,
    ST_MASKED,
    ST_MAX
};

struct FrameQueue {
    std::queue<cv::Mat> frameQueue;
    std::mutex frameMutex;
};

std::vector<FrameQueue> frameQueues(ST_MAX);

struct FrameHolder {
    cv::Mat frame;
    std::mutex frameMutex;
};

std::array<FrameHolder, ST_MAX> frameHolders;

// Convert cv::Mat to JPEG buffer
std::vector<uchar> matToBytes(const cv::Mat& img)
{
    std::vector<uchar> buf;
    cv::imencode(".jpg", img, buf);
    return buf;
}

std::vector<uchar> matToJpeg(const cv::Mat& img)
{
    std::vector<uchar> buf;
    cv::imencode(".jpg", img, buf);
    return buf;
}

void updateFrame(int type, const cv::Mat& frame)
{
    std::lock_guard<std::mutex> lock(frameHolders[type].frameMutex);
    frameHolders[type].frame = frame.clone();
}

cv::Mat compressFrame(const cv::Mat& frame, int quality = 80)
{
    std::vector<uchar> buf;
    std::vector<int> params = { cv::IMWRITE_JPEG_QUALITY, quality };
    cv::imencode(".jpg", frame, buf, params);
    return cv::imdecode(buf, cv::IMREAD_COLOR);
}

void pushToFrameQueue(int type, const cv::Mat& frame) {
    //cv::Mat compressedFrame = compressFrame(frame, 50);

    // clear the queue
    std::queue<cv::Mat> empty;
    std::swap(frameQueues[type].frameQueue, empty);

    frameQueues[type].frameQueue.push(frame.clone());
}

struct PerSocketData {
};

std::mutex wsMutex;
std::vector<uWS::WebSocket<false, true, PerSocketData>*> activeWebSockets;

void SendFramesToAllClients() {
    std::lock_guard<std::mutex> lock(wsMutex);
    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 (!originalFrame.empty() && !maskedFrame.empty()) {
        std::vector<uchar> originalJpeg = matToJpeg(originalFrame);
        std::vector<uchar> maskedJpeg = matToJpeg(maskedFrame);

        uint32_t originalSize = originalJpeg.size();
        uint32_t maskedSize = maskedJpeg.size();

        // 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);
    }
}

void RunWebSocketServer(int websocketPort) {
    auto loop = uWS::Loop::get();

    struct us_timer_t* frameTimer = us_create_timer((struct us_loop_t*)loop, 0, 0);

    us_timer_set(frameTimer, [](struct us_timer_t*) {
        SendFramesToAllClients();
        }, 33, 33);  // 33ms interval for approximately 30 FPS

    uWS::App().ws<PerSocketData>("/*", {
        .compression = uWS::SHARED_COMPRESSOR,
        .maxPayloadLength = 100 * 1024 * 1024,
        .idleTimeout = 120,
        .maxBackpressure = 1 * 1024 * 1024,
        .closeOnBackpressureLimit = false,
        .resetIdleTimeoutOnSend = false,
        .sendPingsAutomatically = true,
        .open = [](auto* ws) {
		    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
        },
        .close = [](auto* ws, int code, std::string_view message) {
			Logger("WebSocket connection closed, current connections: %d\n", activeWebSockets.size() - 1);
            std::lock_guard<std::mutex> lock(wsMutex);
            activeWebSockets.erase(std::remove(activeWebSockets.begin(), activeWebSockets.end(), ws), activeWebSockets.end());
        }
        }).listen(websocketPort, [websocketPort](auto* listen_socket) {
            if (listen_socket) {
				Logger("WebSocket listening on port %i\n", websocketPort);
            }
            }).run();
}