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.
401 lines
16 KiB
401 lines
16 KiB
2 months ago
|
/*###############################################################################
|
||
|
#
|
||
|
# Copyright (c) 2020 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 <stdarg.h>
|
||
|
#include <stdio.h>
|
||
|
#include <string.h>
|
||
|
|
||
|
#include <string>
|
||
|
|
||
|
#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)
|
||
|
|
||
|
|
||
|
bool FLAG_verbose = false;
|
||
|
float FLAG_strength = 0.f,
|
||
|
FLAG_scale = 1.0;
|
||
|
int FLAG_mode = 0,
|
||
|
FLAG_resolution = 0;
|
||
|
std::string FLAG_outFile,
|
||
|
FLAG_modelDir,
|
||
|
FLAG_effect;
|
||
|
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, float *val) {
|
||
|
const char *valStr;
|
||
|
bool success = GetFlagArgVal(flag, arg, &valStr);
|
||
|
if (success)
|
||
|
*val = strtof(valStr, NULL);
|
||
|
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 void Usage() {
|
||
|
printf(
|
||
|
"BatchEffectApp [flags ...] inFile1 [ inFileN ...]\n"
|
||
|
" where flags is:\n"
|
||
|
" --out_file=<path> output image files to be written, default \"BatchOut_%%02u.png\"\n"
|
||
|
" --effect=<effect> the effect to apply\n"
|
||
|
" --strength=<value> strength of the upscaling effect, [0.0, 1.0]\n"
|
||
|
" --scale=<scale> scale factor to be applied: 1.5, 2, 3, maybe 1.3333333\n"
|
||
|
" --resolution=<height> the desired height (either --scale or --resolution may be used)\n"
|
||
|
" --mode=<mode> mode 0 or 1\n"
|
||
|
" --model_dir=<path> the path to the directory that contains the models\n"
|
||
|
" --verbose verbose output\n"
|
||
|
" and inFile1 ... are identically sized image files, e.g. png, jpg\n"
|
||
|
);
|
||
|
|
||
|
const char* cStr;
|
||
|
NvCV_Status err = NvVFX_GetString(nullptr, NVVFX_INFO, &cStr);
|
||
|
if (NVCV_SUCCESS != err)
|
||
|
printf("Cannot get effects: %s\n", NvCV_GetErrorStringFromCode(err));
|
||
|
printf("where effects are:\n%s", cStr);
|
||
|
}
|
||
|
|
||
|
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("effect", arg, &FLAG_effect) ||
|
||
|
GetFlagArgVal("strength", arg, &FLAG_strength) ||
|
||
|
GetFlagArgVal("scale", arg, &FLAG_scale) ||
|
||
|
GetFlagArgVal("mode", arg, &FLAG_mode) ||
|
||
|
GetFlagArgVal("model_dir", arg, &FLAG_modelDir) ||
|
||
|
GetFlagArgVal("out_file", arg, &FLAG_outFile)
|
||
|
) {
|
||
|
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;
|
||
|
}
|
||
|
|
||
|
static bool HasSuffix(const char *str, const char *suf) {
|
||
|
size_t strSize = strlen(str),
|
||
|
sufSize = strlen(suf);
|
||
|
if (strSize < sufSize)
|
||
|
return false;
|
||
|
return (0 == strcasecmp(suf, str + strSize - sufSize));
|
||
|
}
|
||
|
|
||
|
static bool HasOneOfTheseSuffixes(const char *str, ...) {
|
||
|
bool matches = false;
|
||
|
const char *suf;
|
||
|
va_list ap;
|
||
|
va_start(ap, str);
|
||
|
while (nullptr != (suf = va_arg(ap, const char*))) {
|
||
|
if (HasSuffix(str, suf)) {
|
||
|
matches = true;
|
||
|
break;
|
||
|
}
|
||
|
}
|
||
|
va_end(ap);
|
||
|
return matches;
|
||
|
}
|
||
|
|
||
|
static bool IsLossyImageFile(const char *str) {
|
||
|
return HasOneOfTheseSuffixes(str, ".jpg", ".jpeg", nullptr);
|
||
|
}
|
||
|
|
||
|
class App {
|
||
|
public:
|
||
|
NvVFX_Handle _eff;
|
||
|
NvCVImage _src, _dst, _stg;
|
||
|
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, const NvCVImage *src) {
|
||
|
NvCV_Status err = NVCV_ERR_UNIMPLEMENTED;
|
||
|
unsigned dw, dh;
|
||
|
|
||
|
if (FLAG_resolution) {
|
||
|
dw = FLAG_resolution * src->width / src->height, // No rounding
|
||
|
dh = FLAG_resolution;
|
||
|
}
|
||
|
else {
|
||
|
dw = lroundf(src->width * FLAG_scale),
|
||
|
dh = lroundf(src->height * FLAG_scale);
|
||
|
}
|
||
|
|
||
|
_batchSize = batchSize;
|
||
|
BAIL_IF_ERR(err = NvVFX_CreateEffect(effectName, &_eff));
|
||
|
|
||
|
if (!strcmp(effectName, NVVFX_FX_TRANSFER)) {
|
||
|
BAIL_IF_ERR(err = AllocateBatchBuffer(&_src, _batchSize, src->width, src->height, NVCV_RGB, NVCV_U8, NVCV_CHUNKY, NVCV_CUDA, 0));
|
||
|
BAIL_IF_ERR(err = AllocateBatchBuffer(&_dst, _batchSize, src->width, src->height, NVCV_RGB, NVCV_U8, NVCV_CHUNKY, NVCV_CUDA, 0));
|
||
|
}
|
||
|
#ifdef NVVFX_FX_SR_UPSCALE
|
||
|
else if (!strcmp(effectName, NVVFX_FX_SR_UPSCALE)) {
|
||
|
BAIL_IF_ERR(err = AllocateBatchBuffer(&_src, _batchSize, src->width, src->height, NVCV_RGBA, NVCV_U8, NVCV_CHUNKY, NVCV_CUDA, 32)); // n*32, n>=0
|
||
|
BAIL_IF_ERR(err = AllocateBatchBuffer(&_dst, _batchSize, dw, dh, NVCV_RGBA, NVCV_U8, NVCV_CHUNKY, NVCV_CUDA, 32));
|
||
|
BAIL_IF_ERR(err = NvVFX_SetF32(_eff, NVVFX_STRENGTH, FLAG_strength));
|
||
|
}
|
||
|
#endif // NVVFX_FX_SR_UPSCALE
|
||
|
#ifdef NVVFX_FX_ARTIFACT_REDUCTION
|
||
|
else if (!strcmp(effectName, NVVFX_FX_ARTIFACT_REDUCTION)) {
|
||
|
BAIL_IF_ERR(err = AllocateBatchBuffer(&_src, _batchSize, src->width, src->height, NVCV_BGR, NVCV_F32, NVCV_PLANAR, NVCV_CUDA, 1));
|
||
|
BAIL_IF_ERR(err = AllocateBatchBuffer(&_dst, _batchSize, src->width, src->height, NVCV_BGR, NVCV_F32, NVCV_PLANAR, NVCV_CUDA, 1));
|
||
|
BAIL_IF_ERR(err = NvVFX_SetString(_eff, NVVFX_MODEL_DIRECTORY, FLAG_modelDir.c_str()));
|
||
|
BAIL_IF_ERR(err = NvVFX_SetU32(_eff, NVVFX_MODE, FLAG_mode));
|
||
|
}
|
||
|
#endif // NVVFX_FX_ARTIFACT_REDUCTION
|
||
|
#ifdef NVVFX_FX_SUPER_RES
|
||
|
else if (!strcmp(effectName, NVVFX_FX_SUPER_RES)) {
|
||
|
BAIL_IF_ERR(err = AllocateBatchBuffer(&_src, _batchSize, src->width, src->height, NVCV_BGR, NVCV_F32, NVCV_PLANAR, NVCV_CUDA, 1));
|
||
|
BAIL_IF_ERR(err = AllocateBatchBuffer(&_dst, _batchSize, dw, dh, NVCV_BGR, NVCV_F32, NVCV_PLANAR, NVCV_CUDA, 1));
|
||
|
BAIL_IF_ERR(err = NvVFX_SetString(_eff, NVVFX_MODEL_DIRECTORY, FLAG_modelDir.c_str()));
|
||
|
BAIL_IF_ERR(err = NvVFX_SetU32(_eff, NVVFX_MODE, FLAG_mode));
|
||
|
BAIL_IF_ERR(err = NvVFX_SetF32(_eff, NVVFX_STRENGTH, FLAG_strength));
|
||
|
}
|
||
|
#endif // NVVFX_FX_SUPER_RES
|
||
|
else {
|
||
|
BAIL(err, NVCV_ERR_UNIMPLEMENTED);
|
||
|
}
|
||
|
|
||
|
{ // Set common parameters.
|
||
|
NvCVImage nth;
|
||
|
BAIL_IF_ERR(err = NvVFX_SetImage(_eff, NVVFX_INPUT_IMAGE, NthImage(0, src->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));
|
||
|
|
||
|
// The batch size parameter is interpreted at two times:
|
||
|
// (1) during Load(), an appropriate batch-size model is chosen and loaded;
|
||
|
// (2) during Run(), the specified number of images in the batch are processed.
|
||
|
// The optimum throughput results from submitting a batch which is an integral multiple of the batched model
|
||
|
// chosen in Load().
|
||
|
//
|
||
|
// To request a particular batch-sized model, set the batch size before calling Load(),
|
||
|
// then get the batch size afterward to find out what batch-size model was chosen. If you do not specify the
|
||
|
// desired batchSize before calling Load(), it will choose the batchSize=1 model, since that is the default
|
||
|
// value for batchSize.
|
||
|
//
|
||
|
// After calling Load(), you can subsequently change the batch size to any number, even larger or smaller
|
||
|
// than the batch size of the chosen model. If a larger batch size is chosen, smaller batches are submitted
|
||
|
// until the entire larger batch has been processed. In any event, the batch size should be set at least twice:
|
||
|
// once before Load() and once before the initial Run(). In many server applications, it is expected that
|
||
|
// the batch size is changing constantly as some videos complete and other are added, so setting the batchSize
|
||
|
// before every Run() call would be typical.
|
||
|
unsigned gotBatch;
|
||
|
BAIL_IF_ERR(err = NvVFX_SetU32(_eff, NVVFX_MODEL_BATCH, _batchSize)); // Try to choose a model tuned to this batch size
|
||
|
err = NvVFX_Load(_eff); // This will load a new batched model -- a weighty process
|
||
|
if (!(NVCV_SUCCESS == err || NVCV_ERR_MODELSUBSTITUTION == err)) goto bail;
|
||
|
BAIL_IF_ERR(err = NvVFX_GetU32(_eff, NVVFX_MODEL_BATCH, &gotBatch)); // This tells us the batch size of the chosen model
|
||
|
if (FLAG_verbose && gotBatch != _batchSize) {
|
||
|
printf("Effect %s has no batch=%u model; processing in multiple batches of size %u%s instead\n",
|
||
|
effectName, _batchSize, gotBatch, (gotBatch > 1 ? " or less" : ""));
|
||
|
BAIL_IF_ERR(err = NvVFX_SetU32(_eff, NVVFX_BATCH_SIZE, _batchSize)); // This is lightweight, and usually done each Run
|
||
|
}
|
||
|
}
|
||
|
|
||
|
bail:
|
||
|
return err;
|
||
|
}
|
||
|
};
|
||
|
|
||
|
|
||
|
NvCV_Status BatchProcessImages(const char* effectName, const std::vector<const char*>& srcImages, const char *outfilePattern) {
|
||
|
NvCV_Status err = NVCV_SUCCESS;
|
||
|
unsigned batchSize = (unsigned)srcImages.size();
|
||
|
App app;
|
||
|
cv::Mat ocv;
|
||
|
NvCVImage nvx;
|
||
|
unsigned srcWidth, srcHeight, dstHeight, i;
|
||
|
|
||
|
// Read in the first image, to determine the resolution for init()
|
||
|
BAIL_IF_FALSE(srcImages.size() > 0, err, NVCV_ERR_MISSINGINPUT);
|
||
|
ocv = cv::imread(srcImages[0]);
|
||
|
if (!ocv.data) {
|
||
|
printf("Cannot read image file \"%s\"\n", srcImages[0]);
|
||
|
BAIL(err, NVCV_ERR_READ);
|
||
|
}
|
||
|
NVWrapperForCVMat(&ocv, &nvx);
|
||
|
srcWidth = nvx.width;
|
||
|
srcHeight = nvx.height;
|
||
|
BAIL_IF_ERR(err = app.init(effectName, batchSize, &nvx)); // Init effect and buffers
|
||
|
|
||
|
// Transfer the first image to the batch src.
|
||
|
// Note, in all transfers, the scale factor only applies to floating-point pixels.
|
||
|
BAIL_IF_ERR(err = TransferToNthImage(0, &nvx, &app._src, 1.f/255.f, app._stream, &app._stg));
|
||
|
ocv.release();
|
||
|
|
||
|
// Read the remaining images and transfer to the batch src
|
||
|
for (i = 1; i < batchSize; ++i) {
|
||
|
ocv = cv::imread(srcImages[i]);
|
||
|
if (!ocv.data) {
|
||
|
printf("Cannot read image file \"%s\"\n", srcImages[i]);
|
||
|
BAIL(err, NVCV_ERR_READ);
|
||
|
}
|
||
|
NVWrapperForCVMat(&ocv, &nvx);
|
||
|
if (!(nvx.width == srcWidth && nvx.height == srcHeight)) {
|
||
|
printf("Input image file \"%s\" %ux%u does not match %ux%u\n", srcImages[i], nvx.width, nvx.height, srcWidth, srcHeight);
|
||
|
BAIL(err, NVCV_ERR_MISMATCH);
|
||
|
}
|
||
|
BAIL_IF_ERR(err = TransferToNthImage(i, &nvx, &app._src, 1.f / 255.f, app._stream, &app._stg));
|
||
|
ocv.release();
|
||
|
}
|
||
|
|
||
|
// Run batch
|
||
|
BAIL_IF_ERR(err = NvVFX_SetU32(app._eff, NVVFX_BATCH_SIZE, (unsigned)srcImages.size())); // The batchSize can change every Run
|
||
|
BAIL_IF_ERR(err = NvVFX_Run(app._eff, 0));
|
||
|
|
||
|
// Retrieve and write images
|
||
|
dstHeight = app._dst.height / batchSize;
|
||
|
BAIL_IF_ERR(err = NvCVImage_Alloc(&nvx, app._dst.width, dstHeight, ((app._dst.numComponents == 1) ? NVCV_Y : NVCV_BGR), NVCV_U8, NVCV_CHUNKY, NVCV_CPU, 0));
|
||
|
CVWrapperForNvCVImage(&nvx, &ocv);
|
||
|
if(IsLossyImageFile(outfilePattern))
|
||
|
fprintf(stderr, "WARNING: JPEG output file format will reduce image quality\n");
|
||
|
for (i = 0; i < batchSize; ++i) {
|
||
|
char fileName[1024];
|
||
|
snprintf(fileName, sizeof(fileName), outfilePattern, i);
|
||
|
BAIL_IF_ERR(err = TransferFromNthImage(i, &app._dst, &nvx, 255.f, app._stream, &app._stg));
|
||
|
if (!cv::imwrite(fileName, ocv)) {
|
||
|
printf("Cannot write image file \"%s\"\n", fileName);
|
||
|
BAIL(err, NVCV_ERR_WRITE);
|
||
|
}
|
||
|
}
|
||
|
// NvCVImage_Dealloc() is called in the destructors
|
||
|
|
||
|
bail:
|
||
|
return err;
|
||
|
}
|
||
|
|
||
|
|
||
|
int main(int argc, char** argv) {
|
||
|
int nErrs;
|
||
|
NvCV_Status vfxErr;
|
||
|
|
||
|
nErrs = ParseMyArgs(argc, argv);
|
||
|
if (nErrs)
|
||
|
return nErrs;
|
||
|
|
||
|
if (FLAG_outFile.empty())
|
||
|
FLAG_outFile = "BatchOut_%02u.png";
|
||
|
else if (std::string::npos == FLAG_outFile.find_first_of('%'))
|
||
|
FLAG_outFile.insert(FLAG_outFile.size() - 4, "_%02u"); // assuming .xxx, i.e. .jpg, .png
|
||
|
|
||
|
vfxErr = BatchProcessImages(FLAG_effect.c_str(), FLAG_inFiles, FLAG_outFile.c_str());
|
||
|
if (NVCV_SUCCESS != vfxErr) {
|
||
|
printf("Error: %s\n", NvCV_GetErrorStringFromCode(vfxErr));
|
||
|
nErrs = (int)vfxErr;
|
||
|
}
|
||
|
|
||
|
return nErrs;
|
||
|
}
|