16 #include <visp3/core/vpConfig.h>
18 #if defined(VISP_HAVE_TENSORRT) && defined(VISP_HAVE_OPENCV)
20 #if defined(HAVE_OPENCV_CUDEV) && defined(HAVE_OPENCV_CUDAWARPING) && defined(HAVE_OPENCV_CUDAARITHM) && \
21 defined(HAVE_OPENCV_DNN) && defined(HAVE_OPENCV_VIDEOIO)
22 #include <visp3/core/vpImageConvert.h>
23 #include <visp3/core/vpIoTools.h>
24 #include <visp3/gui/vpDisplayX.h>
26 #include <opencv2/videoio.hpp>
29 #include <opencv2/core/cuda.hpp>
30 #include <opencv2/cudaarithm.hpp>
31 #include <opencv2/cudawarping.hpp>
32 #include <opencv2/dnn.hpp>
36 #include <cuda_runtime_api.h>
41 #include <NvOnnxParser.h>
47 void preprocessImage(cv::Mat &img,
float *gpu_input,
const nvinfer1::Dims &dims,
float meanR,
float meanG,
float meanB)
50 std::cerr <<
"Image is empty." << std::endl;
54 cv::cuda::GpuMat gpu_frame;
56 gpu_frame.upload(img);
59 auto input_width = dims.d[3];
60 auto input_height = dims.d[2];
61 auto channels = dims.d[1];
62 auto input_size = cv::Size(input_width, input_height);
65 cv::cuda::GpuMat resized;
66 cv::cuda::resize(gpu_frame, resized, input_size, 0, 0, cv::INTER_NEAREST);
69 cv::cuda::GpuMat flt_image;
70 resized.convertTo(flt_image, CV_32FC3);
71 cv::cuda::subtract(flt_image, cv::Scalar(meanR, meanG, meanB), flt_image, cv::noArray(), -1);
72 cv::cuda::divide(flt_image, cv::Scalar(127.5f, 127.5f, 127.5f), flt_image, 1, -1);
75 std::vector<cv::cuda::GpuMat> chw;
76 for (
int i = 0; i < channels; ++i)
77 chw.emplace_back(cv::cuda::GpuMat(input_size, CV_32FC1, gpu_input + i * input_width * input_height));
78 cv::cuda::split(flt_image, chw);
83 size_t getSizeByDim(
const nvinfer1::Dims &dims)
86 for (
int i = 0; i < dims.nbDims; ++i)
93 std::vector<cv::Rect> postprocessResults(std::vector<void *> buffers,
const std::vector<nvinfer1::Dims> &output_dims,
94 int batch_size,
int image_width,
int image_height,
float confThresh,
95 float nmsThresh, std::vector<int> &classIds)
98 std::vector<cv::Rect> m_boxes, m_boxesNMS;
99 std::vector<int> m_classIds;
100 std::vector<float> m_confidences;
101 std::vector<int> m_indices;
104 std::vector<std::vector<float> > cpu_outputs;
105 for (
size_t i = 0; i < output_dims.size(); i++) {
106 cpu_outputs.push_back(std::vector<float>(getSizeByDim(output_dims[i]) * batch_size));
107 cudaMemcpy(cpu_outputs[i].data(), (
float *)buffers[1 + i], cpu_outputs[i].size() *
sizeof(
float),
108 cudaMemcpyDeviceToHost);
112 int N = output_dims[0].d[1], C = output_dims[0].d[2];
114 for (
int i = 0; i < N; i++)
116 uint32_t maxClass = 0;
117 float maxScore = -1000.0f;
119 for (
int j = 1; j < C; j++)
121 const float score = cpu_outputs[0][i * C + j];
123 if (score < confThresh)
126 if (score > maxScore) {
132 if (maxScore > confThresh) {
133 int left = (int)(cpu_outputs[1][4 * i] * image_width);
134 int top = (int)(cpu_outputs[1][4 * i + 1] * image_height);
135 int right = (int)(cpu_outputs[1][4 * i + 2] * image_width);
136 int bottom = (int)(cpu_outputs[1][4 * i + 3] * image_height);
137 int width = right - left + 1;
138 int height = bottom - top + 1;
140 m_boxes.push_back(cv::Rect(left, top, width, height));
141 m_classIds.push_back(maxClass);
142 m_confidences.push_back(maxScore);
146 cv::dnn::NMSBoxes(m_boxes, m_confidences, confThresh, nmsThresh, m_indices);
147 m_boxesNMS.resize(m_indices.size());
148 for (
size_t i = 0; i < m_indices.size(); ++i) {
149 int idx = m_indices[i];
150 m_boxesNMS[i] = m_boxes[idx];
153 classIds = m_classIds;
158 class Logger :
public nvinfer1::ILogger
161 void log(Severity severity,
const char *msg) noexcept
163 if ((severity == Severity::kERROR) || (severity == Severity::kINTERNAL_ERROR) || (severity == Severity::kVERBOSE))
164 std::cout << msg << std::endl;
171 template <
class T>
void operator()(T *obj)
const
178 template <
class T>
using TRTUniquePtr = std::unique_ptr<T, TRTDestroy>;
181 bool parseOnnxModel(
const std::string &model_path, TRTUniquePtr<nvinfer1::ICudaEngine> &engine,
182 TRTUniquePtr<nvinfer1::IExecutionContext> &context)
186 char cache_prefix[FILENAME_MAX];
187 char cache_path[FILENAME_MAX];
189 snprintf(cache_prefix, FILENAME_MAX,
"%s", model_path.c_str());
190 snprintf(cache_path, FILENAME_MAX,
"%s.engine", cache_prefix);
192 std::cout <<
"attempting to open engine cache file " << cache_path << std::endl;
196 char *engineStream =
nullptr;
197 size_t engineSize = 0;
200 struct stat filestat;
201 stat(cache_path, &filestat);
202 engineSize = filestat.st_size;
205 engineStream = (
char *)malloc(engineSize);
208 FILE *cacheFile =
nullptr;
209 cacheFile = fopen(cache_path,
"rb");
212 const size_t bytesRead = fread(engineStream, 1, engineSize, cacheFile);
214 if (bytesRead != engineSize)
216 std::cerr <<
"Error reading serialized engine into memory." << std::endl;
224 TRTUniquePtr<nvinfer1::IRuntime> infer { nvinfer1::createInferRuntime(gLogger) };
225 engine.reset(infer->deserializeCudaEngine(engineStream, engineSize,
nullptr));
226 context.reset(engine->createExecutionContext());
235 std::cerr <<
"Could not parse ONNX model. File not found" << std::endl;
239 TRTUniquePtr<nvinfer1::IBuilder> builder { nvinfer1::createInferBuilder(gLogger) };
240 TRTUniquePtr<nvinfer1::INetworkDefinition> network {
241 builder->createNetworkV2(1U << (uint32_t)nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH) };
242 TRTUniquePtr<nvonnxparser::IParser> parser { nvonnxparser::createParser(*network, gLogger) };
245 if (!parser->parseFromFile(model_path.c_str(),
static_cast<int>(nvinfer1::ILogger::Severity::kINFO))) {
246 std::cerr <<
"ERROR: could not parse the model." << std::endl;
250 TRTUniquePtr<nvinfer1::IBuilderConfig> config { builder->createBuilderConfig() };
252 config->setMaxWorkspaceSize(32 << 20);
254 if (builder->platformHasFastFp16()) {
255 config->setFlag(nvinfer1::BuilderFlag::kFP16);
258 builder->setMaxBatchSize(1);
260 engine.reset(builder->buildEngineWithConfig(*network, *config));
261 context.reset(engine->createExecutionContext());
263 TRTUniquePtr<nvinfer1::IHostMemory> serMem { engine->serialize() };
266 std::cout <<
"Failed to serialize CUDA engine." << std::endl;
270 const char *serData = (
char *)serMem->data();
271 const size_t serSize = serMem->size();
274 char *engineMemory = (
char *)malloc(serSize);
277 std::cout <<
"Failed to allocate memory to store CUDA engine." << std::endl;
281 memcpy(engineMemory, serData, serSize);
284 FILE *cacheFile =
nullptr;
285 cacheFile = fopen(cache_path,
"wb");
287 fwrite(engineMemory, 1, serSize, cacheFile);
295 int main(
int argc,
char **argv)
298 unsigned int opt_scale = 1;
299 std::string input =
"";
302 std::string config =
"";
303 float meanR = 127.5f, meanG = 127.5f, meanB = 127.5f;
304 float confThresh = 0.5f;
305 float nmsThresh = 0.4f;
307 for (
int i = 1; i < argc; i++) {
308 if (std::string(argv[i]) ==
"--device" && i + 1 < argc) {
309 opt_device = atoi(argv[i + 1]);
311 else if (std::string(argv[i]) ==
"--input" && i + 1 < argc) {
312 input = std::string(argv[i + 1]);
314 else if (std::string(argv[i]) ==
"--model" && i + 1 < argc) {
315 modelFile = std::string(argv[i + 1]);
317 else if (std::string(argv[i]) ==
"--config" && i + 1 < argc) {
318 config = std::string(argv[i + 1]);
320 else if (std::string(argv[i]) ==
"--input-scale" && i + 1 < argc) {
321 opt_scale = atoi(argv[i + 1]);
323 else if (std::string(argv[i]) ==
"--mean" && i + 3 < argc) {
324 meanR = atof(argv[i + 1]);
325 meanG = atof(argv[i + 2]);
326 meanB = atof(argv[i + 3]);
328 else if (std::string(argv[i]) ==
"--confThresh" && i + 1 < argc) {
329 confThresh = (float)atof(argv[i + 1]);
331 else if (std::string(argv[i]) ==
"--nmsThresh" && i + 1 < argc) {
332 nmsThresh = (float)atof(argv[i + 1]);
334 else if (std::string(argv[i]) ==
"--labels" && i + 1 < argc) {
335 labelFile = std::string(argv[i + 1]);
337 else if (std::string(argv[i]) ==
"--help" || std::string(argv[i]) ==
"-h") {
339 <<
" [--device <camera device number>] [--input <path to image or video>"
340 " (camera is used if input is empty)] [--model <path to net trained weights>]"
341 " [--config <path to net config file>]"
342 " [--input-scale <input scale factor>] [--mean <meanR meanG meanB>]"
343 " [--confThresh <confidence threshold>]"
344 " [--nmsThresh <NMS threshold>] [--labels <path to label file>]"
350 std::string model_path(modelFile);
353 std::vector<std::string> labels;
354 if (!labelFile.empty()) {
355 std::ifstream f_label(labelFile);
357 while (std::getline(f_label, line)) {
358 labels.push_back(line);
364 TRTUniquePtr<nvinfer1::ICudaEngine> engine {
nullptr };
365 TRTUniquePtr<nvinfer1::IExecutionContext> context {
nullptr };
366 if (!parseOnnxModel(model_path, engine, context))
368 std::cout <<
"Make sure the model file exists. To see available models, plese visit: "
369 "\n\twww.github.com/lagadic/visp-images/dnn/object_detection/"
375 std::vector<nvinfer1::Dims> input_dims;
376 std::vector<nvinfer1::Dims> output_dims;
377 std::vector<void *> buffers(engine->getNbBindings());
380 for (
int i = 0; i < engine->getNbBindings(); ++i) {
381 auto binding_size = getSizeByDim(engine->getBindingDimensions(i)) * batch_size *
sizeof(float);
382 cudaMalloc(&buffers[i], binding_size);
384 if (engine->bindingIsInput(i)) {
385 input_dims.emplace_back(engine->getBindingDimensions(i));
388 output_dims.emplace_back(engine->getBindingDimensions(i));
392 if (input_dims.empty() || output_dims.empty()) {
393 std::cerr <<
"Expect at least one input and one output for network" << std::endl;
399 cv::VideoCapture capture;
402 capture.open(opt_device);
408 if (!capture.isOpened()) {
409 std::cout <<
"Failed to open the camera" << std::endl;
413 int cap_width = (int)capture.get(cv::CAP_PROP_FRAME_WIDTH);
414 int cap_height = (int)capture.get(cv::CAP_PROP_FRAME_HEIGHT);
415 capture.set(cv::CAP_PROP_FRAME_WIDTH, cap_width / opt_scale);
416 capture.set(cv::CAP_PROP_FRAME_HEIGHT, cap_height / opt_scale);
425 while ((i++ < 20) && !capture.read(frame)) {
432 std::cout <<
"Image size: " << width <<
" x " << height << std::endl;
434 std::vector<cv::Rect> boxesNMS;
435 std::vector<int> classIds;
449 preprocessImage(frame, (
float *)buffers[0], input_dims[0], meanR, meanG, meanB);
452 context->enqueue(batch_size, buffers.data(), 0,
nullptr);
455 boxesNMS = postprocessResults(buffers, output_dims, batch_size, width, height, confThresh, nmsThresh, classIds);
462 for (
unsigned int i = 0; i < boxesNMS.size(); i++) {
472 for (
void *buf : buffers)
480 std::cout <<
"OpenCV is not built with CUDA." << std::endl;
489 std::cout <<
"ViSP is not built with TensorRT." << std::endl;
Use the X11 console to display images on unix-like OS. Thus to enable this class X11 should be instal...
static bool getClick(const vpImage< unsigned char > &I, bool blocking=true)
static void display(const vpImage< unsigned char > &I)
static void flush(const vpImage< unsigned char > &I)
static void displayRectangle(const vpImage< unsigned char > &I, const vpImagePoint &topLeft, unsigned int width, unsigned int height, const vpColor &color, bool fill=false, unsigned int thickness=1)
static void displayText(const vpImage< unsigned char > &I, const vpImagePoint &ip, const std::string &s, const vpColor &color)
static void convert(const vpImage< unsigned char > &src, vpImage< vpRGBa > &dest)
unsigned int getWidth() const
unsigned int getHeight() const
Defines a rectangle in the plane.
VISP_EXPORT double measureTimeMs()