chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Waiting to run
Docker Image CI / build-ubuntu2004 (push) Waiting to run
This commit is contained in:
@@ -0,0 +1,22 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
add_plugin_source(
|
||||
generateDetectionPlugin.cpp
|
||||
generateDetectionPlugin.h
|
||||
)
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
---
|
||||
name: GenerateDetection_TRT
|
||||
interface: "IPluginV2Ext"
|
||||
versions:
|
||||
"1":
|
||||
attributes:
|
||||
- num_classes
|
||||
- keep_topk
|
||||
- score_threshold
|
||||
- iou_threshold
|
||||
- image_size
|
||||
attribute_types:
|
||||
num_classes: int32
|
||||
keep_topk: int32
|
||||
score_threshold: float32
|
||||
iou_threshold: float32
|
||||
image_size: int32
|
||||
attribute_length:
|
||||
num_classes: 1
|
||||
keep_topk: 1
|
||||
score_threshold: 1
|
||||
iou_threshold: 1
|
||||
image_size: 3
|
||||
attribute_options:
|
||||
num_classes:
|
||||
min: "0"
|
||||
max: "=pinf"
|
||||
keep_topk:
|
||||
min: "0"
|
||||
max: "=pinf"
|
||||
score_threshold:
|
||||
min: "=0"
|
||||
max: "=pinf"
|
||||
iou_threshold:
|
||||
min: "0"
|
||||
max: "=pinf"
|
||||
image_size:
|
||||
min: "0, 0, 0"
|
||||
max: "=pinf, =pinf, =pinf"
|
||||
attributes_required:
|
||||
- num_classes
|
||||
- keep_topk
|
||||
- score_threshold
|
||||
- iou_threshold
|
||||
...
|
||||
@@ -0,0 +1,76 @@
|
||||
# generateDetection Plugin [DEPRECATED]
|
||||
|
||||
**This plugin is deprecated since TensorRT 10.12 and will be removed in a future release. No alternatives are planned to be provided.**
|
||||
|
||||
**Table Of Contents**
|
||||
- [Description](#description)
|
||||
* [Structure](#structure)
|
||||
- [Parameters](#parameters)
|
||||
- [Additional resources](#additional-resources)
|
||||
- [License](#license)
|
||||
- [Changelog](#changelog)
|
||||
- [Known issues](#known-issues)
|
||||
|
||||
## Description
|
||||
|
||||
The `generateDetection` plugin performs bounding boxe refinement of MaskRCNN's detection head and generates the final detection output of MaskRCNN.
|
||||
|
||||
|
||||
### Structure
|
||||
|
||||
This plugin supports the NCHW format. It takes three input tensors: `delta_bbox`, `score` and `roi`
|
||||
|
||||
`delta_bbox` is the refinement information of roi boxes generated from the `MultilevelProposeROI` plugin. `delta_bbox` tensor's shape is `[N, rois, num_classes*4, 1, 1]` where `N` is batch size,
|
||||
`rois` is the total number of ROI boxes candidates per image, and `num_classes*4` means 4 refinement elements (`[dy, dx, dh, dw]`) for each roi box as different classes.
|
||||
|
||||
`score` is the predicted class scores of ROI boxes generated from MaskRCNN detection head of shape `[N, rois, num_classes, 1, 1]`. There is an `argmax` operation in `generateDetection` to determine the final class of detection
|
||||
candidates.
|
||||
|
||||
`roi` is the coordinates of ROI boxes candidates from the `MultilevelProposeROI` plugin of shape `[N, rois, 4]`.
|
||||
|
||||
This plugin generates output of shape `[N, keep_topk, 6]` where `keep_topk` is the maximum number of detections left after NMS and '6' means 6 elements of an detection `[y1, x1, y2, x2,
|
||||
class_label, score]`
|
||||
|
||||
## Parameters
|
||||
|
||||
This plugin has the plugin creator class `generateDetectionPluginCreator` and the plugin class `generateDetection`.
|
||||
|
||||
The following parameters were used to create `generateDetection` instance:
|
||||
|
||||
| Type | Parameter | Description
|
||||
|--------------------|------------------------------------|--------------------------------------------------------
|
||||
|`int` |`num_classes` |Number of detection classes(including `background`). `num_classes=91` for COCO dataset
|
||||
|`int` |`keep_topk` |Number of detections will be kept after NMS.
|
||||
|`float` |`score_threshold` |Confidence threshold value. This plugin will drop a detection if its class confidence(score) is under "score_threshold".
|
||||
|`float` |`iou_threshold` |IOU threshold value used in NMS.
|
||||
|`int[3]` |`image_size` |Input image size in CHW. Defaults to [3,832,1344]
|
||||
|
||||
## Limitations
|
||||
|
||||
The number of anchors is capped at 2048 to support embedded devices with smaller shared memory capacity.
|
||||
|
||||
To enable support for a device with higher memory, calls to `sortPerClass` and `KeepTopKGather` can be modified in `DetectionPostProcess` ([maskRCNNKernels.cu](https://github.com/NVIDIA/TensorRT/blob/main/plugin/common/kernels/maskRCNNKernels.cu)).
|
||||
|
||||
## Additional resources
|
||||
|
||||
|
||||
|
||||
## License
|
||||
|
||||
For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html)
|
||||
documentation.
|
||||
|
||||
|
||||
## Changelog
|
||||
|
||||
May 2025
|
||||
Add deprecation note.
|
||||
|
||||
January 2022: The [Limitations](#limitations) section was added to this `README.md` file to document limitations of the plugin related to the maximum number of anchors it can support.
|
||||
|
||||
June 2020: First release of this `README.md` file.
|
||||
|
||||
|
||||
## Known issues
|
||||
|
||||
There are no known issues in this plugin.
|
||||
@@ -0,0 +1,350 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include "generateDetectionPlugin.h"
|
||||
#include "common/plugin.h"
|
||||
#include <algorithm>
|
||||
#include <cuda_runtime_api.h>
|
||||
#include <string_view>
|
||||
|
||||
using namespace nvinfer1;
|
||||
using namespace plugin;
|
||||
using nvinfer1::plugin::GenerateDetection;
|
||||
using nvinfer1::plugin::GenerateDetectionPluginCreator;
|
||||
|
||||
#include <fstream>
|
||||
|
||||
namespace
|
||||
{
|
||||
char const* const kGENERATEDETECTION_PLUGIN_VERSION{"1"};
|
||||
char const* const kGENERATEDETECTION_PLUGIN_NAME{"GenerateDetection_TRT"};
|
||||
} // namespace
|
||||
|
||||
GenerateDetectionPluginCreator::GenerateDetectionPluginCreator() noexcept
|
||||
{
|
||||
mPluginAttributes.clear();
|
||||
mPluginAttributes.emplace_back(PluginField("num_classes", nullptr, PluginFieldType::kINT32, 1));
|
||||
mPluginAttributes.emplace_back(PluginField("keep_topk", nullptr, PluginFieldType::kINT32, 1));
|
||||
mPluginAttributes.emplace_back(PluginField("score_threshold", nullptr, PluginFieldType::kFLOAT32, 1));
|
||||
mPluginAttributes.emplace_back(PluginField("iou_threshold", nullptr, PluginFieldType::kFLOAT32, 1));
|
||||
mPluginAttributes.emplace_back(PluginField("image_size", nullptr, PluginFieldType::kINT32, 3));
|
||||
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
char const* GenerateDetectionPluginCreator::getPluginName() const noexcept
|
||||
{
|
||||
return kGENERATEDETECTION_PLUGIN_NAME;
|
||||
}
|
||||
|
||||
char const* GenerateDetectionPluginCreator::getPluginVersion() const noexcept
|
||||
{
|
||||
return kGENERATEDETECTION_PLUGIN_VERSION;
|
||||
}
|
||||
|
||||
PluginFieldCollection const* GenerateDetectionPluginCreator::getFieldNames() noexcept
|
||||
{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
IPluginV2Ext* GenerateDetectionPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
using namespace std::string_view_literals;
|
||||
auto image_size = TLTMaskRCNNConfig::IMAGE_SHAPE;
|
||||
PluginField const* fields = fc->fields;
|
||||
plugin::validateRequiredAttributesExist({"num_classes", "keep_topk", "score_threshold", "iou_threshold"}, fc);
|
||||
|
||||
for (int32_t i = 0; i < fc->nbFields; ++i)
|
||||
{
|
||||
std::string_view const attrName = fields[i].name;
|
||||
if (attrName == "num_classes"sv)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
|
||||
mNbClasses = *(static_cast<int32_t const*>(fields[i].data));
|
||||
}
|
||||
if (attrName == "keep_topk"sv)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
|
||||
mKeepTopK = *(static_cast<int32_t const*>(fields[i].data));
|
||||
}
|
||||
if (attrName == "score_threshold"sv)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kFLOAT32);
|
||||
mScoreThreshold = *(static_cast<float const*>(fields[i].data));
|
||||
}
|
||||
if (attrName == "iou_threshold"sv)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kFLOAT32);
|
||||
mIOUThreshold = *(static_cast<float const*>(fields[i].data));
|
||||
}
|
||||
if (attrName == "image_size"sv)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
|
||||
auto const dims = static_cast<int32_t const*>(fields[i].data);
|
||||
std::copy_n(dims, 3, image_size.d);
|
||||
}
|
||||
}
|
||||
return new GenerateDetection(mNbClasses, mKeepTopK, mScoreThreshold, mIOUThreshold, image_size);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV2Ext* GenerateDetectionPluginCreator::deserializePlugin(
|
||||
char const* name, void const* data, size_t length) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
return new GenerateDetection(data, length);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
GenerateDetection::GenerateDetection(int32_t num_classes, int32_t keep_topk, float score_threshold, float iou_threshold,
|
||||
nvinfer1::Dims const& image_size)
|
||||
: mNbClasses(num_classes)
|
||||
, mKeepTopK(keep_topk)
|
||||
, mScoreThreshold(score_threshold)
|
||||
, mIOUThreshold(iou_threshold)
|
||||
, mImageSize(image_size)
|
||||
{
|
||||
mBackgroundLabel = 0;
|
||||
PLUGIN_VALIDATE(mNbClasses > 0);
|
||||
PLUGIN_VALIDATE(mKeepTopK > 0);
|
||||
PLUGIN_VALIDATE(score_threshold >= 0.0F);
|
||||
PLUGIN_VALIDATE(iou_threshold > 0.0F);
|
||||
PLUGIN_VALIDATE(mImageSize.nbDims == 3);
|
||||
PLUGIN_VALIDATE(mImageSize.d[0] > 0 && mImageSize.d[1] > 0 && mImageSize.d[2] > 0);
|
||||
|
||||
mParam.backgroundLabelId = 0;
|
||||
mParam.numClasses = mNbClasses;
|
||||
mParam.keepTopK = mKeepTopK;
|
||||
mParam.scoreThreshold = mScoreThreshold;
|
||||
mParam.iouThreshold = mIOUThreshold;
|
||||
|
||||
mType = DataType::kFLOAT;
|
||||
}
|
||||
|
||||
int32_t GenerateDetection::getNbOutputs() const noexcept
|
||||
{
|
||||
return 1;
|
||||
}
|
||||
|
||||
int32_t GenerateDetection::initialize() noexcept
|
||||
{
|
||||
// Init the regWeight [10, 10, 5, 5]
|
||||
mRegWeightDevice = std::make_shared<CudaBind<float>>(4);
|
||||
PLUGIN_CUASSERT(cudaMemcpy(static_cast<void*>(mRegWeightDevice->mPtr),
|
||||
static_cast<void const*>(TLTMaskRCNNConfig::DETECTION_REG_WEIGHTS), sizeof(float) * 4, cudaMemcpyHostToDevice));
|
||||
|
||||
//@Init the mValidCnt and mDecodedBboxes for max batch size
|
||||
std::vector<int32_t> tempValidCnt(mMaxBatchSize, mAnchorsCnt);
|
||||
|
||||
mValidCnt = std::make_shared<CudaBind<int32_t>>(mMaxBatchSize);
|
||||
|
||||
PLUGIN_CUASSERT(cudaMemcpy(mValidCnt->mPtr, static_cast<void*>(tempValidCnt.data()),
|
||||
sizeof(int32_t) * mMaxBatchSize, cudaMemcpyHostToDevice));
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
void GenerateDetection::terminate() noexcept {}
|
||||
|
||||
void GenerateDetection::destroy() noexcept
|
||||
{
|
||||
delete this;
|
||||
}
|
||||
|
||||
bool GenerateDetection::supportsFormat(DataType type, PluginFormat format) const noexcept
|
||||
{
|
||||
return (type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
|
||||
}
|
||||
|
||||
char const* GenerateDetection::getPluginType() const noexcept
|
||||
{
|
||||
return "GenerateDetection_TRT";
|
||||
}
|
||||
|
||||
char const* GenerateDetection::getPluginVersion() const noexcept
|
||||
{
|
||||
return "1";
|
||||
}
|
||||
|
||||
IPluginV2Ext* GenerateDetection::clone() const noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
return new GenerateDetection(*this);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
void GenerateDetection::setPluginNamespace(char const* libNamespace) noexcept
|
||||
{
|
||||
mNameSpace = libNamespace;
|
||||
}
|
||||
|
||||
char const* GenerateDetection::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mNameSpace.c_str();
|
||||
}
|
||||
|
||||
size_t GenerateDetection::getSerializationSize() const noexcept
|
||||
{
|
||||
return sizeof(int32_t) * 2 + sizeof(float) * 2 + sizeof(int32_t) * 2 + sizeof(nvinfer1::Dims);
|
||||
}
|
||||
|
||||
void GenerateDetection::serialize(void* buffer) const noexcept
|
||||
{
|
||||
char *d = reinterpret_cast<char*>(buffer), *a = d;
|
||||
write(d, mNbClasses);
|
||||
write(d, mKeepTopK);
|
||||
write(d, mScoreThreshold);
|
||||
write(d, mIOUThreshold);
|
||||
write(d, mMaxBatchSize);
|
||||
write(d, mAnchorsCnt);
|
||||
write(d, mImageSize);
|
||||
PLUGIN_ASSERT(d == a + getSerializationSize());
|
||||
}
|
||||
|
||||
GenerateDetection::GenerateDetection(void const* data, size_t length)
|
||||
{
|
||||
deserialize(static_cast<int8_t const*>(data), length);
|
||||
}
|
||||
|
||||
void GenerateDetection::deserialize(int8_t const* data, size_t length)
|
||||
{
|
||||
auto const* d{data};
|
||||
int32_t num_classes = read<int32_t>(d);
|
||||
int32_t keep_topk = read<int32_t>(d);
|
||||
float score_threshold = read<float>(d);
|
||||
float iou_threshold = read<float>(d);
|
||||
mMaxBatchSize = read<int32_t>(d);
|
||||
mAnchorsCnt = read<int32_t>(d);
|
||||
mImageSize = read<nvinfer1::Dims3>(d);
|
||||
PLUGIN_VALIDATE(d == data + length);
|
||||
|
||||
mNbClasses = num_classes;
|
||||
mKeepTopK = keep_topk;
|
||||
mScoreThreshold = score_threshold;
|
||||
mIOUThreshold = iou_threshold;
|
||||
|
||||
mParam.backgroundLabelId = 0;
|
||||
mParam.numClasses = mNbClasses;
|
||||
mParam.keepTopK = mKeepTopK;
|
||||
mParam.scoreThreshold = mScoreThreshold;
|
||||
mParam.iouThreshold = mIOUThreshold;
|
||||
|
||||
mType = DataType::kFLOAT;
|
||||
}
|
||||
|
||||
void GenerateDetection::check_valid_inputs(nvinfer1::Dims const* inputs, int32_t nbInputDims) noexcept
|
||||
{
|
||||
// classifier_delta_bbox[N, anchors, num_classes*4, 1, 1]
|
||||
// classifier_class[N, anchors, num_classes, 1, 1]
|
||||
// rpn_rois[N, anchors, 4]
|
||||
PLUGIN_ASSERT(nbInputDims == 3);
|
||||
|
||||
// score
|
||||
PLUGIN_ASSERT(inputs[1].nbDims == 4 && inputs[1].d[1] == mNbClasses);
|
||||
// delta_bbox
|
||||
PLUGIN_ASSERT(inputs[0].nbDims == 4 && inputs[0].d[1] == mNbClasses * 4);
|
||||
// roi
|
||||
PLUGIN_ASSERT(inputs[2].nbDims == 2 && inputs[2].d[1] == 4);
|
||||
}
|
||||
|
||||
size_t GenerateDetection::getWorkspaceSize(int32_t batch_size) const noexcept
|
||||
{
|
||||
RefineDetectionWorkSpace refine(batch_size, mAnchorsCnt, mParam, mType);
|
||||
return refine.totalSize;
|
||||
}
|
||||
|
||||
Dims GenerateDetection::getOutputDimensions(int32_t index, Dims const* inputs, int32_t nbInputDims) noexcept
|
||||
{
|
||||
|
||||
check_valid_inputs(inputs, nbInputDims);
|
||||
PLUGIN_ASSERT(index == 0);
|
||||
|
||||
return {2, {mKeepTopK, 6}};
|
||||
}
|
||||
|
||||
int32_t GenerateDetection::enqueue(
|
||||
int32_t batch_size, void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
|
||||
{
|
||||
|
||||
void* detections = outputs[0];
|
||||
|
||||
// refine detection
|
||||
RefineDetectionWorkSpace refDetcWorkspace(batch_size, mAnchorsCnt, mParam, mType);
|
||||
cudaError_t status
|
||||
= DetectionPostProcess(stream, batch_size, mAnchorsCnt, static_cast<float*>(mRegWeightDevice->mPtr),
|
||||
static_cast<float>(mImageSize.d[1]), // Image Height
|
||||
static_cast<float>(mImageSize.d[2]), // Image Width
|
||||
DataType::kFLOAT, // mType,
|
||||
mParam, refDetcWorkspace, workspace,
|
||||
inputs[1], // inputs[InScore]
|
||||
inputs[0], // inputs[InDelta],
|
||||
mValidCnt->mPtr, // inputs[InCountValid],
|
||||
inputs[2], // inputs[ROI]
|
||||
detections);
|
||||
|
||||
PLUGIN_ASSERT(status == cudaSuccess);
|
||||
return status;
|
||||
}
|
||||
|
||||
DataType GenerateDetection::getOutputDataType(
|
||||
int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept
|
||||
{
|
||||
// Only DataType::kFLOAT is acceptable by the plugin layer
|
||||
return DataType::kFLOAT;
|
||||
}
|
||||
|
||||
// Configure the layer with input and output data types.
|
||||
void GenerateDetection::configurePlugin(Dims const* inputDims, int32_t nbInputs, Dims const* outputDims,
|
||||
int32_t nbOutputs, DataType const* inputTypes, DataType const* outputTypes, bool const* inputIsBroadcast,
|
||||
bool const* outputIsBroadcast, PluginFormat floatFormat, int32_t maxBatchSize) noexcept
|
||||
{
|
||||
check_valid_inputs(inputDims, nbInputs);
|
||||
PLUGIN_ASSERT(inputDims[0].d[0] == inputDims[1].d[0] && inputDims[1].d[0] == inputDims[2].d[0]);
|
||||
|
||||
mAnchorsCnt = inputDims[2].d[0];
|
||||
mType = inputTypes[0];
|
||||
mMaxBatchSize = maxBatchSize;
|
||||
}
|
||||
|
||||
// Attach the plugin object to an execution context and grant the plugin the access to some context resource.
|
||||
void GenerateDetection::attachToContext(
|
||||
cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) noexcept
|
||||
{
|
||||
}
|
||||
|
||||
// Detach the plugin object from its execution context.
|
||||
void GenerateDetection::detachFromContext() noexcept {}
|
||||
@@ -0,0 +1,138 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#ifndef TRT_GENERATE_DETECTION_PLUGIN_H
|
||||
#define TRT_GENERATE_DETECTION_PLUGIN_H
|
||||
#include <cuda_runtime_api.h>
|
||||
#include <memory>
|
||||
#include <string.h>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "NvInfer.h"
|
||||
#include "NvInferPlugin.h"
|
||||
#include "common/kernels/maskRCNNKernels.h"
|
||||
#include "multilevelProposeROI/tlt_mrcnn_config.h"
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
|
||||
class GenerateDetection : public IPluginV2Ext
|
||||
{
|
||||
public:
|
||||
GenerateDetection(int32_t num_classes, int32_t keep_topk, float score_threshold, float iou_threshold,
|
||||
nvinfer1::Dims const& image_size);
|
||||
|
||||
GenerateDetection(void const* data, size_t length);
|
||||
|
||||
~GenerateDetection() noexcept override = default;
|
||||
|
||||
int32_t getNbOutputs() const noexcept override;
|
||||
|
||||
Dims getOutputDimensions(int32_t index, Dims const* inputs, int32_t nbInputDims) noexcept override;
|
||||
|
||||
int32_t initialize() noexcept override;
|
||||
|
||||
void terminate() noexcept override;
|
||||
|
||||
void destroy() noexcept override;
|
||||
|
||||
size_t getWorkspaceSize(int32_t maxBatchSize) const noexcept override;
|
||||
|
||||
int32_t enqueue(int32_t batch_size, void const* const* inputs, void* const* outputs, void* workspace,
|
||||
cudaStream_t stream) noexcept override;
|
||||
|
||||
size_t getSerializationSize() const noexcept override;
|
||||
|
||||
void serialize(void* buffer) const noexcept override;
|
||||
|
||||
bool supportsFormat(DataType type, PluginFormat format) const noexcept override;
|
||||
|
||||
char const* getPluginType() const noexcept override;
|
||||
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
|
||||
IPluginV2Ext* clone() const noexcept override;
|
||||
|
||||
void setPluginNamespace(char const* libNamespace) noexcept override;
|
||||
|
||||
char const* getPluginNamespace() const noexcept override;
|
||||
|
||||
DataType getOutputDataType(
|
||||
int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept override;
|
||||
|
||||
void attachToContext(
|
||||
cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) noexcept override;
|
||||
|
||||
void configurePlugin(Dims const* inputDims, int32_t nbInputs, Dims const* outputDims, int32_t nbOutputs,
|
||||
DataType const* inputTypes, DataType const* outputTypes, bool const* inputIsBroadcast,
|
||||
bool const* outputIsBroadcast, PluginFormat floatFormat, int32_t maxBatchSize) noexcept override;
|
||||
|
||||
void detachFromContext() noexcept override;
|
||||
|
||||
private:
|
||||
void deserialize(int8_t const* data, size_t length);
|
||||
void check_valid_inputs(nvinfer1::Dims const* inputs, int32_t nbInputDims) noexcept;
|
||||
|
||||
int32_t mBackgroundLabel{};
|
||||
int32_t mNbClasses{};
|
||||
int32_t mKeepTopK{};
|
||||
float mScoreThreshold{};
|
||||
float mIOUThreshold{};
|
||||
|
||||
int32_t mMaxBatchSize{};
|
||||
int32_t mAnchorsCnt{};
|
||||
std::shared_ptr<CudaBind<int32_t>> mValidCnt; // valid cnt = number of input rois for every image.
|
||||
nvinfer1::DataType mType{};
|
||||
RefineNMSParameters mParam{};
|
||||
std::shared_ptr<CudaBind<float>> mRegWeightDevice;
|
||||
|
||||
nvinfer1::Dims mImageSize{};
|
||||
|
||||
std::string mNameSpace;
|
||||
};
|
||||
|
||||
class GenerateDetectionPluginCreator : public nvinfer1::pluginInternal::BaseCreator
|
||||
{
|
||||
public:
|
||||
GenerateDetectionPluginCreator() noexcept;
|
||||
|
||||
~GenerateDetectionPluginCreator() noexcept override {}
|
||||
|
||||
char const* getPluginName() const noexcept override;
|
||||
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
|
||||
PluginFieldCollection const* getFieldNames() noexcept override;
|
||||
|
||||
IPluginV2Ext* createPlugin(char const* name, PluginFieldCollection const* fc) noexcept override;
|
||||
|
||||
IPluginV2Ext* deserializePlugin(char const* name, void const* data, size_t length) noexcept override;
|
||||
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
int32_t mNbClasses{};
|
||||
int32_t mKeepTopK{};
|
||||
float mScoreThreshold{};
|
||||
float mIOUThreshold{};
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
};
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
#endif // TRT_GENERATE_DETECTION_PLUGIN_H
|
||||
Reference in New Issue
Block a user