New optimized NMS

2022-08-15 02:48:23 -03:00
parent 8ad99d3f8d
commit c8a4a49f16
24 changed files with 206 additions and 394 deletions
--- a/nvdsinfer_custom_impl_Yolo/layers/route_layer.cpp
+++ b/nvdsinfer_custom_impl_Yolo/layers/route_layer.cpp
@@ -45,6 +45,9 @@ nvinfer1::ITensor* routeLayer(
    }
    layers += std::to_string(idxLayers[idxLayers.size() - 1]);

+    if (concatInputs.size() == 1)
+        return concatInputs[0];
+
    int axis = 0;
    if (block.find("axis") != block.end())
        axis = std::stoi(block.at("axis"));