diff --git a/yolov8/src/model.cpp b/yolov8/src/model.cpp index 316b1301..0b105c97 100644 --- a/yolov8/src/model.cpp +++ b/yolov8/src/model.cpp @@ -20,6 +20,15 @@ static int get_depth(int x, float gd) { return std::max(r, 1); } +void calculateStrides(nvinfer1::IElementWiseLayer* conv_layers[], int size, int reference_size, int strides[]) { + for (int i = 0; i < size; ++i) { + nvinfer1::ILayer* layer = conv_layers[i]; + nvinfer1::Dims dims = layer->getOutput(0)->getDimensions(); + int feature_map_size = dims.d[1]; + strides[i] = reference_size / feature_map_size; + } +} + static nvinfer1::IElementWiseLayer* Proto(nvinfer1::INetworkDefinition* network, std::map& weightMap, nvinfer1::ITensor& input, std::string lname, float gw, int max_channels) { @@ -220,7 +229,9 @@ nvinfer1::IHostMemory* buildEngineYolov8Det(nvinfer1::IBuilder* builder, nvinfer ********************************************* YOLOV8 DETECT ****************************************** *******************************************************************************************************/ - int strides[] = {8, 16, 32}; + nvinfer1::IElementWiseLayer* conv_layers[] = {conv3, conv5, conv7}; + int strides[sizeof(conv_layers) / sizeof(conv_layers[0])]; + calculateStrides(conv_layers, sizeof(conv_layers) / sizeof(conv_layers[0]), kInputH, strides); int stridesLength = sizeof(strides) / sizeof(int); nvinfer1::IShuffleLayer* shuffle22_0 = network->addShuffle(*cat22_0->getOutput(0)); @@ -507,7 +518,9 @@ nvinfer1::IHostMemory* buildEngineYolov8DetP6(nvinfer1::IBuilder* builder, nvinf /******************************************************************************************************* ********************************************* YOLOV8 DETECT ****************************************** *******************************************************************************************************/ - int strides[] = {8, 16, 32, 64}; + nvinfer1::IElementWiseLayer* conv_layers[] = {conv3, conv5, conv7, conv9}; + int strides[sizeof(conv_layers) / sizeof(conv_layers[0])]; + calculateStrides(conv_layers, sizeof(conv_layers) / sizeof(conv_layers[0]), kInputH, strides); int stridesLength = sizeof(strides) / sizeof(int); // P3 processing steps (remains unchanged) @@ -817,7 +830,9 @@ nvinfer1::IHostMemory* buildEngineYolov8DetP2(nvinfer1::IBuilder* builder, nvinf ********************************************* YOLOV8 DETECT ****************************************** *******************************************************************************************************/ - int strides[] = {4, 8, 16, 32}; + nvinfer1::IElementWiseLayer* conv_layers[] = {conv1, conv3, conv5, conv7}; + int strides[sizeof(conv_layers) / sizeof(conv_layers[0])]; + calculateStrides(conv_layers, sizeof(conv_layers) / sizeof(conv_layers[0]), kInputH, strides); int stridesLength = sizeof(strides) / sizeof(int); // P2 processing steps (remains unchanged) @@ -1148,7 +1163,9 @@ nvinfer1::IHostMemory* buildEngineYolov8Seg(nvinfer1::IBuilder* builder, nvinfer ********************************************* YOLOV8 DETECT ****************************************** *******************************************************************************************************/ - int strides[] = {8, 16, 32}; + nvinfer1::IElementWiseLayer* conv_layers[] = {conv3, conv5, conv7}; + int strides[sizeof(conv_layers) / sizeof(conv_layers[0])]; + calculateStrides(conv_layers, sizeof(conv_layers) / sizeof(conv_layers[0]), kInputH, strides); int stridesLength = sizeof(strides) / sizeof(int); nvinfer1::IShuffleLayer* shuffle22_0 = network->addShuffle(*cat22_0->getOutput(0));