Fix assertion and documetation for batchedNMSPlugin (#407)

* Fix assertion and documetation for batchedNMSPlugin Update the assertion to validate input dimensions to handle UFF. Clarify the purpose of clipBoxes field in plugin documentation. Signed-off-by: Rajeev Rao <[email protected]> * Address review comments Co-authored-by: Rajeev Rao <[email protected]>
NVIDIA · Mar 3, 2020 · b447156 · b447156
1 parent f9dd81a
commit b447156
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/plugin/batchedNMSPlugin/README.md b/plugin/batchedNMSPlugin/README.md
@@ -63,6 +63,7 @@ The `batchedNMSPlugin` is created using `BatchedNMSPluginCreator` with `NMSParam
 |`float`   |`scoreThreshold`          |The scalar threshold for score (low scoring boxes are removed).
 |`float`   |`iouThreshold`            |The scalar threshold for IOU (new boxes that have high IOU overlap with previously selected boxes are removed).
 |`bool`    |`isNormalized`            |Set to `false` if the box coordinates are not normalized, meaning they are not in the range `[0,1]`. Defaults to `true`.
+|`bool`    |`clipBoxes`               |Forcibly restrict bounding boxes to the normalized range `[0,1]`. Only applicable if `isNormalized` is also `true`. Defaults to `true`.
 
 
 ## Algorithms

diff --git a/plugin/batchedNMSPlugin/batchedNMSPlugin.cpp b/plugin/batchedNMSPlugin/batchedNMSPlugin.cpp
@@ -68,7 +68,7 @@ Dims BatchedNMSPlugin::getOutputDimensions(int index, const Dims* inputs, int nb
     ASSERT(nbInputDims == 2);
     ASSERT(index >= 0 && index < this->getNbOutputs());
     ASSERT(inputs[0].nbDims == 3);
-    ASSERT(inputs[1].nbDims == 2);
+    ASSERT(inputs[1].nbDims == 2 || (inputs[1].nbDims == 3 && inputs[1].d[2] == 1));
     // boxesSize: number of box coordinates for one sample
     boxesSize = inputs[0].d[0] * inputs[0].d[1] * inputs[0].d[2];
     // scoresSize: number of scores for one sample
@@ -141,7 +141,7 @@ void BatchedNMSPlugin::configurePlugin(const Dims* inputDims, int nbInputs, cons
     ASSERT(nbInputs == 2);
     ASSERT(nbOutputs == 4);
     ASSERT(inputDims[0].nbDims == 3);
-    ASSERT(inputDims[1].nbDims == 2);
+    ASSERT(inputDims[1].nbDims == 2 || (inputDims[1].nbDims == 3 && inputDims[1].d[2] == 1));
     ASSERT(std::none_of(inputIsBroadcast, inputIsBroadcast + nbInputs, [](bool b) { return b; }));
     ASSERT(std::none_of(outputIsBroadcast, outputIsBroadcast + nbInputs, [](bool b) { return b; }));