Optional nms methods supporting for rcnn (#1274)
* Support to output the instance segmentation results on an original input image, rather than on a padded one. * modifications according to freedenS's suggestions * modification for padding * merge preprocessImg and preprocessImg_ and delete all the debug codes. * bug of reference passing * make the vars consistent with the previous rcnn.cpp and update the README * support different nms methods * cxx11 * add contributors and delete interpretations * update README * fix some problems * fix the bugs --------- Co-authored-by: benkang wang <benkang.wang@uisee.com>
This commit is contained in:
HaiyangP
GitHub
parent
30de4562bc
commit
877bb24389
@ -26,7 +26,7 @@ namespace cub = thrust::cuda_cub::cub;
|
||||
namespace nvinfer1 {
|
||||
|
||||
__global__ void batched_nms_kernel(
|
||||
const float threshold, const int num_detections,
|
||||
const int nms_method, const float threshold, const int num_detections,
|
||||
const int *indices, float *scores, const float *classes, const float4 *boxes) {
|
||||
|
||||
// Go through detections by descending score
|
||||
@ -50,18 +50,44 @@ __global__ void batched_nms_kernel(
|
||||
float marea = (mbox.z - mbox.x) * (mbox.w - mbox.y);
|
||||
float inter = w * h;
|
||||
float overlap = inter / (iarea + marea - inter);
|
||||
if (overlap > threshold) {
|
||||
scores[i] = 0.0f;
|
||||
float sigma = 0.5; // this is an empirical value
|
||||
// printf("nms_method: %d", nms_method);
|
||||
//nms methods selection in the second stage
|
||||
// 0: original nms
|
||||
// 1: soft-nms (linear)
|
||||
// 2: soft-nms (gaussian)
|
||||
// printf("nms_method: ", nms_method);
|
||||
switch (nms_method)
|
||||
{
|
||||
case 0:
|
||||
if (overlap > threshold) {
|
||||
scores[i] = 0.0f;
|
||||
}
|
||||
break;
|
||||
case 1:
|
||||
if (overlap > threshold) {
|
||||
scores[i] = (1 - overlap) * scores[i];
|
||||
}
|
||||
break;
|
||||
case 2:
|
||||
if (overlap > threshold) {
|
||||
scores[i] = std::exp(-(overlap * overlap) / sigma) * scores[i];
|
||||
}
|
||||
break;
|
||||
default:
|
||||
if (overlap > threshold) {
|
||||
scores[i] = 0.0f;
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Sync discarded detections
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
int batchedNms(int batch_size,
|
||||
int batchedNms(int nms_method, int batch_size,
|
||||
const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs,
|
||||
size_t count, int detections_per_im, float nms_thresh,
|
||||
void *workspace, size_t workspace_size, cudaStream_t stream) {
|
||||
@ -111,8 +137,9 @@ int batchedNms(int batch_size,
|
||||
// Launch actual NMS kernel - 1 block with each thread handling n detections
|
||||
// TODO: different device has differnet max threads
|
||||
const int max_threads = 1024;
|
||||
|
||||
int num_per_thread = ceil(static_cast<float>(num_detections) / max_threads);
|
||||
batched_nms_kernel << <num_per_thread, max_threads, 0, stream >> > (nms_thresh, num_detections,
|
||||
batched_nms_kernel << <num_per_thread, max_threads, 0, stream >> > (nms_method, nms_thresh, num_detections,
|
||||
indices_sorted, scores_sorted, in_classes, in_boxes);
|
||||
|
||||
// Re-sort with updated scores
|
||||
|
||||
@ -13,7 +13,7 @@ using namespace nvinfer1;
|
||||
#define PLUGIN_NAMESPACE ""
|
||||
|
||||
namespace nvinfer1 {
|
||||
int batchedNms(int batchSize,
|
||||
int batchedNms(int nms_method, int batchSize,
|
||||
const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs,
|
||||
size_t count, int detections_per_im, float nms_thresh,
|
||||
void *workspace, size_t workspace_size, cudaStream_t stream);
|
||||
@ -28,6 +28,7 @@ int batchedNms(int batchSize,
|
||||
Description: implement batched nms
|
||||
*/
|
||||
class BatchedNmsPlugin : public IPluginV2Ext {
|
||||
int _nms_method;
|
||||
float _nms_thresh;
|
||||
int _detections_per_im;
|
||||
|
||||
@ -36,32 +37,36 @@ class BatchedNmsPlugin : public IPluginV2Ext {
|
||||
protected:
|
||||
void deserialize(void const* data, size_t length) {
|
||||
const char* d = static_cast<const char*>(data);
|
||||
read(d, _nms_method);
|
||||
read(d, _nms_thresh);
|
||||
read(d, _detections_per_im);
|
||||
read(d, _count);
|
||||
}
|
||||
|
||||
size_t getSerializationSize() const TRT_NOEXCEPT override {
|
||||
return sizeof(_nms_thresh) + sizeof(_detections_per_im)
|
||||
size_t getSerializationSize() const override {
|
||||
return sizeof(_nms_method) + sizeof(_nms_thresh) + sizeof(_detections_per_im)
|
||||
+ sizeof(_count);
|
||||
}
|
||||
|
||||
void serialize(void *buffer) const TRT_NOEXCEPT override {
|
||||
char* d = static_cast<char*>(buffer);
|
||||
write(d, _nms_method);
|
||||
write(d, _nms_thresh);
|
||||
write(d, _detections_per_im);
|
||||
write(d, _count);
|
||||
}
|
||||
|
||||
public:
|
||||
BatchedNmsPlugin(float nms_thresh, int detections_per_im)
|
||||
: _nms_thresh(nms_thresh), _detections_per_im(detections_per_im) {
|
||||
BatchedNmsPlugin(int nms_method, float nms_thresh, int detections_per_im)
|
||||
: _nms_method(nms_method), _nms_thresh(nms_thresh), _detections_per_im(detections_per_im) {
|
||||
assert(nms_method >= 0);
|
||||
assert(nms_thresh > 0);
|
||||
assert(detections_per_im > 0);
|
||||
}
|
||||
|
||||
BatchedNmsPlugin(float nms_thresh, int detections_per_im, size_t count)
|
||||
: _nms_thresh(nms_thresh), _detections_per_im(detections_per_im), _count(count) {
|
||||
BatchedNmsPlugin(int nms_method, float nms_thresh, int detections_per_im, size_t count)
|
||||
: _nms_method(nms_method), _nms_thresh(nms_thresh), _detections_per_im(detections_per_im), _count(count) {
|
||||
assert(nms_method >= 0);
|
||||
assert(nms_thresh > 0);
|
||||
assert(detections_per_im > 0);
|
||||
assert(count > 0);
|
||||
@ -101,7 +106,7 @@ class BatchedNmsPlugin : public IPluginV2Ext {
|
||||
size_t getWorkspaceSize(int maxBatchSize) const TRT_NOEXCEPT override {
|
||||
static int size = -1;
|
||||
if (size < 0) {
|
||||
size = batchedNms(maxBatchSize, nullptr, nullptr, _count,
|
||||
size = batchedNms(_nms_method, maxBatchSize, nullptr, nullptr, _count,
|
||||
_detections_per_im, _nms_thresh,
|
||||
nullptr, 0, nullptr);
|
||||
}
|
||||
@ -111,7 +116,7 @@ class BatchedNmsPlugin : public IPluginV2Ext {
|
||||
int enqueue(int batchSize,
|
||||
const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs,
|
||||
void *workspace, cudaStream_t stream) TRT_NOEXCEPT override {
|
||||
return batchedNms(batchSize, inputs, outputs, _count,
|
||||
return batchedNms(_nms_method, batchSize, inputs, outputs, _count,
|
||||
_detections_per_im, _nms_thresh,
|
||||
workspace, getWorkspaceSize(batchSize), stream);
|
||||
}
|
||||
@ -152,7 +157,7 @@ class BatchedNmsPlugin : public IPluginV2Ext {
|
||||
}
|
||||
|
||||
IPluginV2Ext *clone() const TRT_NOEXCEPT override {
|
||||
return new BatchedNmsPlugin(_nms_thresh, _detections_per_im, _count);
|
||||
return new BatchedNmsPlugin(_nms_method, _nms_thresh, _detections_per_im, _count);
|
||||
}
|
||||
|
||||
private:
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
# Rcnn
|
||||
|
||||
The Pytorch implementation is [facebookresearch/detectron2](https://github.com/facebookresearch/detectron2). Now, outputting instance segmentation results on the original image size is available, which is more convenient for engineering applications.
|
||||
The Pytorch implementation is [facebookresearch/detectron2](https://github.com/facebookresearch/detectron2). Now, outputting instance segmentation results on the original image size and selecting different nms methods are available, which is more convenient for engineering applications.
|
||||
|
||||
## Models
|
||||
|
||||
@ -21,7 +21,9 @@ TensorRT 8.x is supported and you can use it.
|
||||
|
||||
## Contributors
|
||||
|
||||
<a href="https://github.com/HaiyangPeng"><img src="https://avatars.githubusercontent.com/u/46739135?v=4" width="40px;" alt=""/></a>
|
||||
<a href="https://github.com/nengwp"><img src="https://avatars.githubusercontent.com/u/44516353?s=96&v=4" width="40px;" alt=""/></a>
|
||||
<a href="https://github.com/freedenS"><img src="https://avatars.githubusercontent.com/u/26213470?v=4" width="40px;" alt=""/></a>
|
||||
|
||||
## How to Run
|
||||
|
||||
|
||||
@ -52,6 +52,13 @@ static const char* INPUT_NODE_NAME = "images";
|
||||
static const std::vector<std::string> OUTPUT_NAMES = { "scores", "boxes",
|
||||
"labels", "masks" };
|
||||
|
||||
//nms methods selection in the second stage
|
||||
// 0: original nms
|
||||
// 1: soft-nms (linear)
|
||||
// 2: soft-nms (gaussian)
|
||||
static int NMS_METHOD = 1;
|
||||
static std::vector<int> NMS_METHOD_VEC = {0, 1, 2};
|
||||
|
||||
std::vector<float> GenerateAnchors(const std::vector<float>& anchor_sizes,
|
||||
const std::vector<float>& aspect_ratios) {
|
||||
std::vector<float> res;
|
||||
@ -185,7 +192,7 @@ void BoxHead(INetworkDefinition *network, std::map<std::string, Weights>& weight
|
||||
// nms
|
||||
std::vector<ITensor*> nmsInput = { predictorDecodeLayer->getOutput(0),
|
||||
predictorDecodeLayer->getOutput(1), predictorDecodeLayer->getOutput(2) };
|
||||
auto batchedNmsPlugin = BatchedNmsPlugin(NMS_THRESH_TEST, DETECTIONS_PER_IMAGE);
|
||||
auto batchedNmsPlugin = BatchedNmsPlugin(NMS_METHOD, NMS_THRESH_TEST, DETECTIONS_PER_IMAGE);
|
||||
auto batchedNmsLayer = network->addPluginV2(nmsInput.data(), nmsInput.size(), batchedNmsPlugin);
|
||||
|
||||
// instances
|
||||
@ -375,6 +382,20 @@ bool parse_args(int argc, char** argv, std::string& wtsFile, std::string& engine
|
||||
}
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
|
||||
int flag = 0;
|
||||
for (int &item : NMS_METHOD_VEC) {
|
||||
if (item == NMS_METHOD) {
|
||||
flag = 1;
|
||||
printf("The nms method %d is applied.\n", NMS_METHOD);
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (flag == 0) {
|
||||
printf("[WARNING] The nms_method %d is not supported, please choose from [0, 1, 2].\n", NMS_METHOD);
|
||||
printf("[WARNING] To make the nms robust, the default nms method 0 is applied.\n");
|
||||
NMS_METHOD = 0;
|
||||
}
|
||||
|
||||
// calculate size
|
||||
calculateSize();
|
||||
|
||||
Loading…
Reference in New Issue
Block a user