nvdsinfer_custom_impl.h

Go to the documentation of this file.

/*
 * Copyright (c) 2017-2019, NVIDIA CORPORATION.  All rights reserved.
 *
 * NVIDIA Corporation and its licensors retain all intellectual property
 * and proprietary rights in and to this software, related documentation
 * and any modifications thereto.  Any use, reproduction, disclosure or
 * distribution of this software and related documentation without an express
 * license agreement from NVIDIA Corporation is strictly prohibited.
 *
 */

#ifndef _NVDSINFER_CUSTOM_IMPL_H_
#define _NVDSINFER_CUSTOM_IMPL_H_

#include <string>
#include <vector>
#include "NvCaffeParser.h"
#include "NvUffParser.h"

#include "nvdsinfer.h"

#ifdef __cplusplus
extern "C"
{
#endif

typedef struct
{
  unsigned int numClassesConfigured;
  std::vector<float> perClassThreshold;
} NvDsInferParseDetectionParams;

typedef bool (* NvDsInferParseCustomFunc) (
        std::vector<NvDsInferLayerInfo> const &outputLayersInfo,
        NvDsInferNetworkInfo  const &networkInfo,
        NvDsInferParseDetectionParams const &detectionParams,
        std::vector<NvDsInferObjectDetectionInfo> &objectList);

#define CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(customParseFunc) \
    static void checkFunc_ ## customParseFunc (NvDsInferParseCustomFunc func = customParseFunc) \
        { checkFunc_ ## customParseFunc (); }; \
    extern "C" bool customParseFunc (std::vector<NvDsInferLayerInfo> const &outputLayersInfo, \
           NvDsInferNetworkInfo  const &networkInfo, \
           NvDsInferParseDetectionParams const &detectionParams, \
           std::vector<NvDsInferObjectDetectionInfo> &objectList);

typedef bool (* NvDsInferClassiferParseCustomFunc) (
        std::vector<NvDsInferLayerInfo> const &outputLayersInfo,
        NvDsInferNetworkInfo  const &networkInfo,
        float classifierThreshold,
        std::vector<NvDsInferAttribute> &attrList,
        std::string &descString);

#define CHECK_CUSTOM_CLASSIFIER_PARSE_FUNC_PROTOTYPE(customParseFunc) \
    static void checkFunc_ ## customParseFunc (NvDsInferClassiferParseCustomFunc func = customParseFunc) \
        { checkFunc_ ## customParseFunc (); }; \
    extern "C" bool customParseFunc (std::vector<NvDsInferLayerInfo> const &outputLayersInfo, \
           NvDsInferNetworkInfo  const &networkInfo, \
           float classifierThreshold, \
           std::vector<NvDsInferAttribute> &attrList, \
           std::string &descString);

typedef enum
{
  PLUGIN_FACTORY,
  PLUGIN_FACTORY_EXT,
  PLUGIN_FACTORY_V2
} NvDsInferPluginFactoryType;

typedef union
{
  nvcaffeparser1::IPluginFactory *pluginFactory;
  nvcaffeparser1::IPluginFactoryExt *pluginFactoryExt;
  nvcaffeparser1::IPluginFactoryV2 *pluginFactoryV2;
} NvDsInferPluginFactoryCaffe;

typedef union
{
  nvuffparser::IPluginFactory *pluginFactory;
  nvuffparser::IPluginFactoryExt *pluginFactoryExt;
} NvDsInferPluginFactoryUff;

bool NvDsInferPluginFactoryCaffeGet (NvDsInferPluginFactoryCaffe &pluginFactory,
    NvDsInferPluginFactoryType &type);

void NvDsInferPluginFactoryCaffeDestroy (NvDsInferPluginFactoryCaffe &pluginFactory);

bool NvDsInferPluginFactoryUffGet (NvDsInferPluginFactoryUff &pluginFactory,
    NvDsInferPluginFactoryType &type);

void NvDsInferPluginFactoryUffDestroy (NvDsInferPluginFactoryUff &pluginFactory);

bool NvDsInferPluginFactoryRuntimeGet (nvinfer1::IPluginFactory *& pluginFactory);

void NvDsInferPluginFactoryRuntimeDestroy (nvinfer1::IPluginFactory * pluginFactory);

bool NvDsInferInitializeInputLayers (std::vector<NvDsInferLayerInfo> const &inputLayersInfo,
        NvDsInferNetworkInfo const &networkInfo,
        unsigned int maxBatchSize);

typedef struct _NvDsInferContextInitParams NvDsInferContextInitParams;

bool NvDsInferCudaEngineGet(nvinfer1::IBuilder *builder,
        NvDsInferContextInitParams *initParams,
        nvinfer1::DataType dataType,
        nvinfer1::ICudaEngine *& cudaEngine);

#ifdef __cplusplus
}
#endif

#endif