You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
371 lines
13 KiB
371 lines
13 KiB
/*###############################################################################
|
|
#
|
|
# Copyright (c) 2022 NVIDIA Corporation
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy of
|
|
# this software and associated documentation files (the "Software"), to deal in
|
|
# the Software without restriction, including without limitation the rights to
|
|
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
|
|
# the Software, and to permit persons to whom the Software is furnished to do so,
|
|
# subject to the following conditions:
|
|
#
|
|
# The above copyright notice and this permission notice shall be included in all
|
|
# copies or substantial portions of the Software.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
|
|
# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
|
|
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
|
|
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
|
|
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
|
#
|
|
###############################################################################*/
|
|
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include <cuda_runtime_api.h>
|
|
#include "BatchUtilities.h"
|
|
#include "nvCVOpenCV.h"
|
|
#include "nvVideoEffects.h"
|
|
#include "opencv2/opencv.hpp"
|
|
|
|
#ifdef _MSC_VER
|
|
#define strcasecmp _stricmp
|
|
#endif // _MSC_VER
|
|
|
|
#define BAIL_IF_ERR(err) do { if (0 != (err)) { goto bail; } } while(0)
|
|
#define BAIL_IF_NULL(x, err, code) do { if ((void*)(x) == NULL) { err = code; goto bail; } } while(0)
|
|
#define BAIL_IF_FALSE(x, err, code) do { if (!(x)) { err = code; goto bail; } } while(0)
|
|
#define BAIL(err, code) do { err = code; goto bail; } while(0)
|
|
|
|
#ifdef _WIN32
|
|
#define DEFAULT_CODEC "avc1"
|
|
#else // !_WIN32
|
|
#define DEFAULT_CODEC "H264"
|
|
#endif // _WIN32
|
|
|
|
bool FLAG_verbose = false;
|
|
int FLAG_mode = 0;
|
|
std::string FLAG_outFile,
|
|
FLAG_modelDir,
|
|
FLAG_codec = DEFAULT_CODEC;
|
|
std::vector<const char*> FLAG_inFiles;
|
|
|
|
// Set this when using OTA Updates
|
|
// This path is used by nvVideoEffectsProxy.cpp to load the SDK dll
|
|
// when using OTA Updates
|
|
char *g_nvVFXSDKPath = NULL;
|
|
|
|
static bool GetFlagArgVal(const char *flag, const char *arg, const char **val) {
|
|
if (*arg != '-')
|
|
return false;
|
|
while (*++arg == '-')
|
|
continue;
|
|
const char *s = strchr(arg, '=');
|
|
if (s == NULL) {
|
|
if (strcmp(flag, arg) != 0)
|
|
return false;
|
|
*val = NULL;
|
|
return true;
|
|
}
|
|
size_t n = s - arg;
|
|
if ((strlen(flag) != n) || (strncmp(flag, arg, n) != 0))
|
|
return false;
|
|
*val = s + 1;
|
|
return true;
|
|
}
|
|
|
|
static bool GetFlagArgVal(const char *flag, const char *arg, std::string *val) {
|
|
const char *valStr;
|
|
if (!GetFlagArgVal(flag, arg, &valStr))
|
|
return false;
|
|
val->assign(valStr ? valStr : "");
|
|
return true;
|
|
}
|
|
|
|
static bool GetFlagArgVal(const char *flag, const char *arg, bool *val) {
|
|
const char *valStr;
|
|
bool success = GetFlagArgVal(flag, arg, &valStr);
|
|
if (success) {
|
|
*val = (valStr == NULL ||
|
|
strcasecmp(valStr, "true") == 0 ||
|
|
strcasecmp(valStr, "on") == 0 ||
|
|
strcasecmp(valStr, "yes") == 0 ||
|
|
strcasecmp(valStr, "1") == 0
|
|
);
|
|
}
|
|
return success;
|
|
}
|
|
|
|
static bool GetFlagArgVal(const char *flag, const char *arg, long *val) {
|
|
const char *valStr;
|
|
bool success = GetFlagArgVal(flag, arg, &valStr);
|
|
if (success)
|
|
*val = strtol(valStr, NULL, 10);
|
|
return success;
|
|
}
|
|
|
|
static bool GetFlagArgVal(const char *flag, const char *arg, int *val) {
|
|
long longVal;
|
|
bool success = GetFlagArgVal(flag, arg, &longVal);
|
|
if (success)
|
|
*val = (int)longVal;
|
|
return success;
|
|
}
|
|
|
|
static int StringToFourcc(const std::string &str) {
|
|
union chint {
|
|
int i;
|
|
char c[4];
|
|
};
|
|
chint x = {0};
|
|
for (int n = (str.size() < 4) ? (int)str.size() : 4; n--;) x.c[n] = str[n];
|
|
return x.i;
|
|
}
|
|
|
|
static void Usage() {
|
|
printf(
|
|
"BatchAigsEffectApp [flags ...] inFile1 [ inFileN ...]\n"
|
|
" where flags is:\n"
|
|
" --out_file=<path> output video files to be written (a pattern with one %%u or %%d), default \"BatchOut_%%02u.mp4\"\n"
|
|
" --model_dir=<path> the path to the directory that contains the models\n"
|
|
" --mode=<value> which model to pick for processing (default: 0)\n"
|
|
" --verbose verbose output\n"
|
|
" --codec=<fourcc> the fourcc code for the desired codec (default " DEFAULT_CODEC ")\n"
|
|
" and inFile1 ... are identically sized video files\n"
|
|
);
|
|
}
|
|
|
|
static int ParseMyArgs(int argc, char **argv) {
|
|
int errs = 0;
|
|
for (--argc, ++argv; argc--; ++argv) {
|
|
bool help;
|
|
const char *arg = *argv;
|
|
if (arg[0] == '-') {
|
|
if (arg[1] == '-') { // double-dash
|
|
if (GetFlagArgVal("verbose", arg, &FLAG_verbose) ||
|
|
GetFlagArgVal("mode", arg, &FLAG_mode) ||
|
|
GetFlagArgVal("model_dir", arg, &FLAG_modelDir) ||
|
|
GetFlagArgVal("out_file", arg, &FLAG_outFile) ||
|
|
GetFlagArgVal("codec", arg, &FLAG_codec)
|
|
) {
|
|
continue;
|
|
} else if (GetFlagArgVal("help", arg, &help)) { // --help
|
|
Usage();
|
|
errs = 1;
|
|
}
|
|
}
|
|
else { // single dash
|
|
for (++arg; *arg; ++arg) {
|
|
if (*arg == 'v') {
|
|
FLAG_verbose = true;
|
|
} else {
|
|
printf("Unknown flag ignored: \"-%c\"\n", *arg);
|
|
}
|
|
}
|
|
continue;
|
|
}
|
|
}
|
|
else { // no dash
|
|
FLAG_inFiles.push_back(arg);
|
|
}
|
|
}
|
|
return errs;
|
|
}
|
|
|
|
|
|
class App {
|
|
public:
|
|
NvVFX_Handle _eff;
|
|
NvCVImage _src, _stg, _dst;
|
|
CUstream _stream;
|
|
unsigned _batchSize;
|
|
|
|
|
|
App() : _eff(nullptr), _stream(0), _batchSize(0) {}
|
|
~App() {
|
|
NvVFX_DestroyEffect(_eff); if (_stream) NvVFX_CudaStreamDestroy(_stream);
|
|
}
|
|
|
|
NvCV_Status init(const char* effectName, unsigned batchSize, unsigned int mode, const NvCVImage *srcImg) {
|
|
NvCV_Status err = NVCV_ERR_UNIMPLEMENTED;
|
|
|
|
_batchSize = batchSize;
|
|
BAIL_IF_ERR(err = NvVFX_CreateEffect(effectName, &_eff));
|
|
|
|
BAIL_IF_ERR(err = AllocateBatchBuffer(&_src, _batchSize, srcImg->width, srcImg->height, NVCV_BGR, NVCV_U8, NVCV_CHUNKY, NVCV_GPU, 1));
|
|
BAIL_IF_ERR(err = AllocateBatchBuffer(&_dst, _batchSize, srcImg->width, srcImg->height, NVCV_A, NVCV_U8, NVCV_CHUNKY, NVCV_GPU, 1));
|
|
BAIL_IF_ERR(err = NvVFX_SetString(_eff, NVVFX_MODEL_DIRECTORY, FLAG_modelDir.c_str()));
|
|
|
|
|
|
{ // Set parameters.
|
|
NvCVImage nth;
|
|
BAIL_IF_ERR(err = NvVFX_SetImage(_eff, NVVFX_INPUT_IMAGE, NthImage(0, srcImg->height, &_src, &nth))); // Set the first of the batched images in ...
|
|
BAIL_IF_ERR(err = NvVFX_SetImage(_eff, NVVFX_OUTPUT_IMAGE, NthImage(0, _dst.height / _batchSize, &_dst, &nth))); // ... and out
|
|
BAIL_IF_ERR(err = NvVFX_CudaStreamCreate(&_stream));
|
|
BAIL_IF_ERR(err = NvVFX_SetCudaStream(_eff, NVVFX_CUDA_STREAM, _stream));
|
|
BAIL_IF_ERR(err = NvVFX_SetU32(_eff, NVVFX_MODE, mode));
|
|
}
|
|
|
|
bail:
|
|
return err;
|
|
}
|
|
};
|
|
|
|
|
|
NvCV_Status BatchProcess(const char* effectName, unsigned int mode,
|
|
const std::vector<const char*>& srcVideos, const char *outfilePattern, std::string codec) {
|
|
NvCV_Status err = NVCV_SUCCESS;
|
|
App app;
|
|
cv::Mat ocv1, ocv2;
|
|
NvCVImage nvx1, nvx2;
|
|
unsigned srcWidth, srcHeight, dstHeight;
|
|
|
|
std::vector<NvVFX_StateObjectHandle> arrayOfStates;
|
|
NvVFX_StateObjectHandle* batchOfStates = nullptr;
|
|
|
|
unsigned int numOfVideoStreams = static_cast<unsigned int>(srcVideos.size());
|
|
|
|
// If valid states are passed for inference, then -
|
|
// 1. Effect can only process a batch which is equal to maximum number of video streams
|
|
// 2. Multiple frames from the same video stream should not be present in the same batch
|
|
unsigned batchSize = numOfVideoStreams;
|
|
|
|
std::vector<cv::VideoCapture> srcCaptures(numOfVideoStreams);
|
|
std::vector<cv::VideoWriter> dstWriters(numOfVideoStreams);
|
|
for (unsigned int i = 0; i < numOfVideoStreams; i++) {
|
|
srcCaptures[i].open(srcVideos[i]);
|
|
if (srcCaptures[i].isOpened()==false) BAIL(err, NVCV_ERR_READ);
|
|
|
|
int width, height;
|
|
double fps;
|
|
width = (int)srcCaptures[i].get(cv::CAP_PROP_FRAME_WIDTH);
|
|
height = (int)srcCaptures[i].get(cv::CAP_PROP_FRAME_HEIGHT);
|
|
fps = srcCaptures[i].get(cv::CAP_PROP_FPS);
|
|
|
|
const int fourcc = StringToFourcc(codec);
|
|
char fileName[1024];
|
|
snprintf(fileName, sizeof(fileName), outfilePattern, i);
|
|
dstWriters[i].open(fileName, fourcc, fps, cv::Size2i(width,height), false);
|
|
if (dstWriters[i].isOpened() == false) BAIL(err, NVCV_ERR_WRITE);
|
|
}
|
|
|
|
// Read in the first image, to determine the resolution for init()
|
|
BAIL_IF_FALSE(srcVideos.size() > 0, err, NVCV_ERR_MISSINGINPUT);
|
|
srcCaptures[0] >> ocv1;
|
|
srcCaptures[0].set(cv::CAP_PROP_POS_FRAMES, 0); //resetting to first frame
|
|
if (!ocv1.data) {
|
|
printf("Cannot read video file \"%s\"\n", srcVideos[0]);
|
|
BAIL(err, NVCV_ERR_READ);
|
|
}
|
|
NVWrapperForCVMat(&ocv1, &nvx1);
|
|
srcWidth = nvx1.width;
|
|
srcHeight = nvx1.height;
|
|
|
|
BAIL_IF_ERR(err = app.init(effectName, batchSize, mode, &nvx1)); // Init effect and buffers
|
|
BAIL_IF_ERR(err = NvVFX_SetU32(app._eff, NVVFX_MAX_NUMBER_STREAMS, numOfVideoStreams));
|
|
BAIL_IF_ERR(err = NvVFX_SetU32(app._eff, NVVFX_MODEL_BATCH, numOfVideoStreams>1?8:1));
|
|
BAIL_IF_ERR(err = NvVFX_Load(app._eff));
|
|
|
|
// Creating state objects, one per stream.
|
|
for (unsigned int i = 0; i < numOfVideoStreams; i++) {
|
|
NvVFX_StateObjectHandle state;
|
|
BAIL_IF_ERR(err = NvVFX_AllocateState(app._eff, &state));
|
|
arrayOfStates.push_back(state);
|
|
}
|
|
|
|
//Creating batch array to hold states
|
|
batchOfStates = (NvVFX_StateObjectHandle*)malloc(sizeof(NvVFX_StateObjectHandle) * batchSize);
|
|
if (batchOfStates == nullptr) {
|
|
err = NVCV_ERR_MEMORY;
|
|
goto bail;
|
|
}
|
|
|
|
dstHeight = app._dst.height / batchSize;
|
|
BAIL_IF_ERR(err = NvCVImage_Alloc(&nvx2, app._dst.width, dstHeight, NVCV_A, NVCV_U8, NVCV_CHUNKY, NVCV_CPU, 0));
|
|
CVWrapperForNvCVImage(&nvx2, &ocv2);
|
|
for(int j=0;;j++)
|
|
{
|
|
for (unsigned int i = 0; i < batchSize; i++) {
|
|
int capIdx = i%numOfVideoStreams; // interlacing frames from different video stream, but can in any order
|
|
srcCaptures[capIdx] >> ocv1;
|
|
if (ocv1.empty()) goto bail;
|
|
batchOfStates[i] = arrayOfStates[capIdx];
|
|
|
|
NVWrapperForCVMat(&ocv1, &nvx1);
|
|
if (!(nvx1.width == srcWidth && nvx1.height == srcHeight)) {
|
|
printf("Input video file \"%s\" %ux%u does not match %ux%u\n"
|
|
"Batching requires all video frames to be of the same size\n", srcVideos[i], nvx1.width, nvx1.height, srcWidth, srcHeight);
|
|
BAIL(err, NVCV_ERR_MISMATCH);
|
|
}
|
|
BAIL_IF_ERR(err = TransferToNthImage(i, &nvx1, &app._src, 1.f, app._stream, NULL));
|
|
ocv1.release();
|
|
}
|
|
|
|
// Run batch
|
|
BAIL_IF_ERR(err = NvVFX_SetU32(app._eff, NVVFX_BATCH_SIZE, (unsigned)batchSize)); // The batchSize can change every Run
|
|
BAIL_IF_ERR(err = NvVFX_SetStateObjectHandleArray(app._eff, NVVFX_STATE, batchOfStates)); // The batch of states can change every Run
|
|
BAIL_IF_ERR(err = NvVFX_Run(app._eff, 0));
|
|
|
|
|
|
for (unsigned int i = 0; i < batchSize; ++i) {
|
|
int writerIdx = i % numOfVideoStreams;
|
|
BAIL_IF_ERR(err = TransferFromNthImage(i, &app._dst, &nvx2, 1.0f, app._stream, NULL));
|
|
dstWriters[writerIdx] << ocv2;
|
|
}
|
|
// NvCVImage_Dealloc() is called in the destructors
|
|
}
|
|
bail:
|
|
// If DeallocateState fails, all memory allocated in the SDK returns to the heap when the effect handle is destroyed.
|
|
for (unsigned int i = 0; i < arrayOfStates.size(); i++) {
|
|
NvVFX_DeallocateState(app._eff, arrayOfStates[i]);
|
|
}
|
|
arrayOfStates.clear();
|
|
|
|
if (batchOfStates) {
|
|
free(batchOfStates);
|
|
batchOfStates = nullptr;
|
|
}
|
|
|
|
for (auto& cap : srcCaptures) {
|
|
if (cap.isOpened()) cap.release();
|
|
}
|
|
|
|
for (auto& writer : dstWriters) {
|
|
if (writer.isOpened()) writer.release();
|
|
}
|
|
|
|
return err;
|
|
}
|
|
|
|
|
|
int main(int argc, char** argv) {
|
|
int nErrs;
|
|
NvCV_Status vfxErr;
|
|
|
|
nErrs = ParseMyArgs(argc, argv);
|
|
if (nErrs)
|
|
return nErrs;
|
|
|
|
// If the outFile is missing a stream index
|
|
// insert one, assuming a period followed by a three-character extension
|
|
if (FLAG_outFile.empty())
|
|
FLAG_outFile = "BatchOut_%02u.mp4";
|
|
else if (std::string::npos == FLAG_outFile.find_first_of('%'))
|
|
FLAG_outFile.insert(FLAG_outFile.size() - 4, "_%02u");
|
|
|
|
#ifdef NVVFX_FX_GREEN_SCREEN
|
|
vfxErr = BatchProcess(NVVFX_FX_GREEN_SCREEN, FLAG_mode, FLAG_inFiles, FLAG_outFile.c_str(), FLAG_codec);
|
|
#elif defined(NVVFX_FX_GREEN_SCREEN_I)
|
|
vfxErr = BatchProcess(NVVFX_FX_GREEN_SCREEN_I, FLAG_mode, FLAG_inFiles, FLAG_outFile.c_str(), FLAG_codec);
|
|
#endif
|
|
if (NVCV_SUCCESS != vfxErr) {
|
|
Usage();
|
|
printf("Error: %s\n", NvCV_GetErrorStringFromCode(vfxErr));
|
|
nErrs = (int)vfxErr;
|
|
}
|
|
|
|
return nErrs;
|
|
}
|