diff --git a/samples/AiMaskStreamApp/AiMaskStreamApp.cpp b/samples/AiMaskStreamApp/AiMaskStreamApp.cpp index 11e74af..dda4500 100644 --- a/samples/AiMaskStreamApp/AiMaskStreamApp.cpp +++ b/samples/AiMaskStreamApp/AiMaskStreamApp.cpp @@ -38,16 +38,9 @@ #include "opencv2/opencv.hpp" #define UWS_NO_ZLIB - -#include "httplib.h" #include "App.h" -//#include - -#include -#include #pragma comment(lib, "Ws2_32.lib") -//#pragma comment(lib, "uSockets.lib") enum StreamType { ST_ORIGINAL, @@ -70,39 +63,42 @@ struct FrameHolder { std::array frameHolders; // Convert cv::Mat to JPEG buffer -std::vector matToBytes(const cv::Mat& img) { - +std::vector matToBytes(const cv::Mat& img) +{ std::vector buf; cv::imencode(".jpg", img, buf); return buf; } -std::vector matToJpeg(const cv::Mat& img) { +std::vector matToJpeg(const cv::Mat& img) +{ std::vector buf; cv::imencode(".jpg", img, buf); return buf; } - -void updateFrame(int type, const cv::Mat& frame) { +void updateFrame(int type, const cv::Mat& frame) +{ std::lock_guard lock(frameHolders[type].frameMutex); frameHolders[type].frame = frame.clone(); } +cv::Mat compressFrame(const cv::Mat& frame, int quality = 80) +{ + std::vector buf; + std::vector 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) { - std::vector buffer; - std::vector params = { cv::IMWRITE_JPEG_QUALITY, 80 }; - cv::imencode(".jpg", frame, buffer, params); - - std::lock_guard lock(frameQueues[type].frameMutex); - - cv::Mat finalimg = cv::imdecode(buffer, cv::IMREAD_COLOR); + //cv::Mat compressedFrame = compressFrame(frame, 50); + // clear the queue std::queue empty; std::swap(frameQueues[type].frameQueue, empty); - frameQueues[type].frameQueue.push(finalimg.clone()); + frameQueues[type].frameQueue.push(frame.clone()); } struct PerSocketData { @@ -155,6 +151,8 @@ std::string FLAG_outDir; std::string FLAG_outFile; std::string FLAG_bgFile; int FLAG_camIndex = 0; +int FLAG_wsPort = 9001; +bool FLAG_showWindow = false; static bool GetFlagArgVal(const char *flag, const char *arg, const char **val) { if (*arg != '-') return false; @@ -254,7 +252,8 @@ static int ParseMyArgs(int argc, char **argv) { GetFlagArgVal("cam_res", arg, &FLAG_camRes) || GetFlagArgVal("mode", arg, &FLAG_mode) || GetFlagArgVal("progress", arg, &FLAG_progress) || GetFlagArgVal("show", arg, &FLAG_show) || GetFlagArgVal("comp_mode", arg, &FLAG_compMode) || GetFlagArgVal("blur_strength", arg, &FLAG_blurStrength) || - GetFlagArgVal("cuda_graph", arg, &FLAG_cudaGraph) || GetFlagArgVal("cam_index", arg, &FLAG_camIndex))) { + GetFlagArgVal("cuda_graph", arg, &FLAG_cudaGraph) || GetFlagArgVal("cam_index", arg, &FLAG_camIndex) || + GetFlagArgVal("ws_port", arg, &FLAG_wsPort) || GetFlagArgVal("show_window", arg, &FLAG_showWindow))) { continue; } else if (GetFlagArgVal("help", arg, &help)) { return NVCV_ERR_HELP; @@ -424,7 +423,6 @@ struct FXApp { void destroyEffect(); NvCV_Status allocBuffers(unsigned width, unsigned height); NvCV_Status allocTempBuffers(); - Err processImage(const char *inFile, const char *outFile); Err processMovie(const char *inFile, const char *outFile); Err processKey(int key); void nextCompMode(); @@ -533,30 +531,10 @@ FXApp::Err FXApp::processKey(int key) { case 'q': case ESC_KEY: return errQuit; - case 'c': - case 'C': - nextCompMode(); - break; case 'f': case 'F': _showFPS = !_showFPS; break; - case 'p': - case 'P': - case '%': - _progress = !_progress; - break; - case 'm': - _blurStrength += 0.05f; - if (_blurStrength > 1.0) { - _blurStrength = 1.0; - } - break; - case 'n': - _blurStrength -= 0.05f; - if (_blurStrength < 0.0) { - _blurStrength = 0.0; - } break; default: break; @@ -712,78 +690,6 @@ static NvCV_Status WriteRGBA(const NvCVImage *bgr, const NvCVImage *a, const std return cv::imwrite(name, ocv) ? NVCV_SUCCESS : NVCV_ERR_WRITE; } -FXApp::Err FXApp::processImage(const char *inFile, const char *outFile) { - NvCV_Status vfxErr; - bool ok; - cv::Mat result; - NvCVImage fxSrcChunkyGPU, fxDstChunkyGPU; - - // Allocate space for batchOfStates to hold state variable addresses - // Assume that MODEL_BATCH Size is enough for this scenario - unsigned int modelBatch = 1; - BAIL_IF_ERR(vfxErr = NvVFX_GetU32(_eff, NVVFX_MODEL_BATCH, &modelBatch)); - _batchOfStates = (NvVFX_StateObjectHandle*) malloc(sizeof(NvVFX_StateObjectHandle) * modelBatch); - if (_batchOfStates == nullptr) { - vfxErr = NVCV_ERR_MEMORY; - goto bail; - } - - if (!_eff) return errEffect; - _srcImg = cv::imread(inFile); - if (!_srcImg.data) return errRead; - - _dstImg = cv::Mat::zeros(_srcImg.size(), CV_8UC1); - if (!_dstImg.data) return errMemory; - - (void)NVWrapperForCVMat(&_srcImg, &_srcVFX); - (void)NVWrapperForCVMat(&_dstImg, &_dstVFX); - - if (!fxSrcChunkyGPU.pixels) - { - BAIL_IF_ERR(vfxErr = - NvCVImage_Alloc(&fxSrcChunkyGPU, _srcImg.cols, _srcImg.rows, NVCV_BGR, NVCV_U8, NVCV_CHUNKY, NVCV_GPU, 1)); - } - - if (!fxDstChunkyGPU.pixels) { - BAIL_IF_ERR(vfxErr = - NvCVImage_Alloc(&fxDstChunkyGPU, _srcImg.cols, _srcImg.rows, NVCV_A, NVCV_U8, NVCV_CHUNKY, NVCV_GPU, 1)); - } - - BAIL_IF_ERR(vfxErr = NvVFX_SetImage(_eff, NVVFX_INPUT_IMAGE, &fxSrcChunkyGPU)); - BAIL_IF_ERR(vfxErr = NvVFX_SetImage(_eff, NVVFX_OUTPUT_IMAGE, &fxDstChunkyGPU)); - BAIL_IF_ERR(vfxErr = NvCVImage_Transfer(&_srcVFX, &fxSrcChunkyGPU, 1.0f, _stream, NULL)); - - // Assign states from stateArray in batchOfStates - // There is only one stream in this app - _batchOfStates[0] = _stateArray[0]; - BAIL_IF_ERR(vfxErr = NvVFX_SetStateObjectHandleArray(_eff, NVVFX_STATE, _batchOfStates)); - - BAIL_IF_ERR(vfxErr = NvVFX_Run(_eff, 0)); - BAIL_IF_ERR(vfxErr = NvCVImage_Transfer(&fxDstChunkyGPU, &_dstVFX, 1.0f, _stream, NULL)); - - overlay(_srcImg, _dstImg, 0.5, result); - if (!std::string(outFile).empty()) { - if(IsLossyImageFile(outFile)) - fprintf(stderr, "WARNING: JPEG output file format will reduce image quality\n"); - vfxErr = WriteRGBA(&_srcVFX, &_dstVFX, outFile); - if (NVCV_SUCCESS != vfxErr) { - printf("%s: \"%s\"\n", NvCV_GetErrorStringFromCode(vfxErr), outFile); - goto bail; - } - ok = cv::imwrite(std::string(outFile) + "_segmentation_mask.png", _dstImg); // save segmentation mask too - if (!ok) { - printf("Error writing: \"%s_segmentation_mask.png\"\n", outFile); - return errWrite; - } - } - if (_show) { - cv::imshow("Output", result); - cv::waitKey(3000); - } -bail: - return (FXApp::Err)vfxErr; -} - FXApp::Err FXApp::processMovie(const char *inFile, const char *outFile) { float ms = 0.0f; FXApp::Err appErr = errNone; @@ -1000,11 +906,14 @@ FXApp::Err FXApp::processMovie(const char *inFile, const char *outFile) { if (_show) { drawFrameRate(result); - static std::vector params = { cv::IMWRITE_JPEG_QUALITY, 90 }; - pushToFrameQueue(ST_ORIGINAL, originalImg); pushToFrameQueue(ST_MASKED, result); + if (FLAG_showWindow) { + cv::imshow("Original", originalImg); + cv::imshow("Masked", result); + } + int key = cv::waitKey(1); if (key > 0) { appErr = processKey(key); @@ -1043,57 +952,19 @@ int chooseGPU() { bool isCompModeEnumValid(const FXApp::CompMode& mode) { - if (mode != FXApp::CompMode::compMatte && - mode != FXApp::CompMode::compLight && - mode != FXApp::CompMode::compGreen && - mode != FXApp::CompMode::compWhite && - mode != FXApp::CompMode::compNone && - mode != FXApp::CompMode::compBG && - mode != FXApp::CompMode::compBlur) + if (mode != FXApp::CompMode::compMatte && + mode != FXApp::CompMode::compLight && + mode != FXApp::CompMode::compGreen && + mode != FXApp::CompMode::compWhite && + mode != FXApp::CompMode::compNone && + mode != FXApp::CompMode::compBG && + mode != FXApp::CompMode::compBlur) { - return false; + return false; } return true; } -void startHttpServer() { - httplib::Server svr; - - auto frameHandler = [](StreamType streamType) { - return [streamType](const httplib::Request&, httplib::Response& res) { - cv::Mat frame; - { - std::lock_guard lock(frameQueues[streamType].frameMutex); - if (!frameQueues[streamType].frameQueue.empty()) { - frame = frameQueues[streamType].frameQueue.front(); - frameQueues[streamType].frameQueue.pop(); - } - } - if (!frame.empty()) { - std::vector jpegData = matToBytes(frame); - res.set_content(std::string(jpegData.begin(), jpegData.end()), "image/jpeg"); - } - else { - res.status = 404; - res.set_content("No frame available", "text/plain"); - } - }; - }; - - const std::array routeNames = { - "original", "masked" - }; - - for (int i = 0; i < ST_MAX; ++i) { - std::string route = "/video_" + routeNames[i]; - svr.Get(route, frameHandler((StreamType)i)); - printf("Serving %s frames at http://localhost:8080%s\n", routeNames[i].c_str(), route.c_str()); - } - - svr.listen("0.0.0.0", 8090); -} - - void sendFramesToAllClients() { std::lock_guard lock(wsMutex); for (auto ws : activeWebSockets) { @@ -1122,8 +993,7 @@ void sendFramesToAllClients() { } } - -void runWebSocketServer() { +void RunWebSocketServer() { auto loop = uWS::Loop::get(); struct us_timer_t* frameTimer = us_create_timer((struct us_loop_t*)loop, 0, 0); @@ -1150,12 +1020,13 @@ void runWebSocketServer() { }, .close = [](auto* ws, int code, std::string_view message) { std::cout << "WebSocket connection closed" << std::endl; + std::cout << "Active connections: " << activeWebSockets.size() - 1 << std::endl; std::lock_guard lock(wsMutex); activeWebSockets.erase(std::remove(activeWebSockets.begin(), activeWebSockets.end(), ws), activeWebSockets.end()); } - }).listen(9001, [](auto* listen_socket) { + }).listen(FLAG_wsPort, [](auto* listen_socket) { if (listen_socket) { - std::cout << "Listening on port 9001" << std::endl; + std::cout << "Listening on port " << FLAG_wsPort << std::endl; } }).run(); } @@ -1169,7 +1040,7 @@ int main(int argc, char **argv) { return nErrs; } - std::thread websocketThread(runWebSocketServer); + std::thread websocketThread(RunWebSocketServer); FXApp::Err fxErr = FXApp::errNone; FXApp app; @@ -1209,9 +1080,6 @@ int main(int argc, char **argv) { } else { fxErr = app.appErrFromVfxStatus(app.createAigsEffect()); if (FXApp::errNone == fxErr) { - if (IsImageFile(FLAG_inFile.c_str())) - fxErr = app.processImage(FLAG_inFile.c_str(), FLAG_outFile.c_str()); - else fxErr = app.processMovie(FLAG_inFile.c_str(), FLAG_outFile.c_str()); if (fxErr == FXApp::errNone || fxErr == FXApp::errQuit) { fxErr = FXApp::errNone; // Quitting isn't an error diff --git a/streamWebApp/index.html b/streamWebApp/index.html index 801e0d4..8b574e2 100644 --- a/streamWebApp/index.html +++ b/streamWebApp/index.html @@ -37,6 +37,7 @@

Messages received:

FPS:

+

Total MB received: