b-strike
/
nvidia-vfx-prototyping
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Log amount of data received, extra logging, add args for showing window/changing WS port

Browse Source
master
Benjamin Ruesink 2 months ago
parent 5532b45bbb
commit eb0026ae08
2 changed files with 47 additions and 184 deletions
Show Stats Download Patch File Download Diff File

208
samples/AiMaskStreamApp/AiMaskStreamApp.cpp
Unescape View File

@ -38,16 +38,9 @@
#include "opencv2/opencv.hpp" #include "opencv2/opencv.hpp"
#define UWS_NO_ZLIB #define UWS_NO_ZLIB
#include "httplib.h"
#include "App.h" #include "App.h"
//#include <uWebSockets/App.h>
#include <WinSock2.h>
#include <WS2tcpip.h>
#pragma comment(lib, "Ws2_32.lib") #pragma comment(lib, "Ws2_32.lib")
//#pragma comment(lib, "uSockets.lib")
enum StreamType { enum StreamType {
ST_ORIGINAL, ST_ORIGINAL,
@ -70,39 +63,42 @@ struct FrameHolder {
std::array<FrameHolder, ST_MAX> frameHolders; std::array<FrameHolder, ST_MAX> frameHolders;
// Convert cv::Mat to JPEG buffer // Convert cv::Mat to JPEG buffer
std::vector<uchar> matToBytes(const cv::Mat& img) { std::vector<uchar> matToBytes(const cv::Mat& img)
{
std::vector<uchar> buf; std::vector<uchar> buf;
cv::imencode(".jpg", img, buf); cv::imencode(".jpg", img, buf);
return buf; return buf;
} }
std::vector<uchar> matToJpeg(const cv::Mat& img) { std::vector<uchar> matToJpeg(const cv::Mat& img)
{
std::vector<uchar> buf; std::vector<uchar> buf;
cv::imencode(".jpg", img, buf); cv::imencode(".jpg", img, buf);
return buf; return buf;
} }
void updateFrame(int type, const cv::Mat& frame)
void updateFrame(int type, const cv::Mat& frame) { {
std::lock_guard<std::mutex> lock(frameHolders[type].frameMutex); std::lock_guard<std::mutex> lock(frameHolders[type].frameMutex);
frameHolders[type].frame = frame.clone(); 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) { void pushToFrameQueue(int type, const cv::Mat& frame) {
std::vector<uchar> buffer; //cv::Mat compressedFrame = compressFrame(frame, 50);
std::vector<int> params = { cv::IMWRITE_JPEG_QUALITY, 80 };
cv::imencode(".jpg", frame, buffer, params);
std::lock_guard<std::mutex> lock(frameQueues[type].frameMutex);
cv::Mat finalimg = cv::imdecode(buffer, cv::IMREAD_COLOR);
// clear the queue
std::queue<cv::Mat> empty; std::queue<cv::Mat> empty;
std::swap(frameQueues[type].frameQueue, empty); std::swap(frameQueues[type].frameQueue, empty);
frameQueues[type].frameQueue.push(finalimg.clone()); frameQueues[type].frameQueue.push(frame.clone());
} }
struct PerSocketData { struct PerSocketData {
@ -155,6 +151,8 @@ std::string FLAG_outDir;
std::string FLAG_outFile; std::string FLAG_outFile;
std::string FLAG_bgFile; std::string FLAG_bgFile;
int FLAG_camIndex = 0; int FLAG_camIndex = 0;
int FLAG_wsPort = 9001;
bool FLAG_showWindow = false;
static bool GetFlagArgVal(const char *flag, const char *arg, const char **val) { static bool GetFlagArgVal(const char *flag, const char *arg, const char **val) {
if (*arg != '-') return false; 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("cam_res", arg, &FLAG_camRes) || GetFlagArgVal("mode", arg, &FLAG_mode) ||
GetFlagArgVal("progress", arg, &FLAG_progress) || GetFlagArgVal("show", arg, &FLAG_show) || GetFlagArgVal("progress", arg, &FLAG_progress) || GetFlagArgVal("show", arg, &FLAG_show) ||
GetFlagArgVal("comp_mode", arg, &FLAG_compMode) || GetFlagArgVal("blur_strength", arg, &FLAG_blurStrength) || 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; continue;
} else if (GetFlagArgVal("help", arg, &help)) { } else if (GetFlagArgVal("help", arg, &help)) {
return NVCV_ERR_HELP; return NVCV_ERR_HELP;
@ -424,7 +423,6 @@ struct FXApp {
void destroyEffect(); void destroyEffect();
NvCV_Status allocBuffers(unsigned width, unsigned height); NvCV_Status allocBuffers(unsigned width, unsigned height);
NvCV_Status allocTempBuffers(); NvCV_Status allocTempBuffers();
Err processImage(const char *inFile, const char *outFile);
Err processMovie(const char *inFile, const char *outFile); Err processMovie(const char *inFile, const char *outFile);
Err processKey(int key); Err processKey(int key);
void nextCompMode(); void nextCompMode();
@ -533,30 +531,10 @@ FXApp::Err FXApp::processKey(int key) {
case 'q': case 'q':
case ESC_KEY: case ESC_KEY:
return errQuit; return errQuit;
case 'c':
case 'C':
nextCompMode();
break;
case 'f': case 'f':
case 'F': case 'F':
_showFPS = !_showFPS; _showFPS = !_showFPS;
break; 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; break;
default: default:
break; 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; 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) { FXApp::Err FXApp::processMovie(const char *inFile, const char *outFile) {
float ms = 0.0f; float ms = 0.0f;
FXApp::Err appErr = errNone; FXApp::Err appErr = errNone;
@ -1000,11 +906,14 @@ FXApp::Err FXApp::processMovie(const char *inFile, const char *outFile) {
if (_show) { if (_show) {
drawFrameRate(result); drawFrameRate(result);
static std::vector<int> params = { cv::IMWRITE_JPEG_QUALITY, 90 };
pushToFrameQueue(ST_ORIGINAL, originalImg); pushToFrameQueue(ST_ORIGINAL, originalImg);
pushToFrameQueue(ST_MASKED, result); pushToFrameQueue(ST_MASKED, result);
if (FLAG_showWindow) {
cv::imshow("Original", originalImg);
cv::imshow("Masked", result);
}
int key = cv::waitKey(1); int key = cv::waitKey(1);
if (key > 0) { if (key > 0) {
appErr = processKey(key); appErr = processKey(key);
@ -1043,57 +952,19 @@ int chooseGPU() {
bool isCompModeEnumValid(const FXApp::CompMode& mode) bool isCompModeEnumValid(const FXApp::CompMode& mode)
{ {
if (mode != FXApp::CompMode::compMatte && if (mode != FXApp::CompMode::compMatte &&
mode != FXApp::CompMode::compLight && mode != FXApp::CompMode::compLight &&
mode != FXApp::CompMode::compGreen && mode != FXApp::CompMode::compGreen &&
mode != FXApp::CompMode::compWhite && mode != FXApp::CompMode::compWhite &&
mode != FXApp::CompMode::compNone && mode != FXApp::CompMode::compNone &&
mode != FXApp::CompMode::compBG && mode != FXApp::CompMode::compBG &&
mode != FXApp::CompMode::compBlur) mode != FXApp::CompMode::compBlur)
{ {
return false; return false;
} }
return true; 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<std::mutex> lock(frameQueues[streamType].frameMutex);
if (!frameQueues[streamType].frameQueue.empty()) {
frame = frameQueues[streamType].frameQueue.front();
frameQueues[streamType].frameQueue.pop();
}
}
if (!frame.empty()) {
std::vector<uchar> 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<std::string, ST_MAX> 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() { void sendFramesToAllClients() {
std::lock_guard<std::mutex> lock(wsMutex); std::lock_guard<std::mutex> lock(wsMutex);
for (auto ws : activeWebSockets) { for (auto ws : activeWebSockets) {
@ -1122,8 +993,7 @@ void sendFramesToAllClients() {
} }
} }
void RunWebSocketServer() {
void runWebSocketServer() {
auto loop = uWS::Loop::get(); auto loop = uWS::Loop::get();
struct us_timer_t* frameTimer = us_create_timer((struct us_loop_t*)loop, 0, 0); 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) { .close = [](auto* ws, int code, std::string_view message) {
std::cout << "WebSocket connection closed" << std::endl; std::cout << "WebSocket connection closed" << std::endl;
std::cout << "Active connections: " << activeWebSockets.size() - 1 << std::endl;
std::lock_guard<std::mutex> lock(wsMutex); std::lock_guard<std::mutex> lock(wsMutex);
activeWebSockets.erase(std::remove(activeWebSockets.begin(), activeWebSockets.end(), ws), activeWebSockets.end()); 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) { if (listen_socket) {
std::cout << "Listening on port 9001" << std::endl; std::cout << "Listening on port " << FLAG_wsPort << std::endl;
} }
}).run(); }).run();
} }
@ -1169,7 +1040,7 @@ int main(int argc, char **argv) {
return nErrs; return nErrs;
} }
std::thread websocketThread(runWebSocketServer); std::thread websocketThread(RunWebSocketServer);
FXApp::Err fxErr = FXApp::errNone; FXApp::Err fxErr = FXApp::errNone;
FXApp app; FXApp app;
@ -1209,9 +1080,6 @@ int main(int argc, char **argv) {
} else { } else {
fxErr = app.appErrFromVfxStatus(app.createAigsEffect()); fxErr = app.appErrFromVfxStatus(app.createAigsEffect());
if (FXApp::errNone == fxErr) { 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()); fxErr = app.processMovie(FLAG_inFile.c_str(), FLAG_outFile.c_str());
if (fxErr == FXApp::errNone || fxErr == FXApp::errQuit) { if (fxErr == FXApp::errNone || fxErr == FXApp::errQuit) {
fxErr = FXApp::errNone; // Quitting isn't an error fxErr = FXApp::errNone; // Quitting isn't an error

23
streamWebApp/index.html
Unescape View File

@ -37,6 +37,7 @@
<div class="stats"> <div class="stats">
<h4>Messages received: <label class="message-count">0</label></h4> <h4>Messages received: <label class="message-count">0</label></h4>
<h4>FPS: <label class="fps">0</label></h4> <h4>FPS: <label class="fps">0</label></h4>
<h4>Total MB received: <label class="total-mb">0</label></h4>
</div> </div>
<script> <script>
let websocketUrl = 'ws://localhost:9001'; let websocketUrl = 'ws://localhost:9001';
@ -45,9 +46,10 @@
input.value = websocketUrl; input.value = websocketUrl;
const messageCountElement = document.querySelector('.message-count'); const messageCountElement = document.querySelector('.message-count');
const fpsElement = document.querySelector('.fps'); const fpsElement = document.querySelector('.fps');
const totalMBElement = document.querySelector('.total-mb');
let frameTimestamps = []; let frameTimestamps = [];
let lastFpsUpdateTime = 0; let lastFpsUpdateTime = 0;
let totalBytesReceived = 0;
const connect = () => { const connect = () => {
websocketUrl = input.value; websocketUrl = input.value;
@ -55,7 +57,6 @@
ws = new WebSocket(websocketUrl); ws = new WebSocket(websocketUrl);
ws.onopen = () => { ws.onopen = () => {
console.log('Connected to the server'); console.log('Connected to the server');
requestFrame();
}; };
ws.onmessage = (message) => { ws.onmessage = (message) => {
messageCountElement.textContent = Number(messageCountElement.textContent) + 1; messageCountElement.textContent = Number(messageCountElement.textContent) + 1;
@ -74,7 +75,7 @@
const arrayBuffer = await blob.arrayBuffer(); const arrayBuffer = await blob.arrayBuffer();
const dataView = new DataView(arrayBuffer); const dataView = new DataView(arrayBuffer);
const type = dataView.getUint8(0); const type = dataView.getUint8(0);
const size = dataView.getUint32(1, true); // true for little-endian const size = dataView.getUint32(1, true);
const imageData = arrayBuffer.slice(5, 5 + size); const imageData = arrayBuffer.slice(5, 5 + size);
const imageBlob = new Blob([imageData], { type: 'image/jpeg' }); const imageBlob = new Blob([imageData], { type: 'image/jpeg' });
const imageUrl = URL.createObjectURL(imageBlob); const imageUrl = URL.createObjectURL(imageBlob);
@ -83,30 +84,24 @@
URL.revokeObjectURL(streamElement.src); // Clean up the old object URL URL.revokeObjectURL(streamElement.src); // Clean up the old object URL
}; };
streamElement.src = imageUrl; streamElement.src = imageUrl;
// Update total MB received
totalBytesReceived += arrayBuffer.byteLength;
totalMBElement.textContent = (totalBytesReceived / (1024 * 1024)).toFixed(2);
updateFPS(); updateFPS();
// requestFrame();
} catch (error) { } catch (error) {
console.error('Error processing message:', error); console.error('Error processing message:', error);
// requestFrame();
} }
}; };
// const requestFrame = () => {
// if (ws && ws.readyState === WebSocket.OPEN) {
// ws.send('Next frame');
// }
// };
const updateFPS = () => { const updateFPS = () => {
const now = performance.now(); const now = performance.now();
frameTimestamps.push(now); frameTimestamps.push(now);
// Keep only the last 30 frame timestamps // Keep only the last 30 frame timestamps
if (frameTimestamps.length > 30) { if (frameTimestamps.length > 30) {
frameTimestamps.shift(); frameTimestamps.shift();
} }
// Update FPS // Update FPS
if (now - lastFpsUpdateTime > 100) { if (now - lastFpsUpdateTime > 100) {
const timeElapsed = (frameTimestamps[frameTimestamps.length - 1] - frameTimestamps[0]) / 1000; // in seconds const timeElapsed = (frameTimestamps[frameTimestamps.length - 1] - frameTimestamps[0]) / 1000; // in seconds

Write Preview
Loading…
Cancel
Save