diff --git a/yolov5/README.md b/yolov5/README.md index cdae023..8aef754 100644 --- a/yolov5/README.md +++ b/yolov5/README.md @@ -49,14 +49,8 @@ Currently, we support yolov5 v1.0, v2.0, v3.0, v3.1, v4.0, v5.0, v6.0, v6.2, v7. ## Config -- Choose the model n/s/m/l/x/n6/s6/m6/l6/x6 from command line arguments. -- Input shape defined in yololayer.h -- Number of classes defined in yololayer.h, **DO NOT FORGET TO ADAPT THIS, If using your own model** -- INT8/FP16/FP32 can be selected by the macro in yolov5.cpp, **INT8 need more steps, pls follow `How to Run` first and then go the `INT8 Quantization` below** -- GPU id can be selected by the macro in yolov5.cpp -- NMS thresh in yolov5.cpp -- BBox confidence thresh in yolov5.cpp -- Batch size in yolov5.cpp +- Choose the YOLOv5 sub-model n/s/m/l/x/n6/s6/m6/l6/x6 from command line arguments. +- Other configs please check src/config.h ## Build and Run @@ -83,14 +77,14 @@ cd build cp {ultralytics}/yolov5/yolov5s.wts {tensorrtx}/yolov5/build cmake .. make -sudo ./yolov5_det -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw] // serialize model to plan file -sudo ./yolov5_det -d [.engine] [image folder] // deserialize and run inference, the images in [image folder] will be processed. +./yolov5_det -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw] // serialize model to plan file +./yolov5_det -d [.engine] [image folder] // deserialize and run inference, the images in [image folder] will be processed. // For example yolov5s -sudo ./yolov5_det -s yolov5s.wts yolov5s.engine s -sudo ./yolov5_det -d yolov5s.engine ../samples +./yolov5_det -s yolov5s.wts yolov5s.engine s +./yolov5_det -d yolov5s.engine ../images // For example Custom model with depth_multiple=0.17, width_multiple=0.25 in yolov5.yaml -sudo ./yolov5_det -s yolov5_custom.wts yolov5.engine c 0.17 0.25 -sudo ./yolov5_det -d yolov5.engine ../samples +./yolov5_det -s yolov5_custom.wts yolov5.engine c 0.17 0.25 +./yolov5_det -d yolov5.engine ../images ``` 3. check the images generated, as follows. _zidane.jpg and _bus.jpg @@ -120,7 +114,7 @@ wget https://github.com/joannzhang00/ImageNet-dataset-classes-labels/blob/main/i ./yolov5_cls -s yolov5s-cls.wts yolov5s-cls.engine s # Run inference -./yolov5_cls -d yolov5s-cls.engine ../samples +./yolov5_cls -d yolov5s-cls.engine ../images ``` ### Instance Segmentation @@ -133,7 +127,7 @@ wget https://github.com/joannzhang00/ImageNet-dataset-classes-labels/blob/main/i wget -O coco.txt https://raw.githubusercontent.com/amikelive/coco-labels/master/coco-labels-2014_2017.txt # Run inference with labels file -./yolov5_seg -d yolov5s-seg.engine ../samples coco.txt +./yolov5_seg -d yolov5s-seg.engine ../images coco.txt ```

@@ -146,7 +140,7 @@ wget -O coco.txt https://raw.githubusercontent.com/amikelive/coco-labels/master/ 2. unzip it in yolov5/build -3. set the macro `USE_INT8` in yolov5.cpp and make +3. set the macro `USE_INT8` in src/config.h and make 4. serialize the model and test diff --git a/yolov5/images b/yolov5/images new file mode 120000 index 0000000..02cc755 --- /dev/null +++ b/yolov5/images @@ -0,0 +1 @@ +../yolov3-spp/samples \ No newline at end of file diff --git a/yolov5/plugin/yololayer.cu b/yolov5/plugin/yololayer.cu index 02dcd4c..d80a9a4 100644 --- a/yolov5/plugin/yololayer.cu +++ b/yolov5/plugin/yololayer.cu @@ -1,322 +1,280 @@ -#include -#include -#include #include "yololayer.h" #include "cuda_utils.h" -namespace Tn -{ - template - void write(char*& buffer, const T& val) - { - *reinterpret_cast(buffer) = val; - buffer += sizeof(T); - } +#include +#include +#include - template - void read(const char*& buffer, T& val) - { - val = *reinterpret_cast(buffer); - buffer += sizeof(T); - } +namespace Tn { +template +void write(char*& buffer, const T& val) { + *reinterpret_cast(buffer) = val; + buffer += sizeof(T); } -using namespace Yolo; - -namespace nvinfer1 -{ - YoloLayerPlugin::YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, bool is_segmentation, const std::vector& vYoloKernel) - { - mClassCount = classCount; - mYoloV5NetWidth = netWidth; - mYoloV5NetHeight = netHeight; - mMaxOutObject = maxOut; - is_segmentation_ = is_segmentation; - mYoloKernel = vYoloKernel; - mKernelCount = vYoloKernel.size(); - - CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*))); - size_t AnchorLen = sizeof(float)* CHECK_COUNT * 2; - for (int ii = 0; ii < mKernelCount; ii++) - { - CUDA_CHECK(cudaMalloc(&mAnchor[ii], AnchorLen)); - const auto& yolo = mYoloKernel[ii]; - CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice)); - } - } - YoloLayerPlugin::~YoloLayerPlugin() - { - for (int ii = 0; ii < mKernelCount; ii++) - { - CUDA_CHECK(cudaFree(mAnchor[ii])); - } - CUDA_CHECK(cudaFreeHost(mAnchor)); - } - - // create the plugin at runtime from a byte stream - YoloLayerPlugin::YoloLayerPlugin(const void* data, size_t length) - { - using namespace Tn; - const char *d = reinterpret_cast(data), *a = d; - read(d, mClassCount); - read(d, mThreadCount); - read(d, mKernelCount); - read(d, mYoloV5NetWidth); - read(d, mYoloV5NetHeight); - read(d, mMaxOutObject); - read(d, is_segmentation_); - mYoloKernel.resize(mKernelCount); - auto kernelSize = mKernelCount * sizeof(YoloKernel); - memcpy(mYoloKernel.data(), d, kernelSize); - d += kernelSize; - CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*))); - size_t AnchorLen = sizeof(float)* CHECK_COUNT * 2; - for (int ii = 0; ii < mKernelCount; ii++) - { - CUDA_CHECK(cudaMalloc(&mAnchor[ii], AnchorLen)); - const auto& yolo = mYoloKernel[ii]; - CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice)); - } - assert(d == a + length); - } - - void YoloLayerPlugin::serialize(void* buffer) const TRT_NOEXCEPT - { - using namespace Tn; - char* d = static_cast(buffer), *a = d; - write(d, mClassCount); - write(d, mThreadCount); - write(d, mKernelCount); - write(d, mYoloV5NetWidth); - write(d, mYoloV5NetHeight); - write(d, mMaxOutObject); - write(d, is_segmentation_); - auto kernelSize = mKernelCount * sizeof(YoloKernel); - memcpy(d, mYoloKernel.data(), kernelSize); - d += kernelSize; - - assert(d == a + getSerializationSize()); - } - - size_t YoloLayerPlugin::getSerializationSize() const TRT_NOEXCEPT - { - return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mKernelCount) + sizeof(Yolo::YoloKernel) * mYoloKernel.size() + sizeof(mYoloV5NetWidth) + sizeof(mYoloV5NetHeight) + sizeof(mMaxOutObject) + sizeof(is_segmentation_); - } - - int YoloLayerPlugin::initialize() TRT_NOEXCEPT - { - return 0; - } - - Dims YoloLayerPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRT_NOEXCEPT - { - //output the result to channel - int totalsize = mMaxOutObject * sizeof(Detection) / sizeof(float); - - return Dims3(totalsize + 1, 1, 1); - } - - // Set plugin namespace - void YoloLayerPlugin::setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT - { - mPluginNamespace = pluginNamespace; - } - - const char* YoloLayerPlugin::getPluginNamespace() const TRT_NOEXCEPT - { - return mPluginNamespace; - } - - // Return the DataType of the plugin output at the requested index - DataType YoloLayerPlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT - { - return DataType::kFLOAT; - } - - // Return true if output tensor is broadcast across a batch. - bool YoloLayerPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT - { - return false; - } - - // Return true if plugin can use input that is broadcast across batch without replication. - bool YoloLayerPlugin::canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT - { - return false; - } - - void YoloLayerPlugin::configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) TRT_NOEXCEPT - { - } - - // Attach the plugin object to an execution context and grant the plugin the access to some context resource. - void YoloLayerPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT - { - } - - // Detach the plugin object from its execution context. - void YoloLayerPlugin::detachFromContext() TRT_NOEXCEPT {} - - const char* YoloLayerPlugin::getPluginType() const TRT_NOEXCEPT - { - return "YoloLayer_TRT"; - } - - const char* YoloLayerPlugin::getPluginVersion() const TRT_NOEXCEPT - { - return "1"; - } - - void YoloLayerPlugin::destroy() TRT_NOEXCEPT - { - delete this; - } - - // Clone the plugin - IPluginV2IOExt* YoloLayerPlugin::clone() const TRT_NOEXCEPT - { - YoloLayerPlugin* p = new YoloLayerPlugin(mClassCount, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, is_segmentation_, mYoloKernel); - p->setPluginNamespace(mPluginNamespace); - return p; - } - - __device__ float Logist(float data) { return 1.0f / (1.0f + expf(-data)); }; - - __global__ void CalDetection(const float *input, float *output, int noElements, - const int netwidth, const int netheight, int maxoutobject, int yoloWidth, int yoloHeight, const float anchors[CHECK_COUNT * 2], int classes, int outputElem, bool is_segmentation) - { - - int idx = threadIdx.x + blockDim.x * blockIdx.x; - if (idx >= noElements) return; - - int total_grid = yoloWidth * yoloHeight; - int bnIdx = idx / total_grid; - idx = idx - total_grid * bnIdx; - int info_len_i = 5 + classes; - if (is_segmentation) info_len_i += 32; - const float* curInput = input + bnIdx * (info_len_i * total_grid * CHECK_COUNT); - - for (int k = 0; k < CHECK_COUNT; ++k) { - float box_prob = Logist(curInput[idx + k * info_len_i * total_grid + 4 * total_grid]); - if (box_prob < IGNORE_THRESH) continue; - int class_id = 0; - float max_cls_prob = 0.0; - for (int i = 5; i < 5 + classes; ++i) { - float p = Logist(curInput[idx + k * info_len_i * total_grid + i * total_grid]); - if (p > max_cls_prob) { - max_cls_prob = p; - class_id = i - 5; - } - } - float *res_count = output + bnIdx * outputElem; - int count = (int)atomicAdd(res_count, 1); - if (count >= maxoutobject) return; - char *data = (char*)res_count + sizeof(float) + count * sizeof(Detection); - Detection *det = (Detection*)(data); - - int row = idx / yoloWidth; - int col = idx % yoloWidth; - - //Location - // pytorch: - // y = x[i].sigmoid() - // y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy - // y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh - // X: (sigmoid(tx) + cx)/FeaturemapW * netwidth - det->bbox[0] = (col - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 0 * total_grid])) * netwidth / yoloWidth; - det->bbox[1] = (row - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 1 * total_grid])) * netheight / yoloHeight; - - // W: (Pw * e^tw) / FeaturemapW * netwidth - // v5: https://github.com/ultralytics/yolov5/issues/471 - det->bbox[2] = 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 2 * total_grid]); - det->bbox[2] = det->bbox[2] * det->bbox[2] * anchors[2 * k]; - det->bbox[3] = 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 3 * total_grid]); - det->bbox[3] = det->bbox[3] * det->bbox[3] * anchors[2 * k + 1]; - det->conf = box_prob * max_cls_prob; - det->class_id = class_id; - - for (int i = 0; is_segmentation && i < 32; i++) { - det->mask[i] = curInput[idx + k * info_len_i * total_grid + (i + 5 + classes) * total_grid]; - } - } - } - - void YoloLayerPlugin::forwardGpu(const float* const* inputs, float *output, cudaStream_t stream, int batchSize) - { - int outputElem = 1 + mMaxOutObject * sizeof(Detection) / sizeof(float); - for (int idx = 0; idx < batchSize; ++idx) { - CUDA_CHECK(cudaMemsetAsync(output + idx * outputElem, 0, sizeof(float), stream)); - } - int numElem = 0; - for (unsigned int i = 0; i < mYoloKernel.size(); ++i) { - const auto& yolo = mYoloKernel[i]; - numElem = yolo.width * yolo.height * batchSize; - if (numElem < mThreadCount) mThreadCount = numElem; - - //printf("Net: %d %d \n", mYoloV5NetWidth, mYoloV5NetHeight); - CalDetection << < (numElem + mThreadCount - 1) / mThreadCount, mThreadCount, 0, stream >> > - (inputs[i], output, numElem, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, yolo.width, yolo.height, (float*)mAnchor[i], mClassCount, outputElem, is_segmentation_); - } - } - - - int YoloLayerPlugin::enqueue(int batchSize, const void* const* inputs, void* TRT_CONST_ENQUEUE* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT - { - forwardGpu((const float* const*)inputs, (float*)outputs[0], stream, batchSize); - return 0; - } - - PluginFieldCollection YoloPluginCreator::mFC{}; - std::vector YoloPluginCreator::mPluginAttributes; - - YoloPluginCreator::YoloPluginCreator() - { - mPluginAttributes.clear(); - - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); - } - - const char* YoloPluginCreator::getPluginName() const TRT_NOEXCEPT - { - return "YoloLayer_TRT"; - } - - const char* YoloPluginCreator::getPluginVersion() const TRT_NOEXCEPT - { - return "1"; - } - - const PluginFieldCollection* YoloPluginCreator::getFieldNames() TRT_NOEXCEPT - { - return &mFC; - } - - IPluginV2IOExt* YoloPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) TRT_NOEXCEPT - { - assert(fc->nbFields == 2); - assert(strcmp(fc->fields[0].name, "netinfo") == 0); - assert(strcmp(fc->fields[1].name, "kernels") == 0); - int *p_netinfo = (int*)(fc->fields[0].data); - int class_count = p_netinfo[0]; - int input_w = p_netinfo[1]; - int input_h = p_netinfo[2]; - int max_output_object_count = p_netinfo[3]; - bool is_segmentation = (bool)p_netinfo[4]; - std::vector kernels(fc->fields[1].length); - memcpy(&kernels[0], fc->fields[1].data, kernels.size() * sizeof(Yolo::YoloKernel)); - YoloLayerPlugin* obj = new YoloLayerPlugin(class_count, input_w, input_h, max_output_object_count, is_segmentation, kernels); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - - IPluginV2IOExt* YoloPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRT_NOEXCEPT - { - // This object will be deleted when the network is destroyed, which will - // call YoloLayerPlugin::destroy() - YoloLayerPlugin* obj = new YoloLayerPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } +template +void read(const char*& buffer, T& val) { + val = *reinterpret_cast(buffer); + buffer += sizeof(T); +} +} + +namespace nvinfer1 { +YoloLayerPlugin::YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, bool is_segmentation, const std::vector& vYoloKernel) { + mClassCount = classCount; + mYoloV5NetWidth = netWidth; + mYoloV5NetHeight = netHeight; + mMaxOutObject = maxOut; + is_segmentation_ = is_segmentation; + mYoloKernel = vYoloKernel; + mKernelCount = vYoloKernel.size(); + + CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*))); + size_t AnchorLen = sizeof(float)* kNumAnchor * 2; + for (int ii = 0; ii < mKernelCount; ii++) { + CUDA_CHECK(cudaMalloc(&mAnchor[ii], AnchorLen)); + const auto& yolo = mYoloKernel[ii]; + CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice)); + } +} + +YoloLayerPlugin::~YoloLayerPlugin() { + for (int ii = 0; ii < mKernelCount; ii++) { + CUDA_CHECK(cudaFree(mAnchor[ii])); + } + CUDA_CHECK(cudaFreeHost(mAnchor)); +} + +// create the plugin at runtime from a byte stream +YoloLayerPlugin::YoloLayerPlugin(const void* data, size_t length) { + using namespace Tn; + const char *d = reinterpret_cast(data), *a = d; + read(d, mClassCount); + read(d, mThreadCount); + read(d, mKernelCount); + read(d, mYoloV5NetWidth); + read(d, mYoloV5NetHeight); + read(d, mMaxOutObject); + read(d, is_segmentation_); + mYoloKernel.resize(mKernelCount); + auto kernelSize = mKernelCount * sizeof(YoloKernel); + memcpy(mYoloKernel.data(), d, kernelSize); + d += kernelSize; + CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*))); + size_t AnchorLen = sizeof(float)* kNumAnchor * 2; + for (int ii = 0; ii < mKernelCount; ii++) { + CUDA_CHECK(cudaMalloc(&mAnchor[ii], AnchorLen)); + const auto& yolo = mYoloKernel[ii]; + CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice)); + } + assert(d == a + length); +} + +void YoloLayerPlugin::serialize(void* buffer) const TRT_NOEXCEPT { + using namespace Tn; + char* d = static_cast(buffer), *a = d; + write(d, mClassCount); + write(d, mThreadCount); + write(d, mKernelCount); + write(d, mYoloV5NetWidth); + write(d, mYoloV5NetHeight); + write(d, mMaxOutObject); + write(d, is_segmentation_); + auto kernelSize = mKernelCount * sizeof(YoloKernel); + memcpy(d, mYoloKernel.data(), kernelSize); + d += kernelSize; + + assert(d == a + getSerializationSize()); +} + +size_t YoloLayerPlugin::getSerializationSize() const TRT_NOEXCEPT { + size_t s = sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mKernelCount); + s += sizeof(YoloKernel) * mYoloKernel.size(); + s += sizeof(mYoloV5NetWidth) + sizeof(mYoloV5NetHeight); + s += sizeof(mMaxOutObject) + sizeof(is_segmentation_); + return s; +} + +int YoloLayerPlugin::initialize() TRT_NOEXCEPT { + return 0; +} + +Dims YoloLayerPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRT_NOEXCEPT { + //output the result to channel + int totalsize = mMaxOutObject * sizeof(Detection) / sizeof(float); + return Dims3(totalsize + 1, 1, 1); +} + +// Set plugin namespace +void YoloLayerPlugin::setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT { + mPluginNamespace = pluginNamespace; +} + +const char* YoloLayerPlugin::getPluginNamespace() const TRT_NOEXCEPT { + return mPluginNamespace; +} + +// Return the DataType of the plugin output at the requested index +DataType YoloLayerPlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT { + return DataType::kFLOAT; +} + +// Return true if output tensor is broadcast across a batch. +bool YoloLayerPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT { + return false; +} + +// Return true if plugin can use input that is broadcast across batch without replication. +bool YoloLayerPlugin::canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT { + return false; +} + +void YoloLayerPlugin::configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) TRT_NOEXCEPT {} + +// Attach the plugin object to an execution context and grant the plugin the access to some context resource. +void YoloLayerPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT {} + +// Detach the plugin object from its execution context. +void YoloLayerPlugin::detachFromContext() TRT_NOEXCEPT {} + +const char* YoloLayerPlugin::getPluginType() const TRT_NOEXCEPT { + return "YoloLayer_TRT"; +} + +const char* YoloLayerPlugin::getPluginVersion() const TRT_NOEXCEPT { + return "1"; +} + +void YoloLayerPlugin::destroy() TRT_NOEXCEPT { + delete this; +} + +// Clone the plugin +IPluginV2IOExt* YoloLayerPlugin::clone() const TRT_NOEXCEPT { + YoloLayerPlugin* p = new YoloLayerPlugin(mClassCount, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, is_segmentation_, mYoloKernel); + p->setPluginNamespace(mPluginNamespace); + return p; +} + +__device__ float Logist(float data) { return 1.0f / (1.0f + expf(-data)); }; + +__global__ void CalDetection(const float *input, float *output, int noElements, + const int netwidth, const int netheight, int maxoutobject, int yoloWidth, + int yoloHeight, const float anchors[kNumAnchor * 2], int classes, int outputElem, bool is_segmentation) { + + int idx = threadIdx.x + blockDim.x * blockIdx.x; + if (idx >= noElements) return; + + int total_grid = yoloWidth * yoloHeight; + int bnIdx = idx / total_grid; + idx = idx - total_grid * bnIdx; + int info_len_i = 5 + classes; + if (is_segmentation) info_len_i += 32; + const float* curInput = input + bnIdx * (info_len_i * total_grid * kNumAnchor); + + for (int k = 0; k < kNumAnchor; ++k) { + float box_prob = Logist(curInput[idx + k * info_len_i * total_grid + 4 * total_grid]); + if (box_prob < kIgnoreThresh) continue; + int class_id = 0; + float max_cls_prob = 0.0; + for (int i = 5; i < 5 + classes; ++i) { + float p = Logist(curInput[idx + k * info_len_i * total_grid + i * total_grid]); + if (p > max_cls_prob) { + max_cls_prob = p; + class_id = i - 5; + } + } + float *res_count = output + bnIdx * outputElem; + int count = (int)atomicAdd(res_count, 1); + if (count >= maxoutobject) return; + char *data = (char*)res_count + sizeof(float) + count * sizeof(Detection); + Detection *det = (Detection*)(data); + + int row = idx / yoloWidth; + int col = idx % yoloWidth; + + det->bbox[0] = (col - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 0 * total_grid])) * netwidth / yoloWidth; + det->bbox[1] = (row - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 1 * total_grid])) * netheight / yoloHeight; + + det->bbox[2] = 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 2 * total_grid]); + det->bbox[2] = det->bbox[2] * det->bbox[2] * anchors[2 * k]; + det->bbox[3] = 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 3 * total_grid]); + det->bbox[3] = det->bbox[3] * det->bbox[3] * anchors[2 * k + 1]; + det->conf = box_prob * max_cls_prob; + det->class_id = class_id; + + for (int i = 0; is_segmentation && i < 32; i++) { + det->mask[i] = curInput[idx + k * info_len_i * total_grid + (i + 5 + classes) * total_grid]; + } + } +} + +void YoloLayerPlugin::forwardGpu(const float* const* inputs, float *output, cudaStream_t stream, int batchSize) { + int outputElem = 1 + mMaxOutObject * sizeof(Detection) / sizeof(float); + for (int idx = 0; idx < batchSize; ++idx) { + CUDA_CHECK(cudaMemsetAsync(output + idx * outputElem, 0, sizeof(float), stream)); + } + int numElem = 0; + for (unsigned int i = 0; i < mYoloKernel.size(); ++i) { + const auto& yolo = mYoloKernel[i]; + numElem = yolo.width * yolo.height * batchSize; + if (numElem < mThreadCount) mThreadCount = numElem; + + CalDetection << < (numElem + mThreadCount - 1) / mThreadCount, mThreadCount, 0, stream >> > + (inputs[i], output, numElem, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, yolo.width, yolo.height, (float*)mAnchor[i], mClassCount, outputElem, is_segmentation_); + } +} + + +int YoloLayerPlugin::enqueue(int batchSize, const void* const* inputs, void* TRT_CONST_ENQUEUE* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT { + forwardGpu((const float* const*)inputs, (float*)outputs[0], stream, batchSize); + return 0; +} + +PluginFieldCollection YoloPluginCreator::mFC{}; +std::vector YoloPluginCreator::mPluginAttributes; + +YoloPluginCreator::YoloPluginCreator() { + mPluginAttributes.clear(); + mFC.nbFields = mPluginAttributes.size(); + mFC.fields = mPluginAttributes.data(); +} + +const char* YoloPluginCreator::getPluginName() const TRT_NOEXCEPT { + return "YoloLayer_TRT"; +} + +const char* YoloPluginCreator::getPluginVersion() const TRT_NOEXCEPT { + return "1"; +} + +const PluginFieldCollection* YoloPluginCreator::getFieldNames() TRT_NOEXCEPT { + return &mFC; +} + +IPluginV2IOExt* YoloPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) TRT_NOEXCEPT { + assert(fc->nbFields == 2); + assert(strcmp(fc->fields[0].name, "netinfo") == 0); + assert(strcmp(fc->fields[1].name, "kernels") == 0); + int *p_netinfo = (int*)(fc->fields[0].data); + int class_count = p_netinfo[0]; + int input_w = p_netinfo[1]; + int input_h = p_netinfo[2]; + int max_output_object_count = p_netinfo[3]; + bool is_segmentation = (bool)p_netinfo[4]; + std::vector kernels(fc->fields[1].length); + memcpy(&kernels[0], fc->fields[1].data, kernels.size() * sizeof(YoloKernel)); + YoloLayerPlugin* obj = new YoloLayerPlugin(class_count, input_w, input_h, max_output_object_count, is_segmentation, kernels); + obj->setPluginNamespace(mNamespace.c_str()); + return obj; +} + +IPluginV2IOExt* YoloPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRT_NOEXCEPT { + // This object will be deleted when the network is destroyed, which will + // call YoloLayerPlugin::destroy() + YoloLayerPlugin* obj = new YoloLayerPlugin(serialData, serialLength); + obj->setPluginNamespace(mNamespace.c_str()); + return obj; +} } diff --git a/yolov5/plugin/yololayer.h b/yolov5/plugin/yololayer.h index cbf6427..a73190b 100644 --- a/yolov5/plugin/yololayer.h +++ b/yolov5/plugin/yololayer.h @@ -1,140 +1,106 @@ -#ifndef _YOLO_LAYER_H -#define _YOLO_LAYER_H +#pragma once + +#include "types.h" +#include "macros.h" #include #include -#include -#include "macros.h" -namespace Yolo -{ - static constexpr int CHECK_COUNT = 3; - static constexpr float IGNORE_THRESH = 0.1f; - struct YoloKernel - { - int width; - int height; - float anchors[CHECK_COUNT * 2]; - }; - static constexpr int MAX_OUTPUT_BBOX_COUNT = 1000; - static constexpr int CLASS_NUM = 80; - static constexpr int INPUT_H = 640; // yolov5's input height and width must be divisible by 32. - static constexpr int INPUT_W = 640; +namespace nvinfer1 { +class API YoloLayerPlugin : public IPluginV2IOExt { +public: + YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, bool is_segmentation, const std::vector& vYoloKernel); + YoloLayerPlugin(const void* data, size_t length); + ~YoloLayerPlugin(); - static constexpr int LOCATIONS = 4; - struct alignas(float) Detection { - //center_x center_y w h - float bbox[LOCATIONS]; - float conf; // bbox_conf * cls_conf - float class_id; - float mask[32]; - }; -} + int getNbOutputs() const TRT_NOEXCEPT override { return 1; } -namespace nvinfer1 -{ - class API YoloLayerPlugin : public IPluginV2IOExt - { - public: - YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, bool is_segmentation, const std::vector& vYoloKernel); - YoloLayerPlugin(const void* data, size_t length); - ~YoloLayerPlugin(); + Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRT_NOEXCEPT override; - int getNbOutputs() const TRT_NOEXCEPT override - { - return 1; - } + int initialize() TRT_NOEXCEPT override; - Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRT_NOEXCEPT override; + virtual void terminate() TRT_NOEXCEPT override {}; - int initialize() TRT_NOEXCEPT override; + virtual size_t getWorkspaceSize(int maxBatchSize) const TRT_NOEXCEPT override { return 0; } - virtual void terminate() TRT_NOEXCEPT override {}; + virtual int enqueue(int batchSize, const void* const* inputs, void*TRT_CONST_ENQUEUE* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT override; - virtual size_t getWorkspaceSize(int maxBatchSize) const TRT_NOEXCEPT override { return 0; } + virtual size_t getSerializationSize() const TRT_NOEXCEPT override; - virtual int enqueue(int batchSize, const void* const* inputs, void*TRT_CONST_ENQUEUE* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT override; + virtual void serialize(void* buffer) const TRT_NOEXCEPT override; - virtual size_t getSerializationSize() const TRT_NOEXCEPT override; + bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const TRT_NOEXCEPT override { + return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT; + } - virtual void serialize(void* buffer) const TRT_NOEXCEPT override; + const char* getPluginType() const TRT_NOEXCEPT override; - bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const TRT_NOEXCEPT override { - return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT; - } + const char* getPluginVersion() const TRT_NOEXCEPT override; - const char* getPluginType() const TRT_NOEXCEPT override; + void destroy() TRT_NOEXCEPT override; - const char* getPluginVersion() const TRT_NOEXCEPT override; + IPluginV2IOExt* clone() const TRT_NOEXCEPT override; - void destroy() TRT_NOEXCEPT override; + void setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT override; - IPluginV2IOExt* clone() const TRT_NOEXCEPT override; + const char* getPluginNamespace() const TRT_NOEXCEPT override; - void setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT override; + DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT override; - const char* getPluginNamespace() const TRT_NOEXCEPT override; + bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT override; - DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT override; + bool canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT override; - bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT override; + void attachToContext( + cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT override; - bool canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT override; + void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) TRT_NOEXCEPT override; - void attachToContext( - cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT override; + void detachFromContext() TRT_NOEXCEPT override; - void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) TRT_NOEXCEPT override; - - void detachFromContext() TRT_NOEXCEPT override; - - private: - void forwardGpu(const float* const* inputs, float *output, cudaStream_t stream, int batchSize = 1); - int mThreadCount = 256; - const char* mPluginNamespace; - int mKernelCount; - int mClassCount; - int mYoloV5NetWidth; - int mYoloV5NetHeight; - int mMaxOutObject; - bool is_segmentation_; - std::vector mYoloKernel; - void** mAnchor; - }; - - class API YoloPluginCreator : public IPluginCreator - { - public: - YoloPluginCreator(); - - ~YoloPluginCreator() override = default; - - const char* getPluginName() const TRT_NOEXCEPT override; - - const char* getPluginVersion() const TRT_NOEXCEPT override; - - const PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override; - - IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) TRT_NOEXCEPT override; - - IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRT_NOEXCEPT override; - - void setPluginNamespace(const char* libNamespace) TRT_NOEXCEPT override - { - mNamespace = libNamespace; - } - - const char* getPluginNamespace() const TRT_NOEXCEPT override - { - return mNamespace.c_str(); - } - - private: - std::string mNamespace; - static PluginFieldCollection mFC; - static std::vector mPluginAttributes; - }; - REGISTER_TENSORRT_PLUGIN(YoloPluginCreator); + private: + void forwardGpu(const float* const* inputs, float *output, cudaStream_t stream, int batchSize = 1); + int mThreadCount = 256; + const char* mPluginNamespace; + int mKernelCount; + int mClassCount; + int mYoloV5NetWidth; + int mYoloV5NetHeight; + int mMaxOutObject; + bool is_segmentation_; + std::vector mYoloKernel; + void** mAnchor; +}; + +class API YoloPluginCreator : public IPluginCreator { + public: + YoloPluginCreator(); + + ~YoloPluginCreator() override = default; + + const char* getPluginName() const TRT_NOEXCEPT override; + + const char* getPluginVersion() const TRT_NOEXCEPT override; + + const PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override; + + IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) TRT_NOEXCEPT override; + + IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRT_NOEXCEPT override; + + void setPluginNamespace(const char* libNamespace) TRT_NOEXCEPT override { + mNamespace = libNamespace; + } + + const char* getPluginNamespace() const TRT_NOEXCEPT override { + return mNamespace.c_str(); + } + + private: + std::string mNamespace; + static PluginFieldCollection mFC; + static std::vector mPluginAttributes; +}; +REGISTER_TENSORRT_PLUGIN(YoloPluginCreator); }; -#endif // _YOLO_LAYER_H diff --git a/yolov5/samples b/yolov5/samples deleted file mode 120000 index 3a29e7c..0000000 --- a/yolov5/samples +++ /dev/null @@ -1 +0,0 @@ -../yolov3-spp/samples/ \ No newline at end of file diff --git a/yolov5/src/calibrator.cpp b/yolov5/src/calibrator.cpp index 472582f..ed7ce19 100644 --- a/yolov5/src/calibrator.cpp +++ b/yolov5/src/calibrator.cpp @@ -1,11 +1,35 @@ -#include -#include -#include -#include #include "calibrator.h" #include "cuda_utils.h" #include "utils.h" +#include +#include +#include +#include +#include + +static cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h) { + int w, h, x, y; + float r_w = input_w / (img.cols * 1.0); + float r_h = input_h / (img.rows * 1.0); + if (r_h > r_w) { + w = input_w; + h = r_w * img.rows; + x = 0; + y = (input_h - h) / 2; + } else { + w = r_h * img.cols; + h = input_h; + x = (input_w - w) / 2; + y = 0; + } + cv::Mat re(h, w, CV_8UC3); + cv::resize(img, re, re.size(), 0, 0, cv::INTER_LINEAR); + cv::Mat out(input_h, input_w, CV_8UC3, cv::Scalar(128, 128, 128)); + re.copyTo(out(cv::Rect(x, y, re.cols, re.rows))); + return out; +} + Int8EntropyCalibrator2::Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache) : batchsize_(batchsize), input_w_(input_w), @@ -15,59 +39,59 @@ Int8EntropyCalibrator2::Int8EntropyCalibrator2(int batchsize, int input_w, int i calib_table_name_(calib_table_name), input_blob_name_(input_blob_name), read_cache_(read_cache) { - input_count_ = 3 * input_w * input_h * batchsize; - CUDA_CHECK(cudaMalloc(&device_input_, input_count_ * sizeof(float))); - read_files_in_dir(img_dir, img_files_); + input_count_ = 3 * input_w * input_h * batchsize; + CUDA_CHECK(cudaMalloc(&device_input_, input_count_ * sizeof(float))); + read_files_in_dir(img_dir, img_files_); } Int8EntropyCalibrator2::~Int8EntropyCalibrator2() { - CUDA_CHECK(cudaFree(device_input_)); + CUDA_CHECK(cudaFree(device_input_)); } int Int8EntropyCalibrator2::getBatchSize() const TRT_NOEXCEPT { - return batchsize_; + return batchsize_; } bool Int8EntropyCalibrator2::getBatch(void* bindings[], const char* names[], int nbBindings) TRT_NOEXCEPT { - if (img_idx_ + batchsize_ > (int)img_files_.size()) { - return false; - } + if (img_idx_ + batchsize_ > (int)img_files_.size()) { + return false; + } - std::vector input_imgs_; - for (int i = img_idx_; i < img_idx_ + batchsize_; i++) { - std::cout << img_files_[i] << " " << i << std::endl; - cv::Mat temp = cv::imread(img_dir_ + img_files_[i]); - if (temp.empty()) { - std::cerr << "Fatal error: image cannot open!" << std::endl; - return false; - } - cv::Mat pr_img = preprocess_img(temp, input_w_, input_h_); - input_imgs_.push_back(pr_img); + std::vector input_imgs_; + for (int i = img_idx_; i < img_idx_ + batchsize_; i++) { + std::cout << img_files_[i] << " " << i << std::endl; + cv::Mat temp = cv::imread(img_dir_ + img_files_[i]); + if (temp.empty()) { + std::cerr << "Fatal error: image cannot open!" << std::endl; + return false; } - img_idx_ += batchsize_; - cv::Mat blob = cv::dnn::blobFromImages(input_imgs_, 1.0 / 255.0, cv::Size(input_w_, input_h_), cv::Scalar(0, 0, 0), true, false); + cv::Mat pr_img = preprocess_img(temp, input_w_, input_h_); + input_imgs_.push_back(pr_img); + } + img_idx_ += batchsize_; + cv::Mat blob = cv::dnn::blobFromImages(input_imgs_, 1.0 / 255.0, cv::Size(input_w_, input_h_), cv::Scalar(0, 0, 0), true, false); - CUDA_CHECK(cudaMemcpy(device_input_, blob.ptr(0), input_count_ * sizeof(float), cudaMemcpyHostToDevice)); - assert(!strcmp(names[0], input_blob_name_)); - bindings[0] = device_input_; - return true; + CUDA_CHECK(cudaMemcpy(device_input_, blob.ptr(0), input_count_ * sizeof(float), cudaMemcpyHostToDevice)); + assert(!strcmp(names[0], input_blob_name_)); + bindings[0] = device_input_; + return true; } const void* Int8EntropyCalibrator2::readCalibrationCache(size_t& length) TRT_NOEXCEPT { - std::cout << "reading calib cache: " << calib_table_name_ << std::endl; - calib_cache_.clear(); - std::ifstream input(calib_table_name_, std::ios::binary); - input >> std::noskipws; - if (read_cache_ && input.good()) { - std::copy(std::istream_iterator(input), std::istream_iterator(), std::back_inserter(calib_cache_)); - } - length = calib_cache_.size(); - return length ? calib_cache_.data() : nullptr; + std::cout << "reading calib cache: " << calib_table_name_ << std::endl; + calib_cache_.clear(); + std::ifstream input(calib_table_name_, std::ios::binary); + input >> std::noskipws; + if (read_cache_ && input.good()) { + std::copy(std::istream_iterator(input), std::istream_iterator(), std::back_inserter(calib_cache_)); + } + length = calib_cache_.size(); + return length ? calib_cache_.data() : nullptr; } void Int8EntropyCalibrator2::writeCalibrationCache(const void* cache, size_t length) TRT_NOEXCEPT { - std::cout << "writing calib cache: " << calib_table_name_ << " size: " << length << std::endl; - std::ofstream output(calib_table_name_, std::ios::binary); - output.write(reinterpret_cast(cache), length); + std::cout << "writing calib cache: " << calib_table_name_ << " size: " << length << std::endl; + std::ofstream output(calib_table_name_, std::ios::binary); + output.write(reinterpret_cast(cache), length); } diff --git a/yolov5/src/calibrator.h b/yolov5/src/calibrator.h index afba664..ed77b5f 100644 --- a/yolov5/src/calibrator.h +++ b/yolov5/src/calibrator.h @@ -1,10 +1,8 @@ -#ifndef ENTROPY_CALIBRATOR_H -#define ENTROPY_CALIBRATOR_H +#pragma once -#include +#include "macros.h" #include #include -#include "macros.h" //! \class Int8EntropyCalibrator2 //! @@ -12,28 +10,27 @@ //! CalibrationAlgoType is kENTROPY_CALIBRATION_2. //! class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2 { -public: - Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache = true); + public: + Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache = true); - virtual ~Int8EntropyCalibrator2(); - int getBatchSize() const TRT_NOEXCEPT override; - bool getBatch(void* bindings[], const char* names[], int nbBindings) TRT_NOEXCEPT override; - const void* readCalibrationCache(size_t& length) TRT_NOEXCEPT override; - void writeCalibrationCache(const void* cache, size_t length) TRT_NOEXCEPT override; + virtual ~Int8EntropyCalibrator2(); + int getBatchSize() const TRT_NOEXCEPT override; + bool getBatch(void* bindings[], const char* names[], int nbBindings) TRT_NOEXCEPT override; + const void* readCalibrationCache(size_t& length) TRT_NOEXCEPT override; + void writeCalibrationCache(const void* cache, size_t length) TRT_NOEXCEPT override; -private: - int batchsize_; - int input_w_; - int input_h_; - int img_idx_; - std::string img_dir_; - std::vector img_files_; - size_t input_count_; - std::string calib_table_name_; - const char* input_blob_name_; - bool read_cache_; - void* device_input_; - std::vector calib_cache_; + private: + int batchsize_; + int input_w_; + int input_h_; + int img_idx_; + std::string img_dir_; + std::vector img_files_; + size_t input_count_; + std::string calib_table_name_; + const char* input_blob_name_; + bool read_cache_; + void* device_input_; + std::vector calib_cache_; }; -#endif // ENTROPY_CALIBRATOR_H diff --git a/yolov5/src/common.hpp b/yolov5/src/common.hpp deleted file mode 100644 index e7458de..0000000 --- a/yolov5/src/common.hpp +++ /dev/null @@ -1,344 +0,0 @@ -#ifndef YOLOV5_COMMON_H_ -#define YOLOV5_COMMON_H_ - -#include -#include -#include -#include -#include -#include "NvInfer.h" -#include "yololayer.h" - -using namespace nvinfer1; - -cv::Rect get_rect(cv::Mat& img, float bbox[4]) { - float l, r, t, b; - float r_w = Yolo::INPUT_W / (img.cols * 1.0); - float r_h = Yolo::INPUT_H / (img.rows * 1.0); - if (r_h > r_w) { - l = bbox[0] - bbox[2] / 2.f; - r = bbox[0] + bbox[2] / 2.f; - t = bbox[1] - bbox[3] / 2.f - (Yolo::INPUT_H - r_w * img.rows) / 2; - b = bbox[1] + bbox[3] / 2.f - (Yolo::INPUT_H - r_w * img.rows) / 2; - l = l / r_w; - r = r / r_w; - t = t / r_w; - b = b / r_w; - } else { - l = bbox[0] - bbox[2] / 2.f - (Yolo::INPUT_W - r_h * img.cols) / 2; - r = bbox[0] + bbox[2] / 2.f - (Yolo::INPUT_W - r_h * img.cols) / 2; - t = bbox[1] - bbox[3] / 2.f; - b = bbox[1] + bbox[3] / 2.f; - l = l / r_h; - r = r / r_h; - t = t / r_h; - b = b / r_h; - } - return cv::Rect(round(l), round(t), round(r - l), round(b - t)); -} - -float iou(float lbox[4], float rbox[4]) { - float interBox[] = { - (std::max)(lbox[0] - lbox[2] / 2.f , rbox[0] - rbox[2] / 2.f), //left - (std::min)(lbox[0] + lbox[2] / 2.f , rbox[0] + rbox[2] / 2.f), //right - (std::max)(lbox[1] - lbox[3] / 2.f , rbox[1] - rbox[3] / 2.f), //top - (std::min)(lbox[1] + lbox[3] / 2.f , rbox[1] + rbox[3] / 2.f), //bottom - }; - - if (interBox[2] > interBox[3] || interBox[0] > interBox[1]) - return 0.0f; - - float interBoxS = (interBox[1] - interBox[0])*(interBox[3] - interBox[2]); - return interBoxS / (lbox[2] * lbox[3] + rbox[2] * rbox[3] - interBoxS); -} - -bool cmp(const Yolo::Detection& a, const Yolo::Detection& b) { - return a.conf > b.conf; -} - -void nms(std::vector& res, float *output, float conf_thresh, float nms_thresh = 0.5) { - int det_size = sizeof(Yolo::Detection) / sizeof(float); - std::map> m; - for (int i = 0; i < output[0] && i < Yolo::MAX_OUTPUT_BBOX_COUNT; i++) { - if (output[1 + det_size * i + 4] <= conf_thresh) continue; - Yolo::Detection det; - memcpy(&det, &output[1 + det_size * i], det_size * sizeof(float)); - if (m.count(det.class_id) == 0) m.emplace(det.class_id, std::vector()); - m[det.class_id].push_back(det); - } - for (auto it = m.begin(); it != m.end(); it++) { - //std::cout << it->second[0].class_id << " --- " << std::endl; - auto& dets = it->second; - std::sort(dets.begin(), dets.end(), cmp); - for (size_t m = 0; m < dets.size(); ++m) { - auto& item = dets[m]; - res.push_back(item); - for (size_t n = m + 1; n < dets.size(); ++n) { - if (iou(item.bbox, dets[n].bbox) > nms_thresh) { - dets.erase(dets.begin() + n); - --n; - } - } - } - } -} - -// TensorRT weight files have a simple space delimited format: -// [type] [size] -std::map loadWeights(const std::string file) { - std::cout << "Loading weights: " << file << std::endl; - std::map weightMap; - - // Open weights file - std::ifstream input(file); - assert(input.is_open() && "Unable to load weight file. please check if the .wts file path is right!!!!!!"); - - // Read number of weight blobs - int32_t count; - input >> count; - assert(count > 0 && "Invalid weight map file."); - - while (count--) { - Weights wt{ DataType::kFLOAT, nullptr, 0 }; - uint32_t size; - - // Read name and type of blob - std::string name; - input >> name >> std::dec >> size; - wt.type = DataType::kFLOAT; - - // Load blob - uint32_t* val = reinterpret_cast(malloc(sizeof(val) * size)); - for (uint32_t x = 0, y = size; x < y; ++x) { - input >> std::hex >> val[x]; - } - wt.values = val; - - wt.count = size; - weightMap[name] = wt; - } - - return weightMap; -} - -IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map& weightMap, ITensor& input, std::string lname, float eps) { - float *gamma = (float*)weightMap[lname + ".weight"].values; - float *beta = (float*)weightMap[lname + ".bias"].values; - float *mean = (float*)weightMap[lname + ".running_mean"].values; - float *var = (float*)weightMap[lname + ".running_var"].values; - int len = weightMap[lname + ".running_var"].count; - - float *scval = reinterpret_cast(malloc(sizeof(float) * len)); - for (int i = 0; i < len; i++) { - scval[i] = gamma[i] / sqrt(var[i] + eps); - } - Weights scale{ DataType::kFLOAT, scval, len }; - - float *shval = reinterpret_cast(malloc(sizeof(float) * len)); - for (int i = 0; i < len; i++) { - shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps); - } - Weights shift{ DataType::kFLOAT, shval, len }; - - float *pval = reinterpret_cast(malloc(sizeof(float) * len)); - for (int i = 0; i < len; i++) { - pval[i] = 1.0; - } - Weights power{ DataType::kFLOAT, pval, len }; - - weightMap[lname + ".scale"] = scale; - weightMap[lname + ".shift"] = shift; - weightMap[lname + ".power"] = power; - IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power); - assert(scale_1); - return scale_1; -} - -ILayer* convBlock(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, int ksize, int s, int g, std::string lname) { - Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; - int p = ksize / 3; - IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[lname + ".conv.weight"], emptywts); - assert(conv1); - conv1->setStrideNd(DimsHW{ s, s }); - conv1->setPaddingNd(DimsHW{ p, p }); - conv1->setNbGroups(g); - conv1->setName((lname + ".conv").c_str()); - IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn", 1e-3); - - // silu = x * sigmoid - auto sig = network->addActivation(*bn1->getOutput(0), ActivationType::kSIGMOID); - assert(sig); - auto ew = network->addElementWise(*bn1->getOutput(0), *sig->getOutput(0), ElementWiseOperation::kPROD); - assert(ew); - return ew; -} - -ILayer* focus(INetworkDefinition *network, std::map& weightMap, ITensor& input, int inch, int outch, int ksize, std::string lname) { - ISliceLayer *s1 = network->addSlice(input, Dims3{ 0, 0, 0 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 }); - ISliceLayer *s2 = network->addSlice(input, Dims3{ 0, 1, 0 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 }); - ISliceLayer *s3 = network->addSlice(input, Dims3{ 0, 0, 1 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 }); - ISliceLayer *s4 = network->addSlice(input, Dims3{ 0, 1, 1 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 }); - ITensor* inputTensors[] = { s1->getOutput(0), s2->getOutput(0), s3->getOutput(0), s4->getOutput(0) }; - auto cat = network->addConcatenation(inputTensors, 4); - auto conv = convBlock(network, weightMap, *cat->getOutput(0), outch, ksize, 1, 1, lname + ".conv"); - return conv; -} - -ILayer* bottleneck(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, bool shortcut, int g, float e, std::string lname) { - auto cv1 = convBlock(network, weightMap, input, (int)((float)c2 * e), 1, 1, 1, lname + ".cv1"); - auto cv2 = convBlock(network, weightMap, *cv1->getOutput(0), c2, 3, 1, g, lname + ".cv2"); - if (shortcut && c1 == c2) { - auto ew = network->addElementWise(input, *cv2->getOutput(0), ElementWiseOperation::kSUM); - return ew; - } - return cv2; -} - -ILayer* bottleneckCSP(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int n, bool shortcut, int g, float e, std::string lname) { - Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; - int c_ = (int)((float)c2 * e); - auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1"); - auto cv2 = network->addConvolutionNd(input, c_, DimsHW{ 1, 1 }, weightMap[lname + ".cv2.weight"], emptywts); - ITensor *y1 = cv1->getOutput(0); - for (int i = 0; i < n; i++) { - auto b = bottleneck(network, weightMap, *y1, c_, c_, shortcut, g, 1.0, lname + ".m." + std::to_string(i)); - y1 = b->getOutput(0); - } - auto cv3 = network->addConvolutionNd(*y1, c_, DimsHW{ 1, 1 }, weightMap[lname + ".cv3.weight"], emptywts); - - ITensor* inputTensors[] = { cv3->getOutput(0), cv2->getOutput(0) }; - auto cat = network->addConcatenation(inputTensors, 2); - - IScaleLayer* bn = addBatchNorm2d(network, weightMap, *cat->getOutput(0), lname + ".bn", 1e-4); - auto lr = network->addActivation(*bn->getOutput(0), ActivationType::kLEAKY_RELU); - lr->setAlpha(0.1); - - auto cv4 = convBlock(network, weightMap, *lr->getOutput(0), c2, 1, 1, 1, lname + ".cv4"); - return cv4; -} - -ILayer* C3(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int n, bool shortcut, int g, float e, std::string lname) { - int c_ = (int)((float)c2 * e); - auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1"); - auto cv2 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv2"); - ITensor *y1 = cv1->getOutput(0); - for (int i = 0; i < n; i++) { - auto b = bottleneck(network, weightMap, *y1, c_, c_, shortcut, g, 1.0, lname + ".m." + std::to_string(i)); - y1 = b->getOutput(0); - } - - ITensor* inputTensors[] = { y1, cv2->getOutput(0) }; - auto cat = network->addConcatenation(inputTensors, 2); - - auto cv3 = convBlock(network, weightMap, *cat->getOutput(0), c2, 1, 1, 1, lname + ".cv3"); - return cv3; -} - -ILayer* SPP(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int k1, int k2, int k3, std::string lname) { - int c_ = c1 / 2; - auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1"); - - auto pool1 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k1, k1 }); - pool1->setPaddingNd(DimsHW{ k1 / 2, k1 / 2 }); - pool1->setStrideNd(DimsHW{ 1, 1 }); - auto pool2 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k2, k2 }); - pool2->setPaddingNd(DimsHW{ k2 / 2, k2 / 2 }); - pool2->setStrideNd(DimsHW{ 1, 1 }); - auto pool3 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k3, k3 }); - pool3->setPaddingNd(DimsHW{ k3 / 2, k3 / 2 }); - pool3->setStrideNd(DimsHW{ 1, 1 }); - - ITensor* inputTensors[] = { cv1->getOutput(0), pool1->getOutput(0), pool2->getOutput(0), pool3->getOutput(0) }; - auto cat = network->addConcatenation(inputTensors, 4); - - auto cv2 = convBlock(network, weightMap, *cat->getOutput(0), c2, 1, 1, 1, lname + ".cv2"); - return cv2; -} - -ILayer* SPPF(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int k, std::string lname) { - int c_ = c1 / 2; - auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1"); - - auto pool1 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k, k }); - pool1->setPaddingNd(DimsHW{ k / 2, k / 2 }); - pool1->setStrideNd(DimsHW{ 1, 1 }); - auto pool2 = network->addPoolingNd(*pool1->getOutput(0), PoolingType::kMAX, DimsHW{ k, k }); - pool2->setPaddingNd(DimsHW{ k / 2, k / 2 }); - pool2->setStrideNd(DimsHW{ 1, 1 }); - auto pool3 = network->addPoolingNd(*pool2->getOutput(0), PoolingType::kMAX, DimsHW{ k, k }); - pool3->setPaddingNd(DimsHW{ k / 2, k / 2 }); - pool3->setStrideNd(DimsHW{ 1, 1 }); - ITensor* inputTensors[] = { cv1->getOutput(0), pool1->getOutput(0), pool2->getOutput(0), pool3->getOutput(0) }; - auto cat = network->addConcatenation(inputTensors, 4); - auto cv2 = convBlock(network, weightMap, *cat->getOutput(0), c2, 1, 1, 1, lname + ".cv2"); - return cv2; -} - -ILayer* Proto(INetworkDefinition* network, std::map& weightMap, ITensor& input, int c_, int c2, std::string lname) { - auto cv1 = convBlock(network, weightMap, input, c_, 3, 1, 1, lname + ".cv1"); - - auto upsample = network->addResize(*cv1->getOutput(0)); - assert(upsample); - upsample->setResizeMode(ResizeMode::kNEAREST); - const float scales[] = {1, 2, 2}; - upsample->setScales(scales, 3); - - auto cv2 = convBlock(network, weightMap, *upsample->getOutput(0), c_, 3, 1, 1, lname + ".cv2"); - auto cv3 = convBlock(network, weightMap, *cv2->getOutput(0), c2, 1, 1, 1, lname + ".cv3"); - assert(cv3); - return cv3; -} - -std::vector> getAnchors(std::map& weightMap, std::string lname) { - std::vector> anchors; - Weights wts = weightMap[lname + ".anchor_grid"]; - int anchor_len = Yolo::CHECK_COUNT * 2; - for (int i = 0; i < wts.count / anchor_len; i++) { - auto *p = (const float*)wts.values + i * anchor_len; - std::vector anchor(p, p + anchor_len); - anchors.push_back(anchor); - } - return anchors; -} - -IPluginV2Layer* addYoLoLayer(INetworkDefinition *network, std::map& weightMap, std::string lname, std::vector dets, bool is_segmentation = false) { - auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1"); - auto anchors = getAnchors(weightMap, lname); - PluginField plugin_fields[2]; - int netinfo[5] = {Yolo::CLASS_NUM, Yolo::INPUT_W, Yolo::INPUT_H, Yolo::MAX_OUTPUT_BBOX_COUNT, (int)is_segmentation}; - plugin_fields[0].data = netinfo; - plugin_fields[0].length = 5; - plugin_fields[0].name = "netinfo"; - plugin_fields[0].type = PluginFieldType::kFLOAT32; - - //load strides from Detect layer - assert(weightMap.find(lname + ".strides") != weightMap.end() && "Not found `strides`, please check gen_wts.py!!!"); - Weights strides = weightMap[lname + ".strides"]; - auto *p = (const float*)(strides.values); - std::vector scales(p, p + strides.count); - - std::vector kernels; - for (size_t i = 0; i < anchors.size(); i++) { - Yolo::YoloKernel kernel; - kernel.width = Yolo::INPUT_W / scales[i]; - kernel.height = Yolo::INPUT_H / scales[i]; - memcpy(kernel.anchors, &anchors[i][0], anchors[i].size() * sizeof(float)); - kernels.push_back(kernel); - } - plugin_fields[1].data = &kernels[0]; - plugin_fields[1].length = kernels.size(); - plugin_fields[1].name = "kernels"; - plugin_fields[1].type = PluginFieldType::kFLOAT32; - PluginFieldCollection plugin_data; - plugin_data.nbFields = 2; - plugin_data.fields = plugin_fields; - IPluginV2 *plugin_obj = creator->createPlugin("yololayer", &plugin_data); - std::vector input_tensors; - for (auto det: dets) { - input_tensors.push_back(det->getOutput(0)); - } - auto yolo = network->addPluginV2(&input_tensors[0], input_tensors.size(), *plugin_obj); - return yolo; -} -#endif // YOLOV5_COMMON_H_ - diff --git a/yolov5/src/config.h b/yolov5/src/config.h new file mode 100644 index 0000000..abf5cec --- /dev/null +++ b/yolov5/src/config.h @@ -0,0 +1,55 @@ +#pragma once + +/* -------------------------------------------------------- + * These configs are related to tensorrt model, if these are changed, + * please re-compile and re-serialize the tensorrt model. + * --------------------------------------------------------*/ + +// For INT8, you need prepare the calibration dataset, please refer to +// https://github.com/wang-xinyu/tensorrtx/tree/master/yolov5#int8-quantization +#define USE_FP16 // set USE_INT8 or USE_FP16 or USE_FP32 + +// These are used to define input/output tensor names, +// you can set them to whatever you want. +const static char* kInputTensorName = "data"; +const static char* kOutputTensorName = "prob"; + +// Detection model and Segmentation model' number of classes +constexpr static int kNumClass = 80; + +// Classfication model's number of classes +constexpr static int kClsNumClass = 1000; + +constexpr static int kBatchSize = 1; + +// Yolo's input width and height must by divisible by 32 +constexpr static int kInputH = 640; +constexpr static int kInputW = 640; + +// Classfication model's input shape +constexpr static int kClsInputH = 224; +constexpr static int kClsInputW = 224; + +// Maximum number of output bounding boxes from yololayer plugin. +// That is maximum number of output bounding boxes before NMS. +constexpr static int kMaxNumOutputBbox = 1000; + +constexpr static int kNumAnchor = 3; + +// The bboxes whose confidence is lower than kIgnoreThresh will be ignored in yololayer plugin. +constexpr static float kIgnoreThresh = 0.1f; + +/* -------------------------------------------------------- + * These configs are NOT related to tensorrt model, if these are changed, + * please re-compile, but no need to re-serialize the tensorrt model. + * --------------------------------------------------------*/ + +// NMS overlapping thresh and final detection confidence thresh +const static float kNmsThresh = 0.45f; +const static float kConfThresh = 0.5f; + +const static int kGpuId = 0; + +// If your image size is larger than 4096 * 3112, please increase this value +const static int kMaxInputImageSize = 4096 * 3112; + diff --git a/yolov5/src/macros.h b/yolov5/src/macros.h index 0e5b71b..17339a2 100644 --- a/yolov5/src/macros.h +++ b/yolov5/src/macros.h @@ -1,6 +1,8 @@ #ifndef __MACROS_H #define __MACROS_H +#include + #ifdef API_EXPORTS #if defined(_MSC_VER) #define API __declspec(dllexport) diff --git a/yolov5/src/model.cpp b/yolov5/src/model.cpp new file mode 100644 index 0000000..712d782 --- /dev/null +++ b/yolov5/src/model.cpp @@ -0,0 +1,628 @@ +#include "model.h" +#include "calibrator.h" +#include "config.h" +#include "yololayer.h" + +#include +#include +#include +#include +#include +#include + +using namespace nvinfer1; + +// TensorRT weight files have a simple space delimited format: +// [type] [size] +static std::map loadWeights(const std::string file) { + std::cout << "Loading weights: " << file << std::endl; + std::map weightMap; + + // Open weights file + std::ifstream input(file); + assert(input.is_open() && "Unable to load weight file. please check if the .wts file path is right!!!!!!"); + + // Read number of weight blobs + int32_t count; + input >> count; + assert(count > 0 && "Invalid weight map file."); + + while (count--) { + Weights wt{ DataType::kFLOAT, nullptr, 0 }; + uint32_t size; + + // Read name and type of blob + std::string name; + input >> name >> std::dec >> size; + wt.type = DataType::kFLOAT; + + // Load blob + uint32_t* val = reinterpret_cast(malloc(sizeof(val) * size)); + for (uint32_t x = 0, y = size; x < y; ++x) { + input >> std::hex >> val[x]; + } + wt.values = val; + + wt.count = size; + weightMap[name] = wt; + } + + return weightMap; +} + +static int get_width(int x, float gw, int divisor = 8) { + return int(ceil((x * gw) / divisor)) * divisor; +} + +static int get_depth(int x, float gd) { + if (x == 1) return 1; + int r = round(x * gd); + if (x * gd - int(x * gd) == 0.5 && (int(x * gd) % 2) == 0) { + --r; + } + return std::max(r, 1); +} + +static IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map& weightMap, ITensor& input, std::string lname, float eps) { + float* gamma = (float*)weightMap[lname + ".weight"].values; + float* beta = (float*)weightMap[lname + ".bias"].values; + float* mean = (float*)weightMap[lname + ".running_mean"].values; + float* var = (float*)weightMap[lname + ".running_var"].values; + int len = weightMap[lname + ".running_var"].count; + + float* scval = reinterpret_cast(malloc(sizeof(float) * len)); + for (int i = 0; i < len; i++) { + scval[i] = gamma[i] / sqrt(var[i] + eps); + } + Weights scale{ DataType::kFLOAT, scval, len }; + + float* shval = reinterpret_cast(malloc(sizeof(float) * len)); + for (int i = 0; i < len; i++) { + shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps); + } + Weights shift{ DataType::kFLOAT, shval, len }; + + float* pval = reinterpret_cast(malloc(sizeof(float) * len)); + for (int i = 0; i < len; i++) { + pval[i] = 1.0; + } + Weights power{ DataType::kFLOAT, pval, len }; + + weightMap[lname + ".scale"] = scale; + weightMap[lname + ".shift"] = shift; + weightMap[lname + ".power"] = power; + IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power); + assert(scale_1); + return scale_1; +} + +static ILayer* convBlock(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, int ksize, int s, int g, std::string lname) { + Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; + int p = ksize / 3; + IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[lname + ".conv.weight"], emptywts); + assert(conv1); + conv1->setStrideNd(DimsHW{ s, s }); + conv1->setPaddingNd(DimsHW{ p, p }); + conv1->setNbGroups(g); + conv1->setName((lname + ".conv").c_str()); + IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn", 1e-3); + + // silu = x * sigmoid + auto sig = network->addActivation(*bn1->getOutput(0), ActivationType::kSIGMOID); + assert(sig); + auto ew = network->addElementWise(*bn1->getOutput(0), *sig->getOutput(0), ElementWiseOperation::kPROD); + assert(ew); + return ew; +} + +static ILayer* focus(INetworkDefinition *network, std::map& weightMap, ITensor& input, int inch, int outch, int ksize, std::string lname) { + ISliceLayer* s1 = network->addSlice(input, Dims3{ 0, 0, 0 }, Dims3{ inch, kInputH / 2, kInputW / 2 }, Dims3{ 1, 2, 2 }); + ISliceLayer* s2 = network->addSlice(input, Dims3{ 0, 1, 0 }, Dims3{ inch, kInputH / 2, kInputW / 2 }, Dims3{ 1, 2, 2 }); + ISliceLayer* s3 = network->addSlice(input, Dims3{ 0, 0, 1 }, Dims3{ inch, kInputH / 2, kInputW / 2 }, Dims3{ 1, 2, 2 }); + ISliceLayer* s4 = network->addSlice(input, Dims3{ 0, 1, 1 }, Dims3{ inch, kInputH / 2, kInputW / 2 }, Dims3{ 1, 2, 2 }); + ITensor* inputTensors[] = { s1->getOutput(0), s2->getOutput(0), s3->getOutput(0), s4->getOutput(0) }; + auto cat = network->addConcatenation(inputTensors, 4); + auto conv = convBlock(network, weightMap, *cat->getOutput(0), outch, ksize, 1, 1, lname + ".conv"); + return conv; +} + +static ILayer* bottleneck(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, bool shortcut, int g, float e, std::string lname) { + auto cv1 = convBlock(network, weightMap, input, (int)((float)c2 * e), 1, 1, 1, lname + ".cv1"); + auto cv2 = convBlock(network, weightMap, *cv1->getOutput(0), c2, 3, 1, g, lname + ".cv2"); + if (shortcut && c1 == c2) { + auto ew = network->addElementWise(input, *cv2->getOutput(0), ElementWiseOperation::kSUM); + return ew; + } + return cv2; +} + +static ILayer* bottleneckCSP(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int n, bool shortcut, int g, float e, std::string lname) { + Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; + int c_ = (int)((float)c2 * e); + auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1"); + auto cv2 = network->addConvolutionNd(input, c_, DimsHW{ 1, 1 }, weightMap[lname + ".cv2.weight"], emptywts); + ITensor* y1 = cv1->getOutput(0); + for (int i = 0; i < n; i++) { + auto b = bottleneck(network, weightMap, *y1, c_, c_, shortcut, g, 1.0, lname + ".m." + std::to_string(i)); + y1 = b->getOutput(0); + } + auto cv3 = network->addConvolutionNd(*y1, c_, DimsHW{ 1, 1 }, weightMap[lname + ".cv3.weight"], emptywts); + + ITensor* inputTensors[] = { cv3->getOutput(0), cv2->getOutput(0) }; + auto cat = network->addConcatenation(inputTensors, 2); + + IScaleLayer* bn = addBatchNorm2d(network, weightMap, *cat->getOutput(0), lname + ".bn", 1e-4); + auto lr = network->addActivation(*bn->getOutput(0), ActivationType::kLEAKY_RELU); + lr->setAlpha(0.1); + + auto cv4 = convBlock(network, weightMap, *lr->getOutput(0), c2, 1, 1, 1, lname + ".cv4"); + return cv4; +} + +static ILayer* C3(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int n, bool shortcut, int g, float e, std::string lname) { + int c_ = (int)((float)c2 * e); + auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1"); + auto cv2 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv2"); + ITensor *y1 = cv1->getOutput(0); + for (int i = 0; i < n; i++) { + auto b = bottleneck(network, weightMap, *y1, c_, c_, shortcut, g, 1.0, lname + ".m." + std::to_string(i)); + y1 = b->getOutput(0); + } + + ITensor* inputTensors[] = { y1, cv2->getOutput(0) }; + auto cat = network->addConcatenation(inputTensors, 2); + + auto cv3 = convBlock(network, weightMap, *cat->getOutput(0), c2, 1, 1, 1, lname + ".cv3"); + return cv3; +} + +static ILayer* SPP(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int k1, int k2, int k3, std::string lname) { + int c_ = c1 / 2; + auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1"); + + auto pool1 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k1, k1 }); + pool1->setPaddingNd(DimsHW{ k1 / 2, k1 / 2 }); + pool1->setStrideNd(DimsHW{ 1, 1 }); + auto pool2 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k2, k2 }); + pool2->setPaddingNd(DimsHW{ k2 / 2, k2 / 2 }); + pool2->setStrideNd(DimsHW{ 1, 1 }); + auto pool3 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k3, k3 }); + pool3->setPaddingNd(DimsHW{ k3 / 2, k3 / 2 }); + pool3->setStrideNd(DimsHW{ 1, 1 }); + + ITensor* inputTensors[] = { cv1->getOutput(0), pool1->getOutput(0), pool2->getOutput(0), pool3->getOutput(0) }; + auto cat = network->addConcatenation(inputTensors, 4); + + auto cv2 = convBlock(network, weightMap, *cat->getOutput(0), c2, 1, 1, 1, lname + ".cv2"); + return cv2; +} + +static ILayer* SPPF(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int k, std::string lname) { + int c_ = c1 / 2; + auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1"); + + auto pool1 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k, k }); + pool1->setPaddingNd(DimsHW{ k / 2, k / 2 }); + pool1->setStrideNd(DimsHW{ 1, 1 }); + auto pool2 = network->addPoolingNd(*pool1->getOutput(0), PoolingType::kMAX, DimsHW{ k, k }); + pool2->setPaddingNd(DimsHW{ k / 2, k / 2 }); + pool2->setStrideNd(DimsHW{ 1, 1 }); + auto pool3 = network->addPoolingNd(*pool2->getOutput(0), PoolingType::kMAX, DimsHW{ k, k }); + pool3->setPaddingNd(DimsHW{ k / 2, k / 2 }); + pool3->setStrideNd(DimsHW{ 1, 1 }); + ITensor* inputTensors[] = { cv1->getOutput(0), pool1->getOutput(0), pool2->getOutput(0), pool3->getOutput(0) }; + auto cat = network->addConcatenation(inputTensors, 4); + auto cv2 = convBlock(network, weightMap, *cat->getOutput(0), c2, 1, 1, 1, lname + ".cv2"); + return cv2; +} + +static ILayer* Proto(INetworkDefinition* network, std::map& weightMap, ITensor& input, int c_, int c2, std::string lname) { + auto cv1 = convBlock(network, weightMap, input, c_, 3, 1, 1, lname + ".cv1"); + + auto upsample = network->addResize(*cv1->getOutput(0)); + assert(upsample); + upsample->setResizeMode(ResizeMode::kNEAREST); + const float scales[] = {1, 2, 2}; + upsample->setScales(scales, 3); + + auto cv2 = convBlock(network, weightMap, *upsample->getOutput(0), c_, 3, 1, 1, lname + ".cv2"); + auto cv3 = convBlock(network, weightMap, *cv2->getOutput(0), c2, 1, 1, 1, lname + ".cv3"); + assert(cv3); + return cv3; +} + +static std::vector> getAnchors(std::map& weightMap, std::string lname) { + std::vector> anchors; + Weights wts = weightMap[lname + ".anchor_grid"]; + int anchor_len = kNumAnchor * 2; + for (int i = 0; i < wts.count / anchor_len; i++) { + auto *p = (const float*)wts.values + i * anchor_len; + std::vector anchor(p, p + anchor_len); + anchors.push_back(anchor); + } + return anchors; +} + +static IPluginV2Layer* addYoLoLayer(INetworkDefinition *network, std::map& weightMap, std::string lname, std::vector dets, bool is_segmentation = false) { + auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1"); + auto anchors = getAnchors(weightMap, lname); + PluginField plugin_fields[2]; + int netinfo[5] = {kNumClass, kInputW, kInputH, kMaxNumOutputBbox, (int)is_segmentation}; + plugin_fields[0].data = netinfo; + plugin_fields[0].length = 5; + plugin_fields[0].name = "netinfo"; + plugin_fields[0].type = PluginFieldType::kFLOAT32; + + //load strides from Detect layer + assert(weightMap.find(lname + ".strides") != weightMap.end() && "Not found `strides`, please check gen_wts.py!!!"); + Weights strides = weightMap[lname + ".strides"]; + auto *p = (const float*)(strides.values); + std::vector scales(p, p + strides.count); + + std::vector kernels; + for (size_t i = 0; i < anchors.size(); i++) { + YoloKernel kernel; + kernel.width = kInputW / scales[i]; + kernel.height = kInputH / scales[i]; + memcpy(kernel.anchors, &anchors[i][0], anchors[i].size() * sizeof(float)); + kernels.push_back(kernel); + } + plugin_fields[1].data = &kernels[0]; + plugin_fields[1].length = kernels.size(); + plugin_fields[1].name = "kernels"; + plugin_fields[1].type = PluginFieldType::kFLOAT32; + PluginFieldCollection plugin_data; + plugin_data.nbFields = 2; + plugin_data.fields = plugin_fields; + IPluginV2 *plugin_obj = creator->createPlugin("yololayer", &plugin_data); + std::vector input_tensors; + for (auto det: dets) { + input_tensors.push_back(det->getOutput(0)); + } + auto yolo = network->addPluginV2(&input_tensors[0], input_tensors.size(), *plugin_obj); + return yolo; +} + +ICudaEngine* build_det_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) { + INetworkDefinition* network = builder->createNetworkV2(0U); + + // Create input tensor of shape {3, kInputH, kInputW} + ITensor* data = network->addInput(kInputTensorName, dt, Dims3{ 3, kInputH, kInputW }); + assert(data); + std::map weightMap = loadWeights(wts_name); + + // Backbone + auto conv0 = convBlock(network, weightMap, *data, get_width(64, gw), 6, 2, 1, "model.0"); + assert(conv0); + auto conv1 = convBlock(network, weightMap, *conv0->getOutput(0), get_width(128, gw), 3, 2, 1, "model.1"); + auto bottleneck_CSP2 = C3(network, weightMap, *conv1->getOutput(0), get_width(128, gw), get_width(128, gw), get_depth(3, gd), true, 1, 0.5, "model.2"); + auto conv3 = convBlock(network, weightMap, *bottleneck_CSP2->getOutput(0), get_width(256, gw), 3, 2, 1, "model.3"); + auto bottleneck_csp4 = C3(network, weightMap, *conv3->getOutput(0), get_width(256, gw), get_width(256, gw), get_depth(6, gd), true, 1, 0.5, "model.4"); + auto conv5 = convBlock(network, weightMap, *bottleneck_csp4->getOutput(0), get_width(512, gw), 3, 2, 1, "model.5"); + auto bottleneck_csp6 = C3(network, weightMap, *conv5->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(9, gd), true, 1, 0.5, "model.6"); + auto conv7 = convBlock(network, weightMap, *bottleneck_csp6->getOutput(0), get_width(1024, gw), 3, 2, 1, "model.7"); + auto bottleneck_csp8 = C3(network, weightMap, *conv7->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), true, 1, 0.5, "model.8"); + auto spp9 = SPPF(network, weightMap, *bottleneck_csp8->getOutput(0), get_width(1024, gw), get_width(1024, gw), 5, "model.9"); + + // Head + auto conv10 = convBlock(network, weightMap, *spp9->getOutput(0), get_width(512, gw), 1, 1, 1, "model.10"); + + auto upsample11 = network->addResize(*conv10->getOutput(0)); + assert(upsample11); + upsample11->setResizeMode(ResizeMode::kNEAREST); + upsample11->setOutputDimensions(bottleneck_csp6->getOutput(0)->getDimensions()); + + ITensor* inputTensors12[] = { upsample11->getOutput(0), bottleneck_csp6->getOutput(0) }; + auto cat12 = network->addConcatenation(inputTensors12, 2); + auto bottleneck_csp13 = C3(network, weightMap, *cat12->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.13"); + auto conv14 = convBlock(network, weightMap, *bottleneck_csp13->getOutput(0), get_width(256, gw), 1, 1, 1, "model.14"); + + auto upsample15 = network->addResize(*conv14->getOutput(0)); + assert(upsample15); + upsample15->setResizeMode(ResizeMode::kNEAREST); + upsample15->setOutputDimensions(bottleneck_csp4->getOutput(0)->getDimensions()); + + ITensor* inputTensors16[] = { upsample15->getOutput(0), bottleneck_csp4->getOutput(0) }; + auto cat16 = network->addConcatenation(inputTensors16, 2); + + auto bottleneck_csp17 = C3(network, weightMap, *cat16->getOutput(0), get_width(512, gw), get_width(256, gw), get_depth(3, gd), false, 1, 0.5, "model.17"); + + // Detect + IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (kNumClass + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]); + auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), get_width(256, gw), 3, 2, 1, "model.18"); + ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) }; + auto cat19 = network->addConcatenation(inputTensors19, 2); + auto bottleneck_csp20 = C3(network, weightMap, *cat19->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.20"); + IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (kNumClass + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]); + auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), get_width(512, gw), 3, 2, 1, "model.21"); + ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) }; + auto cat22 = network->addConcatenation(inputTensors22, 2); + auto bottleneck_csp23 = C3(network, weightMap, *cat22->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.23"); + IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (kNumClass + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]); + + auto yolo = addYoLoLayer(network, weightMap, "model.24", std::vector{det0, det1, det2}); + yolo->getOutput(0)->setName(kOutputTensorName); + network->markOutput(*yolo->getOutput(0)); + + // Engine config + builder->setMaxBatchSize(maxBatchSize); + config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB +#if defined(USE_FP16) + config->setFlag(BuilderFlag::kFP16); +#elif defined(USE_INT8) + std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; + assert(builder->platformHasFastInt8()); + config->setFlag(BuilderFlag::kINT8); + Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, kInputW, kInputH, "./coco_calib/", "int8calib.table", kInputTensorName); + config->setInt8Calibrator(calibrator); +#endif + + std::cout << "Building engine, please wait for a while..." << std::endl; + ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); + std::cout << "Build engine successfully!" << std::endl; + + // Don't need the network any more + network->destroy(); + + // Release host memory + for (auto& mem : weightMap) { + free((void*)(mem.second.values)); + } + + return engine; +} + +ICudaEngine* build_det_p6_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) { + INetworkDefinition* network = builder->createNetworkV2(0U); + + // Create input tensor of shape {3, kInputH, kInputW} + ITensor* data = network->addInput(kInputTensorName, dt, Dims3{ 3, kInputH, kInputW }); + assert(data); + + std::map weightMap = loadWeights(wts_name); + + // Backbone + auto conv0 = convBlock(network, weightMap, *data, get_width(64, gw), 6, 2, 1, "model.0"); + auto conv1 = convBlock(network, weightMap, *conv0->getOutput(0), get_width(128, gw), 3, 2, 1, "model.1"); + auto c3_2 = C3(network, weightMap, *conv1->getOutput(0), get_width(128, gw), get_width(128, gw), get_depth(3, gd), true, 1, 0.5, "model.2"); + auto conv3 = convBlock(network, weightMap, *c3_2->getOutput(0), get_width(256, gw), 3, 2, 1, "model.3"); + auto c3_4 = C3(network, weightMap, *conv3->getOutput(0), get_width(256, gw), get_width(256, gw), get_depth(6, gd), true, 1, 0.5, "model.4"); + auto conv5 = convBlock(network, weightMap, *c3_4->getOutput(0), get_width(512, gw), 3, 2, 1, "model.5"); + auto c3_6 = C3(network, weightMap, *conv5->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(9, gd), true, 1, 0.5, "model.6"); + auto conv7 = convBlock(network, weightMap, *c3_6->getOutput(0), get_width(768, gw), 3, 2, 1, "model.7"); + auto c3_8 = C3(network, weightMap, *conv7->getOutput(0), get_width(768, gw), get_width(768, gw), get_depth(3, gd), true, 1, 0.5, "model.8"); + auto conv9 = convBlock(network, weightMap, *c3_8->getOutput(0), get_width(1024, gw), 3, 2, 1, "model.9"); + auto c3_10 = C3(network, weightMap, *conv9->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), true, 1, 0.5, "model.10"); + auto sppf11 = SPPF(network, weightMap, *c3_10->getOutput(0), get_width(1024, gw), get_width(1024, gw), 5, "model.11"); + + // Head + auto conv12 = convBlock(network, weightMap, *sppf11->getOutput(0), get_width(768, gw), 1, 1, 1, "model.12"); + auto upsample13 = network->addResize(*conv12->getOutput(0)); + assert(upsample13); + upsample13->setResizeMode(ResizeMode::kNEAREST); + upsample13->setOutputDimensions(c3_8->getOutput(0)->getDimensions()); + ITensor* inputTensors14[] = { upsample13->getOutput(0), c3_8->getOutput(0) }; + auto cat14 = network->addConcatenation(inputTensors14, 2); + auto c3_15 = C3(network, weightMap, *cat14->getOutput(0), get_width(1536, gw), get_width(768, gw), get_depth(3, gd), false, 1, 0.5, "model.15"); + + auto conv16 = convBlock(network, weightMap, *c3_15->getOutput(0), get_width(512, gw), 1, 1, 1, "model.16"); + auto upsample17 = network->addResize(*conv16->getOutput(0)); + assert(upsample17); + upsample17->setResizeMode(ResizeMode::kNEAREST); + upsample17->setOutputDimensions(c3_6->getOutput(0)->getDimensions()); + ITensor* inputTensors18[] = { upsample17->getOutput(0), c3_6->getOutput(0) }; + auto cat18 = network->addConcatenation(inputTensors18, 2); + auto c3_19 = C3(network, weightMap, *cat18->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.19"); + + auto conv20 = convBlock(network, weightMap, *c3_19->getOutput(0), get_width(256, gw), 1, 1, 1, "model.20"); + auto upsample21 = network->addResize(*conv20->getOutput(0)); + assert(upsample21); + upsample21->setResizeMode(ResizeMode::kNEAREST); + upsample21->setOutputDimensions(c3_4->getOutput(0)->getDimensions()); + ITensor* inputTensors21[] = { upsample21->getOutput(0), c3_4->getOutput(0) }; + auto cat22 = network->addConcatenation(inputTensors21, 2); + auto c3_23 = C3(network, weightMap, *cat22->getOutput(0), get_width(512, gw), get_width(256, gw), get_depth(3, gd), false, 1, 0.5, "model.23"); + + auto conv24 = convBlock(network, weightMap, *c3_23->getOutput(0), get_width(256, gw), 3, 2, 1, "model.24"); + ITensor* inputTensors25[] = { conv24->getOutput(0), conv20->getOutput(0) }; + auto cat25 = network->addConcatenation(inputTensors25, 2); + auto c3_26 = C3(network, weightMap, *cat25->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.26"); + + auto conv27 = convBlock(network, weightMap, *c3_26->getOutput(0), get_width(512, gw), 3, 2, 1, "model.27"); + ITensor* inputTensors28[] = { conv27->getOutput(0), conv16->getOutput(0) }; + auto cat28 = network->addConcatenation(inputTensors28, 2); + auto c3_29 = C3(network, weightMap, *cat28->getOutput(0), get_width(1536, gw), get_width(768, gw), get_depth(3, gd), false, 1, 0.5, "model.29"); + + auto conv30 = convBlock(network, weightMap, *c3_29->getOutput(0), get_width(768, gw), 3, 2, 1, "model.30"); + ITensor* inputTensors31[] = { conv30->getOutput(0), conv12->getOutput(0) }; + auto cat31 = network->addConcatenation(inputTensors31, 2); + auto c3_32 = C3(network, weightMap, *cat31->getOutput(0), get_width(2048, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.32"); + + // Detect + IConvolutionLayer* det0 = network->addConvolutionNd(*c3_23->getOutput(0), 3 * (kNumClass + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.0.weight"], weightMap["model.33.m.0.bias"]); + IConvolutionLayer* det1 = network->addConvolutionNd(*c3_26->getOutput(0), 3 * (kNumClass + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.1.weight"], weightMap["model.33.m.1.bias"]); + IConvolutionLayer* det2 = network->addConvolutionNd(*c3_29->getOutput(0), 3 * (kNumClass + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.2.weight"], weightMap["model.33.m.2.bias"]); + IConvolutionLayer* det3 = network->addConvolutionNd(*c3_32->getOutput(0), 3 * (kNumClass + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.3.weight"], weightMap["model.33.m.3.bias"]); + + auto yolo = addYoLoLayer(network, weightMap, "model.33", std::vector{det0, det1, det2, det3}); + yolo->getOutput(0)->setName(kOutputTensorName); + network->markOutput(*yolo->getOutput(0)); + + // Engine config + builder->setMaxBatchSize(maxBatchSize); + config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB +#if defined(USE_FP16) + config->setFlag(BuilderFlag::kFP16); +#elif defined(USE_INT8) + std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; + assert(builder->platformHasFastInt8()); + config->setFlag(BuilderFlag::kINT8); + Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, kInputW, kInputH, "./coco_calib/", "int8calib.table", kInputTensorName); + config->setInt8Calibrator(calibrator); +#endif + + std::cout << "Building engine, please wait for a while..." << std::endl; + ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); + std::cout << "Build engine successfully!" << std::endl; + + // Don't need the network any more + network->destroy(); + + // Release host memory + for (auto& mem : weightMap) { + free((void*)(mem.second.values)); + } + + return engine; +} + +ICudaEngine* build_cls_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) { + INetworkDefinition* network = builder->createNetworkV2(0U); + + // Create input tensor + ITensor* data = network->addInput(kInputTensorName, dt, Dims3{ 3, kClsInputH, kClsInputW }); + assert(data); + std::map weightMap = loadWeights(wts_name); + + // Backbone + auto conv0 = convBlock(network, weightMap, *data, get_width(64, gw), 6, 2, 1, "model.0"); + assert(conv0); + auto conv1 = convBlock(network, weightMap, *conv0->getOutput(0), get_width(128, gw), 3, 2, 1, "model.1"); + auto bottleneck_CSP2 = C3(network, weightMap, *conv1->getOutput(0), get_width(128, gw), get_width(128, gw), get_depth(3, gd), true, 1, 0.5, "model.2"); + auto conv3 = convBlock(network, weightMap, *bottleneck_CSP2->getOutput(0), get_width(256, gw), 3, 2, 1, "model.3"); + auto bottleneck_csp4 = C3(network, weightMap, *conv3->getOutput(0), get_width(256, gw), get_width(256, gw), get_depth(6, gd), true, 1, 0.5, "model.4"); + auto conv5 = convBlock(network, weightMap, *bottleneck_csp4->getOutput(0), get_width(512, gw), 3, 2, 1, "model.5"); + auto bottleneck_csp6 = C3(network, weightMap, *conv5->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(9, gd), true, 1, 0.5, "model.6"); + auto conv7 = convBlock(network, weightMap, *bottleneck_csp6->getOutput(0), get_width(1024, gw), 3, 2, 1, "model.7"); + auto bottleneck_csp8 = C3(network, weightMap, *conv7->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), true, 1, 0.5, "model.8"); + + // Head + auto conv_class = convBlock(network, weightMap, *bottleneck_csp8->getOutput(0), 1280, 1, 1, 1, "model.9.conv"); + IPoolingLayer* pool2 = network->addPoolingNd(*conv_class->getOutput(0), PoolingType::kAVERAGE, DimsHW{7, 7}); + assert(pool2); + IFullyConnectedLayer* yolo = network->addFullyConnected(*pool2->getOutput(0), kClsNumClass, weightMap["model.9.linear.weight"], weightMap["model.9.linear.bias"]); + assert(yolo); + + yolo->getOutput(0)->setName(kOutputTensorName); + network->markOutput(*yolo->getOutput(0)); + + // Engine config + builder->setMaxBatchSize(maxBatchSize); + config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB + +#if defined(USE_FP16) + config->setFlag(BuilderFlag::kFP16); +#elif defined(USE_INT8) + std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; + assert(builder->platformHasFastInt8()); + config->setFlag(BuilderFlag::kINT8); + Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, kClsInputW, kClsInputW, "./coco_calib/", "int8calib.table", kInputTensorName); + config->setInt8Calibrator(calibrator); +#endif + + std::cout << "Building engine, please wait for a while..." << std::endl; + ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); + std::cout << "Build engine successfully!" << std::endl; + + // Don't need the network any more + network->destroy(); + + // Release host memory + for (auto& mem : weightMap) { + free((void*)(mem.second.values)); + } + + return engine; +} + +ICudaEngine* build_seg_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) { + INetworkDefinition* network = builder->createNetworkV2(0U); + ITensor* data = network->addInput(kInputTensorName, dt, Dims3{ 3, kInputH, kInputW }); + assert(data); + std::map weightMap = loadWeights(wts_name); + + // Backbone + auto conv0 = convBlock(network, weightMap, *data, get_width(64, gw), 6, 2, 1, "model.0"); + assert(conv0); + auto conv1 = convBlock(network, weightMap, *conv0->getOutput(0), get_width(128, gw), 3, 2, 1, "model.1"); + auto bottleneck_CSP2 = C3(network, weightMap, *conv1->getOutput(0), get_width(128, gw), get_width(128, gw), get_depth(3, gd), true, 1, 0.5, "model.2"); + auto conv3 = convBlock(network, weightMap, *bottleneck_CSP2->getOutput(0), get_width(256, gw), 3, 2, 1, "model.3"); + auto bottleneck_csp4 = C3(network, weightMap, *conv3->getOutput(0), get_width(256, gw), get_width(256, gw), get_depth(6, gd), true, 1, 0.5, "model.4"); + auto conv5 = convBlock(network, weightMap, *bottleneck_csp4->getOutput(0), get_width(512, gw), 3, 2, 1, "model.5"); + auto bottleneck_csp6 = C3(network, weightMap, *conv5->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(9, gd), true, 1, 0.5, "model.6"); + auto conv7 = convBlock(network, weightMap, *bottleneck_csp6->getOutput(0), get_width(1024, gw), 3, 2, 1, "model.7"); + auto bottleneck_csp8 = C3(network, weightMap, *conv7->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), true, 1, 0.5, "model.8"); + auto spp9 = SPPF(network, weightMap, *bottleneck_csp8->getOutput(0), get_width(1024, gw), get_width(1024, gw), 5, "model.9"); + + // Head + auto conv10 = convBlock(network, weightMap, *spp9->getOutput(0), get_width(512, gw), 1, 1, 1, "model.10"); + + auto upsample11 = network->addResize(*conv10->getOutput(0)); + assert(upsample11); + upsample11->setResizeMode(ResizeMode::kNEAREST); + upsample11->setOutputDimensions(bottleneck_csp6->getOutput(0)->getDimensions()); + + ITensor* inputTensors12[] = { upsample11->getOutput(0), bottleneck_csp6->getOutput(0) }; + auto cat12 = network->addConcatenation(inputTensors12, 2); + auto bottleneck_csp13 = C3(network, weightMap, *cat12->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.13"); + auto conv14 = convBlock(network, weightMap, *bottleneck_csp13->getOutput(0), get_width(256, gw), 1, 1, 1, "model.14"); + + auto upsample15 = network->addResize(*conv14->getOutput(0)); + assert(upsample15); + upsample15->setResizeMode(ResizeMode::kNEAREST); + upsample15->setOutputDimensions(bottleneck_csp4->getOutput(0)->getDimensions()); + + ITensor* inputTensors16[] = { upsample15->getOutput(0), bottleneck_csp4->getOutput(0) }; + auto cat16 = network->addConcatenation(inputTensors16, 2); + + auto bottleneck_csp17 = C3(network, weightMap, *cat16->getOutput(0), get_width(512, gw), get_width(256, gw), get_depth(3, gd), false, 1, 0.5, "model.17"); + + // Segmentation + IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (32 + kNumClass + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]); + auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), get_width(256, gw), 3, 2, 1, "model.18"); + ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) }; + auto cat19 = network->addConcatenation(inputTensors19, 2); + auto bottleneck_csp20 = C3(network, weightMap, *cat19->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.20"); + IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (32 + kNumClass + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]); + auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), get_width(512, gw), 3, 2, 1, "model.21"); + ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) }; + auto cat22 = network->addConcatenation(inputTensors22, 2); + auto bottleneck_csp23 = C3(network, weightMap, *cat22->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.23"); + IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (32 + kNumClass + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]); + + auto yolo = addYoLoLayer(network, weightMap, "model.24", std::vector{det0, det1, det2}, true); + yolo->getOutput(0)->setName(kOutputTensorName); + network->markOutput(*yolo->getOutput(0)); + + auto proto = Proto(network, weightMap, *bottleneck_csp17->getOutput(0), get_width(256, gw), 32, "model.24.proto"); + proto->getOutput(0)->setName("proto"); + network->markOutput(*proto->getOutput(0)); + + // Engine config + builder->setMaxBatchSize(maxBatchSize); + config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB +#if defined(USE_FP16) + config->setFlag(BuilderFlag::kFP16); +#elif defined(USE_INT8) + std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; + assert(builder->platformHasFastInt8()); + config->setFlag(BuilderFlag::kINT8); + Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, kInputW, kInputH, "./coco_calib/", "int8calib.table", kInputTensorName); + config->setInt8Calibrator(calibrator); +#endif + + std::cout << "Building engine, please wait for a while..." << std::endl; + ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); + std::cout << "Build engine successfully!" << std::endl; + + // Don't need the network any more + network->destroy(); + + // Release host memory + for (auto& mem : weightMap) { + free((void*)(mem.second.values)); + } + + return engine; +} + diff --git a/yolov5/src/model.h b/yolov5/src/model.h new file mode 100644 index 0000000..73a5f62 --- /dev/null +++ b/yolov5/src/model.h @@ -0,0 +1,16 @@ +#pragma once + +#include +#include + +nvinfer1::ICudaEngine* build_det_engine(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, + nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, + float& gd, float& gw, std::string& wts_name); + +nvinfer1::ICudaEngine* build_det_p6_engine(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, + nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, + float& gd, float& gw, std::string& wts_name); + +nvinfer1::ICudaEngine* build_cls_engine(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, float& gd, float& gw, std::string& wts_name); + +nvinfer1::ICudaEngine* build_seg_engine(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, float& gd, float& gw, std::string& wts_name); diff --git a/yolov5/src/postprocess.cpp b/yolov5/src/postprocess.cpp new file mode 100644 index 0000000..70e4033 --- /dev/null +++ b/yolov5/src/postprocess.cpp @@ -0,0 +1,189 @@ +#include "postprocess.h" +#include "utils.h" + +cv::Rect get_rect(cv::Mat& img, float bbox[4]) { + float l, r, t, b; + float r_w = kInputW / (img.cols * 1.0); + float r_h = kInputH / (img.rows * 1.0); + if (r_h > r_w) { + l = bbox[0] - bbox[2] / 2.f; + r = bbox[0] + bbox[2] / 2.f; + t = bbox[1] - bbox[3] / 2.f - (kInputH - r_w * img.rows) / 2; + b = bbox[1] + bbox[3] / 2.f - (kInputH - r_w * img.rows) / 2; + l = l / r_w; + r = r / r_w; + t = t / r_w; + b = b / r_w; + } else { + l = bbox[0] - bbox[2] / 2.f - (kInputW - r_h * img.cols) / 2; + r = bbox[0] + bbox[2] / 2.f - (kInputW - r_h * img.cols) / 2; + t = bbox[1] - bbox[3] / 2.f; + b = bbox[1] + bbox[3] / 2.f; + l = l / r_h; + r = r / r_h; + t = t / r_h; + b = b / r_h; + } + return cv::Rect(round(l), round(t), round(r - l), round(b - t)); +} + +static float iou(float lbox[4], float rbox[4]) { + float interBox[] = { + (std::max)(lbox[0] - lbox[2] / 2.f , rbox[0] - rbox[2] / 2.f), //left + (std::min)(lbox[0] + lbox[2] / 2.f , rbox[0] + rbox[2] / 2.f), //right + (std::max)(lbox[1] - lbox[3] / 2.f , rbox[1] - rbox[3] / 2.f), //top + (std::min)(lbox[1] + lbox[3] / 2.f , rbox[1] + rbox[3] / 2.f), //bottom + }; + + if (interBox[2] > interBox[3] || interBox[0] > interBox[1]) + return 0.0f; + + float interBoxS = (interBox[1] - interBox[0])*(interBox[3] - interBox[2]); + return interBoxS / (lbox[2] * lbox[3] + rbox[2] * rbox[3] - interBoxS); +} + +static bool cmp(const Detection& a, const Detection& b) { + return a.conf > b.conf; +} + +void nms(std::vector& res, float* output, float conf_thresh, float nms_thresh) { + int det_size = sizeof(Detection) / sizeof(float); + std::map> m; + for (int i = 0; i < output[0] && i < kMaxNumOutputBbox; i++) { + if (output[1 + det_size * i + 4] <= conf_thresh) continue; + Detection det; + memcpy(&det, &output[1 + det_size * i], det_size * sizeof(float)); + if (m.count(det.class_id) == 0) m.emplace(det.class_id, std::vector()); + m[det.class_id].push_back(det); + } + for (auto it = m.begin(); it != m.end(); it++) { + auto& dets = it->second; + std::sort(dets.begin(), dets.end(), cmp); + for (size_t m = 0; m < dets.size(); ++m) { + auto& item = dets[m]; + res.push_back(item); + for (size_t n = m + 1; n < dets.size(); ++n) { + if (iou(item.bbox, dets[n].bbox) > nms_thresh) { + dets.erase(dets.begin() + n); + --n; + } + } + } + } +} + +void batch_nms(std::vector>& res_batch, float *output, int batch_size, int output_size, float conf_thresh, float nms_thresh) { + res_batch.resize(batch_size); + for (int i = 0; i < batch_size; i++) { + nms(res_batch[i], &output[i * output_size], conf_thresh, nms_thresh); + } +} + +void draw_bbox(std::vector& img_batch, std::vector>& res_batch) { + for (size_t i = 0; i < img_batch.size(); i++) { + auto& res = res_batch[i]; + cv::Mat img = img_batch[i]; + for (size_t j = 0; j < res.size(); j++) { + cv::Rect r = get_rect(img, res[j].bbox); + cv::rectangle(img, r, cv::Scalar(0x27, 0xC1, 0x36), 2); + cv::putText(img, std::to_string((int)res[j].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2); + } + } +} + +static cv::Rect get_downscale_rect(float bbox[4], float scale) { + float left = bbox[0] - bbox[2] / 2; + float top = bbox[1] - bbox[3] / 2; + float right = bbox[0] + bbox[2] / 2; + float bottom = bbox[1] + bbox[3] / 2; + left /= scale; + top /= scale; + right /= scale; + bottom /= scale; + return cv::Rect(round(left), round(top), round(right - left), round(bottom - top)); +} + +std::vector process_mask(const float* proto, int proto_size, std::vector& dets) { + std::vector masks; + for (size_t i = 0; i < dets.size(); i++) { + cv::Mat mask_mat = cv::Mat::zeros(kInputH / 4, kInputW / 4, CV_32FC1); + auto r = get_downscale_rect(dets[i].bbox, 4); + for (int x = r.x; x < r.x + r.width; x++) { + for (int y = r.y; y < r.y + r.height; y++) { + float e = 0.0f; + for (int j = 0; j < 32; j++) { + e += dets[i].mask[j] * proto[j * proto_size / 32 + y * mask_mat.cols + x]; + } + e = 1.0f / (1.0f + expf(-e)); + mask_mat.at(y, x) = e; + } + } + cv::resize(mask_mat, mask_mat, cv::Size(kInputW, kInputH)); + masks.push_back(mask_mat); + } + return masks; +} + +cv::Mat scale_mask(cv::Mat mask, cv::Mat img) { + int x, y, w, h; + float r_w = kInputW / (img.cols * 1.0); + float r_h = kInputH / (img.rows * 1.0); + if (r_h > r_w) { + w = kInputW; + h = r_w * img.rows; + x = 0; + y = (kInputH - h) / 2; + } else { + w = r_h * img.cols; + h = kInputH; + x = (kInputW - w) / 2; + y = 0; + } + cv::Rect r(x, y, w, h); + cv::Mat res; + cv::resize(mask(r), res, img.size()); + return res; +} + +void draw_mask_bbox(cv::Mat& img, std::vector& dets, std::vector& masks, std::unordered_map& labels_map) { + static std::vector colors = {0xFF3838, 0xFF9D97, 0xFF701F, 0xFFB21D, 0xCFD231, 0x48F90A, + 0x92CC17, 0x3DDB86, 0x1A9334, 0x00D4BB, 0x2C99A8, 0x00C2FF, + 0x344593, 0x6473FF, 0x0018EC, 0x8438FF, 0x520085, 0xCB38FF, + 0xFF95C8, 0xFF37C7}; + for (size_t i = 0; i < dets.size(); i++) { + cv::Mat img_mask = scale_mask(masks[i], img); + auto color = colors[(int)dets[i].class_id % colors.size()]; + auto bgr = cv::Scalar(color & 0xFF, color >> 8 & 0xFF, color >> 16 & 0xFF); + + cv::Rect r = get_rect(img, dets[i].bbox); + for (int x = r.x; x < r.x + r.width; x++) { + for (int y = r.y; y < r.y + r.height; y++) { + float val = img_mask.at(y, x); + if (val <= 0.5) continue; + img.at(y, x)[0] = img.at(y, x)[0] / 2 + bgr[0] / 2; + img.at(y, x)[1] = img.at(y, x)[1] / 2 + bgr[1] / 2; + img.at(y, x)[2] = img.at(y, x)[2] / 2 + bgr[2] / 2; + } + } + + cv::rectangle(img, r, bgr, 2); + + // Get the size of the text + cv::Size textSize = cv::getTextSize(labels_map[(int)dets[i].class_id] + " " + to_string_with_precision(dets[i].conf), cv::FONT_HERSHEY_PLAIN, 1.2, 2, NULL); + // Set the top left corner of the rectangle + cv::Point topLeft(r.x, r.y - textSize.height); + + // Set the bottom right corner of the rectangle + cv::Point bottomRight(r.x + textSize.width, r.y + textSize.height); + + // Set the thickness of the rectangle lines + int lineThickness = 2; + + // Draw the rectangle on the image + cv::rectangle(img, topLeft, bottomRight, bgr, -1); + + cv::putText(img, labels_map[(int)dets[i].class_id] + " " + to_string_with_precision(dets[i].conf), cv::Point(r.x, r.y + 4), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar::all(0xFF), 2); + + } +} + diff --git a/yolov5/src/postprocess.h b/yolov5/src/postprocess.h new file mode 100644 index 0000000..904ccb0 --- /dev/null +++ b/yolov5/src/postprocess.h @@ -0,0 +1,16 @@ +#pragma once + +#include "types.h" +#include + +cv::Rect get_rect(cv::Mat& img, float bbox[4]); + +void nms(std::vector& res, float *output, float conf_thresh, float nms_thresh = 0.5); + +void batch_nms(std::vector>& batch_res, float *output, int batch_size, int output_size, float conf_thresh, float nms_thresh = 0.5); + +void draw_bbox(std::vector& img_batch, std::vector>& res_batch); + +std::vector process_mask(const float* proto, int proto_size, std::vector& dets); + +void draw_mask_bbox(cv::Mat& img, std::vector& dets, std::vector& masks, std::unordered_map& labels_map); diff --git a/yolov5/src/preprocess.cu b/yolov5/src/preprocess.cu index 51359cc..8de0093 100644 --- a/yolov5/src/preprocess.cu +++ b/yolov5/src/preprocess.cu @@ -1,116 +1,153 @@ #include "preprocess.h" -#include +#include "cuda_utils.h" -__global__ void warpaffine_kernel( - uint8_t* src, int src_line_size, int src_width, - int src_height, float* dst, int dst_width, +static uint8_t* img_buffer_host = nullptr; +static uint8_t* img_buffer_device = nullptr; + +struct AffineMatrix { + float value[6]; +}; + +__global__ void warpaffine_kernel( + uint8_t* src, int src_line_size, int src_width, + int src_height, float* dst, int dst_width, int dst_height, uint8_t const_value_st, AffineMatrix d2s, int edge) { - int position = blockDim.x * blockIdx.x + threadIdx.x; - if (position >= edge) return; + int position = blockDim.x * blockIdx.x + threadIdx.x; + if (position >= edge) return; - float m_x1 = d2s.value[0]; - float m_y1 = d2s.value[1]; - float m_z1 = d2s.value[2]; - float m_x2 = d2s.value[3]; - float m_y2 = d2s.value[4]; - float m_z2 = d2s.value[5]; + float m_x1 = d2s.value[0]; + float m_y1 = d2s.value[1]; + float m_z1 = d2s.value[2]; + float m_x2 = d2s.value[3]; + float m_y2 = d2s.value[4]; + float m_z2 = d2s.value[5]; - int dx = position % dst_width; - int dy = position / dst_width; - float src_x = m_x1 * dx + m_y1 * dy + m_z1 + 0.5f; - float src_y = m_x2 * dx + m_y2 * dy + m_z2 + 0.5f; - float c0, c1, c2; + int dx = position % dst_width; + int dy = position / dst_width; + float src_x = m_x1 * dx + m_y1 * dy + m_z1 + 0.5f; + float src_y = m_x2 * dx + m_y2 * dy + m_z2 + 0.5f; + float c0, c1, c2; - if (src_x <= -1 || src_x >= src_width || src_y <= -1 || src_y >= src_height) { - // out of range - c0 = const_value_st; - c1 = const_value_st; - c2 = const_value_st; - } else { - int y_low = floorf(src_y); - int x_low = floorf(src_x); - int y_high = y_low + 1; - int x_high = x_low + 1; + if (src_x <= -1 || src_x >= src_width || src_y <= -1 || src_y >= src_height) { + // out of range + c0 = const_value_st; + c1 = const_value_st; + c2 = const_value_st; + } else { + int y_low = floorf(src_y); + int x_low = floorf(src_x); + int y_high = y_low + 1; + int x_high = x_low + 1; - uint8_t const_value[] = {const_value_st, const_value_st, const_value_st}; - float ly = src_y - y_low; - float lx = src_x - x_low; - float hy = 1 - ly; - float hx = 1 - lx; - float w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx; - uint8_t* v1 = const_value; - uint8_t* v2 = const_value; - uint8_t* v3 = const_value; - uint8_t* v4 = const_value; + uint8_t const_value[] = {const_value_st, const_value_st, const_value_st}; + float ly = src_y - y_low; + float lx = src_x - x_low; + float hy = 1 - ly; + float hx = 1 - lx; + float w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx; + uint8_t* v1 = const_value; + uint8_t* v2 = const_value; + uint8_t* v3 = const_value; + uint8_t* v4 = const_value; - if (y_low >= 0) { - if (x_low >= 0) - v1 = src + y_low * src_line_size + x_low * 3; + if (y_low >= 0) { + if (x_low >= 0) + v1 = src + y_low * src_line_size + x_low * 3; - if (x_high < src_width) - v2 = src + y_low * src_line_size + x_high * 3; - } - - if (y_high < src_height) { - if (x_low >= 0) - v3 = src + y_high * src_line_size + x_low * 3; - - if (x_high < src_width) - v4 = src + y_high * src_line_size + x_high * 3; - } - - c0 = w1 * v1[0] + w2 * v2[0] + w3 * v3[0] + w4 * v4[0]; - c1 = w1 * v1[1] + w2 * v2[1] + w3 * v3[1] + w4 * v4[1]; - c2 = w1 * v1[2] + w2 * v2[2] + w3 * v3[2] + w4 * v4[2]; + if (x_high < src_width) + v2 = src + y_low * src_line_size + x_high * 3; } - //bgr to rgb - float t = c2; - c2 = c0; - c0 = t; + if (y_high < src_height) { + if (x_low >= 0) + v3 = src + y_high * src_line_size + x_low * 3; - //normalization - c0 = c0 / 255.0f; - c1 = c1 / 255.0f; - c2 = c2 / 255.0f; + if (x_high < src_width) + v4 = src + y_high * src_line_size + x_high * 3; + } - //rgbrgbrgb to rrrgggbbb - int area = dst_width * dst_height; - float* pdst_c0 = dst + dy * dst_width + dx; - float* pdst_c1 = pdst_c0 + area; - float* pdst_c2 = pdst_c1 + area; - *pdst_c0 = c0; - *pdst_c1 = c1; - *pdst_c2 = c2; + c0 = w1 * v1[0] + w2 * v2[0] + w3 * v3[0] + w4 * v4[0]; + c1 = w1 * v1[1] + w2 * v2[1] + w3 * v3[1] + w4 * v4[1]; + c2 = w1 * v1[2] + w2 * v2[2] + w3 * v3[2] + w4 * v4[2]; + } + + // bgr to rgb + float t = c2; + c2 = c0; + c0 = t; + + // normalization + c0 = c0 / 255.0f; + c1 = c1 / 255.0f; + c2 = c2 / 255.0f; + + // rgbrgbrgb to rrrgggbbb + int area = dst_width * dst_height; + float* pdst_c0 = dst + dy * dst_width + dx; + float* pdst_c1 = pdst_c0 + area; + float* pdst_c2 = pdst_c1 + area; + *pdst_c0 = c0; + *pdst_c1 = c1; + *pdst_c2 = c2; } -void preprocess_kernel_img( +void cuda_preprocess( uint8_t* src, int src_width, int src_height, float* dst, int dst_width, int dst_height, cudaStream_t stream) { - AffineMatrix s2d,d2s; - float scale = std::min(dst_height / (float)src_height, dst_width / (float)src_width); - s2d.value[0] = scale; - s2d.value[1] = 0; - s2d.value[2] = -scale * src_width * 0.5 + dst_width * 0.5; - s2d.value[3] = 0; - s2d.value[4] = scale; - s2d.value[5] = -scale * src_height * 0.5 + dst_height * 0.5; + int img_size = src_width * src_height * 3; + // copy data to pinned memory + memcpy(img_buffer_host, src, img_size); + // copy data to device memory + CUDA_CHECK(cudaMemcpyAsync(img_buffer_device, img_buffer_host, img_size, cudaMemcpyHostToDevice, stream)); - cv::Mat m2x3_s2d(2, 3, CV_32F, s2d.value); - cv::Mat m2x3_d2s(2, 3, CV_32F, d2s.value); - cv::invertAffineTransform(m2x3_s2d, m2x3_d2s); + AffineMatrix s2d, d2s; + float scale = std::min(dst_height / (float)src_height, dst_width / (float)src_width); - memcpy(d2s.value, m2x3_d2s.ptr(0), sizeof(d2s.value)); + s2d.value[0] = scale; + s2d.value[1] = 0; + s2d.value[2] = -scale * src_width * 0.5 + dst_width * 0.5; + s2d.value[3] = 0; + s2d.value[4] = scale; + s2d.value[5] = -scale * src_height * 0.5 + dst_height * 0.5; - int jobs = dst_height * dst_width; - int threads = 256; - int blocks = ceil(jobs / (float)threads); - warpaffine_kernel<<>>( - src, src_width*3, src_width, - src_height, dst, dst_width, - dst_height, 128, d2s, jobs); + cv::Mat m2x3_s2d(2, 3, CV_32F, s2d.value); + cv::Mat m2x3_d2s(2, 3, CV_32F, d2s.value); + cv::invertAffineTransform(m2x3_s2d, m2x3_d2s); + memcpy(d2s.value, m2x3_d2s.ptr(0), sizeof(d2s.value)); + + int jobs = dst_height * dst_width; + int threads = 256; + int blocks = ceil(jobs / (float)threads); + + warpaffine_kernel<<>>( + img_buffer_device, src_width * 3, src_width, + src_height, dst, dst_width, + dst_height, 128, d2s, jobs); } + +void cuda_batch_preprocess(std::vector& img_batch, + float* dst, int dst_width, int dst_height, + cudaStream_t stream) { + int dst_size = dst_width * dst_height * 3; + for (size_t i = 0; i < img_batch.size(); i++) { + cuda_preprocess(img_batch[i].ptr(), img_batch[i].cols, img_batch[i].rows, &dst[dst_size * i], dst_width, dst_height, stream); + CUDA_CHECK(cudaStreamSynchronize(stream)); + } +} + +void cuda_preprocess_init(int max_image_size) { + // prepare input data in pinned memory + CUDA_CHECK(cudaMallocHost((void**)&img_buffer_host, max_image_size * 3)); + // prepare input data in device memory + CUDA_CHECK(cudaMalloc((void**)&img_buffer_device, max_image_size * 3)); +} + +void cuda_preprocess_destroy() { + CUDA_CHECK(cudaFree(img_buffer_device)); + CUDA_CHECK(cudaFreeHost(img_buffer_host)); +} + diff --git a/yolov5/src/preprocess.h b/yolov5/src/preprocess.h index 236d5f3..c0dc1aa 100644 --- a/yolov5/src/preprocess.h +++ b/yolov5/src/preprocess.h @@ -1,16 +1,15 @@ -#ifndef __PREPROCESS_H -#define __PREPROCESS_H +#pragma once #include #include +#include - -struct AffineMatrix{ - float value[6]; -}; - - -void preprocess_kernel_img(uint8_t* src, int src_width, int src_height, +void cuda_preprocess_init(int max_image_size); +void cuda_preprocess_destroy(); +void cuda_preprocess(uint8_t* src, int src_width, int src_height, + float* dst, int dst_width, int dst_height, + cudaStream_t stream); +void cuda_batch_preprocess(std::vector& img_batch, float* dst, int dst_width, int dst_height, cudaStream_t stream); -#endif // __PREPROCESS_H + diff --git a/yolov5/src/types.h b/yolov5/src/types.h new file mode 100644 index 0000000..8004eda --- /dev/null +++ b/yolov5/src/types.h @@ -0,0 +1,17 @@ +#pragma once + +#include "config.h" + +struct YoloKernel { + int width; + int height; + float anchors[kNumAnchor * 2]; +}; + +struct alignas(float) Detection { + float bbox[4]; // center_x center_y w h + float conf; // bbox_conf * cls_conf + float class_id; + float mask[32]; +}; + diff --git a/yolov5/src/utils.h b/yolov5/src/utils.h index 9483728..2dea946 100644 --- a/yolov5/src/utils.h +++ b/yolov5/src/utils.h @@ -1,34 +1,12 @@ -#ifndef TRTX_YOLOV5_UTILS_H_ -#define TRTX_YOLOV5_UTILS_H_ +#pragma once #include -#include #include #include #include #include - -static inline cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h) { - int w, h, x, y; - float r_w = input_w / (img.cols * 1.0); - float r_h = input_h / (img.rows * 1.0); - if (r_h > r_w) { - w = input_w; - h = r_w * img.rows; - x = 0; - y = (input_h - h) / 2; - } else { - w = r_h * img.cols; - h = input_h; - x = (input_w - w) / 2; - y = 0; - } - cv::Mat re(h, w, CV_8UC3); - cv::resize(img, re, re.size(), 0, 0, cv::INTER_LINEAR); - cv::Mat out(input_h, input_w, CV_8UC3, cv::Scalar(128, 128, 128)); - re.copyTo(out(cv::Rect(x, y, re.cols, re.rows))); - return out; -} +#include +#include static inline int read_files_in_dir(const char* p_dir_name, std::vector& file_names) { DIR *p_dir = opendir(p_dir_name); @@ -61,6 +39,7 @@ static inline std::string trim_leading_whitespace(const std::string& str) { size_t last = str.find_last_not_of(' '); return str.substr(first, (last - first + 1)); } + // Src: https://stackoverflow.com/questions/16605967 static inline std::string to_string_with_precision(const float a_value, const int n = 2) { std::ostringstream out; @@ -89,5 +68,3 @@ static inline int read_labels(const std::string labels_filename, std::unordered_ return 0; } -#endif // TRTX_YOLOV5_UTILS_H_ - diff --git a/yolov5/yolov5_cls.cpp b/yolov5/yolov5_cls.cpp index c4bf1b8..7ce6052 100644 --- a/yolov5/yolov5_cls.cpp +++ b/yolov5/yolov5_cls.cpp @@ -1,333 +1,286 @@ +#include "cuda_utils.h" +#include "logging.h" +#include "utils.h" +#include "model.h" +#include "config.h" + #include #include #include #include -#include "cuda_utils.h" -#include "logging.h" -#include "common.hpp" -#include "utils.h" -#include "calibrator.h" +#include -#define USE_FP32 // set USE_INT8 or USE_FP16 or USE_FP32 -#define DEVICE 0 // GPU id -#define BATCH_SIZE 1 +using namespace nvinfer1; -// stuff we know about the network and the input/output blobs -static const int INPUT_H = 224; -static const int INPUT_W = 224; -static const int CLASS_NUM = 1000; - -static const int OUTPUT_SIZE = CLASS_NUM; -const char* INPUT_BLOB_NAME = "data"; -const char* OUTPUT_BLOB_NAME = "prob"; static Logger gLogger; +const static int kOutputSize = kClsNumClass; -static int get_width(int x, float gw, int divisor = 8) { - return int(ceil((x * gw) / divisor)) * divisor; -} - -static int get_depth(int x, float gd) { - if (x == 1) return 1; - int r = round(x * gd); - if (x * gd - int(x * gd) == 0.5 && (int(x * gd) % 2) == 0) { - --r; +void batch_preprocess(std::vector& imgs, float* output) { + for (size_t b = 0; b < imgs.size(); b++) { + cv::Mat img; + cv::resize(imgs[b], img, cv::Size(kClsInputW, kClsInputH)); + int i = 0; + for (int row = 0; row < img.rows; ++row) { + uchar* uc_pixel = img.data + row * img.step; + for (int col = 0; col < img.cols; ++col) { + output[b * 3 * img.rows * img.cols + i] = ((float)uc_pixel[2] / 255.0 - 0.485) / 0.229; // R - 0.485 + output[b * 3 * img.rows * img.cols + i + img.rows * img.cols] = ((float)uc_pixel[1] / 255.0 - 0.456) / 0.224; + output[b * 3 * img.rows * img.cols + i + 2 * img.rows * img.cols] = ((float)uc_pixel[0] / 255.0 - 0.406) / 0.225; + uc_pixel += 3; + ++i; + } } - return std::max(r, 1); + } } std::vector softmax(float *prob, int n) { - std::vector res; - float sum = 0.0f; - float t; - for (int i = 0; i < n; i++) { - t = expf(prob[i]); - res.push_back(t); - sum += t; - } - for (int i = 0; i < n; i++) { - res[i] /= sum; - } - return res; + std::vector res; + float sum = 0.0f; + float t; + for (int i = 0; i < n; i++) { + t = expf(prob[i]); + res.push_back(t); + sum += t; + } + for (int i = 0; i < n; i++) { + res[i] /= sum; + } + return res; } std::vector topk(const std::vector& vec, int k) { - std::vector topk_index; - std::vector vec_index(vec.size()); - std::iota(vec_index.begin(), vec_index.end(), 0); + std::vector topk_index; + std::vector vec_index(vec.size()); + std::iota(vec_index.begin(), vec_index.end(), 0); - std::sort(vec_index.begin(), vec_index.end(), [&vec](size_t index_1, size_t index_2) { return vec[index_1] > vec[index_2]; }); + std::sort(vec_index.begin(), vec_index.end(), [&vec](size_t index_1, size_t index_2) { return vec[index_1] > vec[index_2]; }); - int k_num = std::min(vec.size(), k); + int k_num = std::min(vec.size(), k); - for (int i = 0; i < k_num; ++i) { - topk_index.push_back(vec_index[i]); - } + for (int i = 0; i < k_num; ++i) { + topk_index.push_back(vec_index[i]); + } - return topk_index; + return topk_index; } std::vector read_classes(std::string file_name) { - std::vector classes; - std::ifstream ifs(file_name, std::ios::in); - if (!ifs.is_open()) { - std::cerr << file_name << " is not found, pls refer to README and download it." << std::endl; - assert(0); - } - std::string s; - while (std::getline(ifs, s)) { - classes.push_back(s); - } - ifs.close(); - return classes; -} - -ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) { - INetworkDefinition* network = builder->createNetworkV2(0U); - - // Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME - ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{ 3, INPUT_H, INPUT_W }); - assert(data); - std::map weightMap = loadWeights(wts_name); - /* ------ yolov5 backbone------ */ - auto conv0 = convBlock(network, weightMap, *data, get_width(64, gw), 6, 2, 1, "model.0"); - assert(conv0); - auto conv1 = convBlock(network, weightMap, *conv0->getOutput(0), get_width(128, gw), 3, 2, 1, "model.1"); - auto bottleneck_CSP2 = C3(network, weightMap, *conv1->getOutput(0), get_width(128, gw), get_width(128, gw), get_depth(3, gd), true, 1, 0.5, "model.2"); - auto conv3 = convBlock(network, weightMap, *bottleneck_CSP2->getOutput(0), get_width(256, gw), 3, 2, 1, "model.3"); - auto bottleneck_csp4 = C3(network, weightMap, *conv3->getOutput(0), get_width(256, gw), get_width(256, gw), get_depth(6, gd), true, 1, 0.5, "model.4"); - auto conv5 = convBlock(network, weightMap, *bottleneck_csp4->getOutput(0), get_width(512, gw), 3, 2, 1, "model.5"); - auto bottleneck_csp6 = C3(network, weightMap, *conv5->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(9, gd), true, 1, 0.5, "model.6"); - auto conv7 = convBlock(network, weightMap, *bottleneck_csp6->getOutput(0), get_width(1024, gw), 3, 2, 1, "model.7"); - auto bottleneck_csp8 = C3(network, weightMap, *conv7->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), true, 1, 0.5, "model.8"); - - /* ------ yolov5 classification head ------ */ - auto conv_class = convBlock(network, weightMap, *bottleneck_csp8->getOutput(0), 1280, 1, 1, 1, "model.9.conv"); - IPoolingLayer* pool2 = network->addPoolingNd(*conv_class->getOutput(0), PoolingType::kAVERAGE, DimsHW{7, 7}); - assert(pool2); - IFullyConnectedLayer* yolo = network->addFullyConnected(*pool2->getOutput(0), CLASS_NUM, weightMap["model.9.linear.weight"], weightMap["model.9.linear.bias"]); - assert(yolo); - - yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME); - network->markOutput(*yolo->getOutput(0)); - // Build engine - builder->setMaxBatchSize(maxBatchSize); - config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB - -#if defined(USE_FP16) - config->setFlag(BuilderFlag::kFP16); -#elif defined(USE_INT8) - std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; - assert(builder->platformHasFastInt8()); - config->setFlag(BuilderFlag::kINT8); - Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, INPUT_W, INPUT_H, "./coco_calib/", "int8calib.table", INPUT_BLOB_NAME); - config->setInt8Calibrator(calibrator); -#endif - - std::cout << "Building engine, please wait for a while..." << std::endl; - ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); - std::cout << "Build engine successfully!" << std::endl; - - // Don't need the network any more - network->destroy(); - - // Release host memory - for (auto& mem : weightMap) { - free((void*)(mem.second.values)); - } - - return engine; -} - -void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream, float& gd, float& gw, std::string& wts_name) { - // Create builder - IBuilder* builder = createInferBuilder(gLogger); - IBuilderConfig* config = builder->createBuilderConfig(); - - // Create model to populate the network, then set the outputs and create an engine - ICudaEngine *engine = nullptr; - - engine = build_engine(maxBatchSize, builder, config, DataType::kFLOAT, gd, gw, wts_name); - - assert(engine != nullptr); - - // Serialize the engine - (*modelStream) = engine->serialize(); - - // Close everything down - engine->destroy(); - builder->destroy(); - config->destroy(); -} - -void doInference(IExecutionContext& context, cudaStream_t& stream, void **buffers, float* input, float* output, int batchSize) { - // infer on the batch asynchronously, and DMA output back to host - CUDA_CHECK(cudaMemcpyAsync(buffers[0], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream)); - context.enqueue(batchSize, buffers, stream, nullptr); - CUDA_CHECK(cudaMemcpyAsync(output, buffers[1], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream)); - cudaStreamSynchronize(stream); + std::vector classes; + std::ifstream ifs(file_name, std::ios::in); + if (!ifs.is_open()) { + std::cerr << file_name << " is not found, pls refer to README and download it." << std::endl; + assert(0); + } + std::string s; + while (std::getline(ifs, s)) { + classes.push_back(s); + } + ifs.close(); + return classes; } bool parse_args(int argc, char** argv, std::string& wts, std::string& engine, float& gd, float& gw, std::string& img_dir) { - if (argc < 4) return false; - if (std::string(argv[1]) == "-s" && (argc == 5 || argc == 7)) { - wts = std::string(argv[2]); - engine = std::string(argv[3]); - auto net = std::string(argv[4]); - if (net[0] == 'n') { - gd = 0.33; - gw = 0.25; - } else if (net[0] == 's') { - gd = 0.33; - gw = 0.50; - } else if (net[0] == 'm') { - gd = 0.67; - gw = 0.75; - } else if (net[0] == 'l') { - gd = 1.0; - gw = 1.0; - } else if (net[0] == 'x') { - gd = 1.33; - gw = 1.25; - } else if (net[0] == 'c' && argc == 7) { - gd = atof(argv[5]); - gw = atof(argv[6]); - } else { - return false; - } - } else if (std::string(argv[1]) == "-d" && argc == 4) { - engine = std::string(argv[2]); - img_dir = std::string(argv[3]); + if (argc < 4) return false; + if (std::string(argv[1]) == "-s" && (argc == 5 || argc == 7)) { + wts = std::string(argv[2]); + engine = std::string(argv[3]); + auto net = std::string(argv[4]); + if (net[0] == 'n') { + gd = 0.33; + gw = 0.25; + } else if (net[0] == 's') { + gd = 0.33; + gw = 0.50; + } else if (net[0] == 'm') { + gd = 0.67; + gw = 0.75; + } else if (net[0] == 'l') { + gd = 1.0; + gw = 1.0; + } else if (net[0] == 'x') { + gd = 1.33; + gw = 1.25; + } else if (net[0] == 'c' && argc == 7) { + gd = atof(argv[5]); + gw = atof(argv[6]); } else { - return false; + return false; } - return true; + } else if (std::string(argv[1]) == "-d" && argc == 4) { + engine = std::string(argv[2]); + img_dir = std::string(argv[3]); + } else { + return false; + } + return true; +} + +void prepare_buffers(ICudaEngine* engine, float** gpu_input_buffer, float** gpu_output_buffer, float** cpu_input_buffer, float** cpu_output_buffer) { + assert(engine->getNbBindings() == 2); + // In order to bind the buffers, we need to know the names of the input and output tensors. + // Note that indices are guaranteed to be less than IEngine::getNbBindings() + const int inputIndex = engine->getBindingIndex(kInputTensorName); + const int outputIndex = engine->getBindingIndex(kOutputTensorName); + assert(inputIndex == 0); + assert(outputIndex == 1); + // Create GPU buffers on device + CUDA_CHECK(cudaMalloc((void**)gpu_input_buffer, kBatchSize * 3 * kClsInputH * kClsInputW * sizeof(float))); + CUDA_CHECK(cudaMalloc((void**)gpu_output_buffer, kBatchSize * kOutputSize * sizeof(float))); + + *cpu_input_buffer = new float[kBatchSize * 3 * kClsInputH * kClsInputW]; + *cpu_output_buffer = new float[kBatchSize * kOutputSize]; +} + +void infer(IExecutionContext& context, cudaStream_t& stream, void **buffers, float* input, float* output, int batchSize) { + CUDA_CHECK(cudaMemcpyAsync(buffers[0], input, batchSize * 3 * kClsInputH * kClsInputW * sizeof(float), cudaMemcpyHostToDevice, stream)); + context.enqueue(batchSize, buffers, stream, nullptr); + CUDA_CHECK(cudaMemcpyAsync(output, buffers[1], batchSize * kOutputSize * sizeof(float), cudaMemcpyDeviceToHost, stream)); + cudaStreamSynchronize(stream); +} + +void serialize_engine(unsigned int max_batchsize, float& gd, float& gw, std::string& wts_name, std::string& engine_name) { + // Create builder + IBuilder* builder = createInferBuilder(gLogger); + IBuilderConfig* config = builder->createBuilderConfig(); + + // Create model to populate the network, then set the outputs and create an engine + ICudaEngine *engine = nullptr; + + engine = build_cls_engine(max_batchsize, builder, config, DataType::kFLOAT, gd, gw, wts_name); + + assert(engine != nullptr); + + // Serialize the engine + IHostMemory* serialized_engine = engine->serialize(); + assert(serialized_engine != nullptr); + + // Save engine to file + std::ofstream p(engine_name, std::ios::binary); + if (!p) { + std::cerr << "Could not open plan output file" << std::endl; + assert(false); + } + p.write(reinterpret_cast(serialized_engine->data()), serialized_engine->size()); + + // Close everything down + engine->destroy(); + builder->destroy(); + config->destroy(); + serialized_engine->destroy(); +} + +void deserialize_engine(std::string& engine_name, IRuntime** runtime, ICudaEngine** engine, IExecutionContext** context) { + std::ifstream file(engine_name, std::ios::binary); + if (!file.good()) { + std::cerr << "read " << engine_name << " error!" << std::endl; + assert(false); + } + size_t size = 0; + file.seekg(0, file.end); + size = file.tellg(); + file.seekg(0, file.beg); + char* serialized_engine = new char[size]; + assert(serialized_engine); + file.read(serialized_engine, size); + file.close(); + + *runtime = createInferRuntime(gLogger); + assert(*runtime); + *engine = (*runtime)->deserializeCudaEngine(serialized_engine, size); + assert(*engine); + *context = (*engine)->createExecutionContext(); + assert(*context); + delete[] serialized_engine; } int main(int argc, char** argv) { - cudaSetDevice(DEVICE); + cudaSetDevice(kGpuId); - std::string wts_name = ""; - std::string engine_name = ""; - float gd = 0.0f, gw = 0.0f; - std::string img_dir; - if (!parse_args(argc, argv, wts_name, engine_name, gd, gw, img_dir)) { - std::cerr << "arguments not right!" << std::endl; - std::cerr << "./yolov5_cls -s [.wts] [.engine] [n/s/m/l/x or c gd gw] // serialize model to plan file" << std::endl; - std::cerr << "./yolov5_cls -d [.engine] ../samples // deserialize plan file and run inference" << std::endl; - return -1; - } + std::string wts_name = ""; + std::string engine_name = ""; + float gd = 0.0f, gw = 0.0f; + std::string img_dir; - // create a model using the API directly and serialize it to a stream - if (!wts_name.empty()) { - IHostMemory* modelStream{ nullptr }; - APIToModel(BATCH_SIZE, &modelStream, gd, gw, wts_name); - assert(modelStream != nullptr); - std::ofstream p(engine_name, std::ios::binary); - if (!p) { - std::cerr << "could not open plan output file" << std::endl; - return -1; - } - p.write(reinterpret_cast(modelStream->data()), modelStream->size()); - modelStream->destroy(); - return 0; - } - - // deserialize the .engine and run inference - std::ifstream file(engine_name, std::ios::binary); - if (!file.good()) { - std::cerr << "read " << engine_name << " error!" << std::endl; - return -1; - } - char *trtModelStream = nullptr; - size_t size = 0; - file.seekg(0, file.end); - size = file.tellg(); - file.seekg(0, file.beg); - trtModelStream = new char[size]; - assert(trtModelStream); - file.read(trtModelStream, size); - file.close(); - - std::vector file_names; - if (read_files_in_dir(img_dir.c_str(), file_names) < 0) { - std::cerr << "read_files_in_dir failed." << std::endl; - return -1; - } - auto classes = read_classes("imagenet_classes.txt"); - - static float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W]; - static float prob[BATCH_SIZE * OUTPUT_SIZE]; - IRuntime* runtime = createInferRuntime(gLogger); - assert(runtime != nullptr); - ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size); - assert(engine != nullptr); - IExecutionContext* context = engine->createExecutionContext(); - assert(context != nullptr); - delete[] trtModelStream; - assert(engine->getNbBindings() == 2); - void* buffers[2]; - // In order to bind the buffers, we need to know the names of the input and output tensors. - // Note that indices are guaranteed to be less than IEngine::getNbBindings() - const int inputIndex = engine->getBindingIndex(INPUT_BLOB_NAME); - const int outputIndex = engine->getBindingIndex(OUTPUT_BLOB_NAME); - assert(inputIndex == 0); - assert(outputIndex == 1); - // Create GPU buffers on device - CUDA_CHECK(cudaMalloc((void**)&buffers[inputIndex], BATCH_SIZE * 3 * INPUT_H * INPUT_W * sizeof(float))); - CUDA_CHECK(cudaMalloc((void**)&buffers[outputIndex], BATCH_SIZE * OUTPUT_SIZE * sizeof(float))); - - // Create stream - cudaStream_t stream; - CUDA_CHECK(cudaStreamCreate(&stream)); - - int fcount = 0; - for (int f = 0; f < (int)file_names.size(); f++) { - fcount++; - if (fcount < BATCH_SIZE && f + 1 != (int)file_names.size()) continue; - for (int b = 0; b < fcount; b++) { - cv::Mat img = cv::imread(img_dir + "/" + file_names[f - fcount + 1 + b]); - if (img.empty()) continue; - cv::Mat pr_img; - cv::resize(img, pr_img, cv::Size(INPUT_W, INPUT_H)); - int i = 0; - for (int row = 0; row < INPUT_H; ++row) { - uchar* uc_pixel = pr_img.data + row * pr_img.step; - for (int col = 0; col < INPUT_W; ++col) { - data[b * 3 * INPUT_H * INPUT_W + i] = ((float)uc_pixel[2] / 255.0 - 0.485) / 0.229; // R - 0.485 - data[b * 3 * INPUT_H * INPUT_W + i + INPUT_H * INPUT_W] = ((float)uc_pixel[1] / 255.0 - 0.456) / 0.224; - data[b * 3 * INPUT_H * INPUT_W + i + 2 * INPUT_H * INPUT_W] = ((float)uc_pixel[0] / 255.0 - 0.406) / 0.225; - uc_pixel += 3; - ++i; - } - } - } - // Run inference - auto start = std::chrono::system_clock::now(); - doInference(*context, stream, buffers, data, prob, BATCH_SIZE); - auto end = std::chrono::system_clock::now(); - std::cout << "inference time: " << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; - for (int b = 0; b < fcount; b++) { - float *p = &prob[b * OUTPUT_SIZE]; - auto res = softmax(p, OUTPUT_SIZE); - auto topk_idx = topk(res, 3); - std::cout << file_names[f - fcount + 1 + b] << std::endl; - for (auto idx: topk_idx) { - std::cout << " " << classes[idx] << " " << res[idx] << std::endl; - } - } - - fcount = 0; - } - - // Release stream and buffers - cudaStreamDestroy(stream); - CUDA_CHECK(cudaFree(buffers[inputIndex])); - CUDA_CHECK(cudaFree(buffers[outputIndex])); - // Destroy the engine - context->destroy(); - engine->destroy(); - runtime->destroy(); + if (!parse_args(argc, argv, wts_name, engine_name, gd, gw, img_dir)) { + std::cerr << "arguments not right!" << std::endl; + std::cerr << "./yolov5_cls -s [.wts] [.engine] [n/s/m/l/x or c gd gw] // serialize model to plan file" << std::endl; + std::cerr << "./yolov5_cls -d [.engine] ../images // deserialize plan file and run inference" << std::endl; + return -1; + } + // Create a model using the API directly and serialize it to a file + if (!wts_name.empty()) { + serialize_engine(kBatchSize, gd, gw, wts_name, engine_name); return 0; + } + + // Deserialize the engine from file + IRuntime* runtime = nullptr; + ICudaEngine* engine = nullptr; + IExecutionContext* context = nullptr; + deserialize_engine(engine_name, &runtime, &engine, &context); + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreate(&stream)); + + // Prepare cpu and gpu buffers + float* gpu_buffers[2]; + float* cpu_input_buffer = nullptr; + float* cpu_output_buffer = nullptr; + prepare_buffers(engine, &gpu_buffers[0], &gpu_buffers[1], &cpu_input_buffer, &cpu_output_buffer); + + // Read images from directory + std::vector file_names; + if (read_files_in_dir(img_dir.c_str(), file_names) < 0) { + std::cerr << "read_files_in_dir failed." << std::endl; + return -1; + } + + // Read imagenet labels + auto classes = read_classes("imagenet_classes.txt"); + + // batch predict + for (size_t i = 0; i < file_names.size(); i += kBatchSize) { + // Get a batch of images + std::vector img_batch; + std::vector img_name_batch; + for (size_t j = i; j < i + kBatchSize && j < file_names.size(); j++) { + cv::Mat img = cv::imread(img_dir + "/" + file_names[j]); + img_batch.push_back(img); + img_name_batch.push_back(file_names[j]); + } + + // Preprocess + batch_preprocess(img_batch, cpu_input_buffer); + + // Run inference + auto start = std::chrono::system_clock::now(); + infer(*context, stream, (void**)gpu_buffers, cpu_input_buffer, cpu_output_buffer, kBatchSize); + auto end = std::chrono::system_clock::now(); + std::cout << "inference time: " << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; + + // Postprocess and get top-k result + for (size_t b = 0; b < img_name_batch.size(); b++) { + float* p = &cpu_output_buffer[b * kOutputSize]; + auto res = softmax(p, kOutputSize); + auto topk_idx = topk(res, 3); + std::cout << img_name_batch[b] << std::endl; + for (auto idx: topk_idx) { + std::cout << " " << classes[idx] << " " << res[idx] << std::endl; + } + } + } + + // Release stream and buffers + cudaStreamDestroy(stream); + CUDA_CHECK(cudaFree(gpu_buffers[0])); + CUDA_CHECK(cudaFree(gpu_buffers[1])); + delete[] cpu_input_buffer; + delete[] cpu_output_buffer; + // Destroy the engine + context->destroy(); + engine->destroy(); + runtime->destroy(); + + return 0; } diff --git a/yolov5/yolov5_cls_trt.py b/yolov5/yolov5_cls_trt.py index c2ff932..bb04020 100644 --- a/yolov5/yolov5_cls_trt.py +++ b/yolov5/yolov5_cls_trt.py @@ -229,7 +229,7 @@ if __name__ == "__main__": try: print('batch size is', yolov5_wrapper.batch_size) - image_dir = "samples/" + image_dir = "images/" image_path_batches = get_img_path_batches( yolov5_wrapper.batch_size, image_dir) diff --git a/yolov5/yolov5_det.cpp b/yolov5/yolov5_det.cpp index 9ad6ed5..f073326 100644 --- a/yolov5/yolov5_det.cpp +++ b/yolov5/yolov5_det.cpp @@ -1,231 +1,18 @@ #include "cuda_utils.h" #include "logging.h" -#include "common.hpp" #include "utils.h" -#include "calibrator.h" #include "preprocess.h" +#include "postprocess.h" +#include "model.h" #include #include #include -#define USE_FP16 // set USE_INT8 or USE_FP16 or USE_FP32 -#define DEVICE 0 // GPU id -#define NMS_THRESH 0.4 -#define CONF_THRESH 0.5 -#define MAX_IMAGE_INPUT_SIZE_THRESH 3000 * 3000 // ensure it exceed the maximum size in the input images ! +using namespace nvinfer1; -// stuff we know about the network and the input/output blobs -static const int kBatchSize = 1; -static const int kInputH = Yolo::INPUT_H; -static const int kInputW = Yolo::INPUT_W; -static const int kNumClass = Yolo::CLASS_NUM; -static const int kOutputSize = Yolo::MAX_OUTPUT_BBOX_COUNT * sizeof(Yolo::Detection) / sizeof(float) + 1; // we assume the yololayer outputs no more than MAX_OUTPUT_BBOX_COUNT boxes that conf >= 0.1 -const char* kInputTensorName = "data"; -const char* kOutputTensorName = "prob"; static Logger gLogger; - -static int get_width(int x, float gw, int divisor = 8) { - return int(ceil((x * gw) / divisor)) * divisor; -} - -static int get_depth(int x, float gd) { - if (x == 1) return 1; - int r = round(x * gd); - if (x * gd - int(x * gd) == 0.5 && (int(x * gd) % 2) == 0) { - --r; - } - return std::max(r, 1); -} - -static ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) { - INetworkDefinition* network = builder->createNetworkV2(0U); - - // Create input tensor of shape {3, kInputH, kInputW} with name kInputTensorName - ITensor* data = network->addInput(kInputTensorName, dt, Dims3{ 3, kInputH, kInputW }); - assert(data); - std::map weightMap = loadWeights(wts_name); - /* ------ yolov5 backbone------ */ - auto conv0 = convBlock(network, weightMap, *data, get_width(64, gw), 6, 2, 1, "model.0"); - assert(conv0); - auto conv1 = convBlock(network, weightMap, *conv0->getOutput(0), get_width(128, gw), 3, 2, 1, "model.1"); - auto bottleneck_CSP2 = C3(network, weightMap, *conv1->getOutput(0), get_width(128, gw), get_width(128, gw), get_depth(3, gd), true, 1, 0.5, "model.2"); - auto conv3 = convBlock(network, weightMap, *bottleneck_CSP2->getOutput(0), get_width(256, gw), 3, 2, 1, "model.3"); - auto bottleneck_csp4 = C3(network, weightMap, *conv3->getOutput(0), get_width(256, gw), get_width(256, gw), get_depth(6, gd), true, 1, 0.5, "model.4"); - auto conv5 = convBlock(network, weightMap, *bottleneck_csp4->getOutput(0), get_width(512, gw), 3, 2, 1, "model.5"); - auto bottleneck_csp6 = C3(network, weightMap, *conv5->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(9, gd), true, 1, 0.5, "model.6"); - auto conv7 = convBlock(network, weightMap, *bottleneck_csp6->getOutput(0), get_width(1024, gw), 3, 2, 1, "model.7"); - auto bottleneck_csp8 = C3(network, weightMap, *conv7->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), true, 1, 0.5, "model.8"); - auto spp9 = SPPF(network, weightMap, *bottleneck_csp8->getOutput(0), get_width(1024, gw), get_width(1024, gw), 5, "model.9"); - /* ------ yolov5 head ------ */ - auto conv10 = convBlock(network, weightMap, *spp9->getOutput(0), get_width(512, gw), 1, 1, 1, "model.10"); - - auto upsample11 = network->addResize(*conv10->getOutput(0)); - assert(upsample11); - upsample11->setResizeMode(ResizeMode::kNEAREST); - upsample11->setOutputDimensions(bottleneck_csp6->getOutput(0)->getDimensions()); - - ITensor* inputTensors12[] = { upsample11->getOutput(0), bottleneck_csp6->getOutput(0) }; - auto cat12 = network->addConcatenation(inputTensors12, 2); - auto bottleneck_csp13 = C3(network, weightMap, *cat12->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.13"); - auto conv14 = convBlock(network, weightMap, *bottleneck_csp13->getOutput(0), get_width(256, gw), 1, 1, 1, "model.14"); - - auto upsample15 = network->addResize(*conv14->getOutput(0)); - assert(upsample15); - upsample15->setResizeMode(ResizeMode::kNEAREST); - upsample15->setOutputDimensions(bottleneck_csp4->getOutput(0)->getDimensions()); - - ITensor* inputTensors16[] = { upsample15->getOutput(0), bottleneck_csp4->getOutput(0) }; - auto cat16 = network->addConcatenation(inputTensors16, 2); - - auto bottleneck_csp17 = C3(network, weightMap, *cat16->getOutput(0), get_width(512, gw), get_width(256, gw), get_depth(3, gd), false, 1, 0.5, "model.17"); - - /* ------ detect ------ */ - IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (kNumClass + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]); - auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), get_width(256, gw), 3, 2, 1, "model.18"); - ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) }; - auto cat19 = network->addConcatenation(inputTensors19, 2); - auto bottleneck_csp20 = C3(network, weightMap, *cat19->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.20"); - IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (kNumClass + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]); - auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), get_width(512, gw), 3, 2, 1, "model.21"); - ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) }; - auto cat22 = network->addConcatenation(inputTensors22, 2); - auto bottleneck_csp23 = C3(network, weightMap, *cat22->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.23"); - IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (kNumClass + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]); - - auto yolo = addYoLoLayer(network, weightMap, "model.24", std::vector{det0, det1, det2}); - yolo->getOutput(0)->setName(kOutputTensorName); - network->markOutput(*yolo->getOutput(0)); - // Build engine - builder->setMaxBatchSize(maxBatchSize); - config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB -#if defined(USE_FP16) - config->setFlag(BuilderFlag::kFP16); -#elif defined(USE_INT8) - std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; - assert(builder->platformHasFastInt8()); - config->setFlag(BuilderFlag::kINT8); - Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, kInputW, kInputH, "./coco_calib/", "int8calib.table", kInputTensorName); - config->setInt8Calibrator(calibrator); -#endif - - std::cout << "Building engine, please wait for a while..." << std::endl; - ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); - std::cout << "Build engine successfully!" << std::endl; - - // Don't need the network any more - network->destroy(); - - // Release host memory - for (auto& mem : weightMap) { - free((void*)(mem.second.values)); - } - - return engine; -} - -static ICudaEngine* build_engine_p6(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) { - INetworkDefinition* network = builder->createNetworkV2(0U); - // Create input tensor of shape {3, kInputH, kInputW} with name kInputTensorName - ITensor* data = network->addInput(kInputTensorName, dt, Dims3{ 3, kInputH, kInputW }); - assert(data); - - std::map weightMap = loadWeights(wts_name); - - /* ------ yolov5 backbone------ */ - auto conv0 = convBlock(network, weightMap, *data, get_width(64, gw), 6, 2, 1, "model.0"); - auto conv1 = convBlock(network, weightMap, *conv0->getOutput(0), get_width(128, gw), 3, 2, 1, "model.1"); - auto c3_2 = C3(network, weightMap, *conv1->getOutput(0), get_width(128, gw), get_width(128, gw), get_depth(3, gd), true, 1, 0.5, "model.2"); - auto conv3 = convBlock(network, weightMap, *c3_2->getOutput(0), get_width(256, gw), 3, 2, 1, "model.3"); - auto c3_4 = C3(network, weightMap, *conv3->getOutput(0), get_width(256, gw), get_width(256, gw), get_depth(6, gd), true, 1, 0.5, "model.4"); - auto conv5 = convBlock(network, weightMap, *c3_4->getOutput(0), get_width(512, gw), 3, 2, 1, "model.5"); - auto c3_6 = C3(network, weightMap, *conv5->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(9, gd), true, 1, 0.5, "model.6"); - auto conv7 = convBlock(network, weightMap, *c3_6->getOutput(0), get_width(768, gw), 3, 2, 1, "model.7"); - auto c3_8 = C3(network, weightMap, *conv7->getOutput(0), get_width(768, gw), get_width(768, gw), get_depth(3, gd), true, 1, 0.5, "model.8"); - auto conv9 = convBlock(network, weightMap, *c3_8->getOutput(0), get_width(1024, gw), 3, 2, 1, "model.9"); - auto c3_10 = C3(network, weightMap, *conv9->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), true, 1, 0.5, "model.10"); - auto sppf11 = SPPF(network, weightMap, *c3_10->getOutput(0), get_width(1024, gw), get_width(1024, gw), 5, "model.11"); - - /* ------ yolov5 head ------ */ - auto conv12 = convBlock(network, weightMap, *sppf11->getOutput(0), get_width(768, gw), 1, 1, 1, "model.12"); - auto upsample13 = network->addResize(*conv12->getOutput(0)); - assert(upsample13); - upsample13->setResizeMode(ResizeMode::kNEAREST); - upsample13->setOutputDimensions(c3_8->getOutput(0)->getDimensions()); - ITensor* inputTensors14[] = { upsample13->getOutput(0), c3_8->getOutput(0) }; - auto cat14 = network->addConcatenation(inputTensors14, 2); - auto c3_15 = C3(network, weightMap, *cat14->getOutput(0), get_width(1536, gw), get_width(768, gw), get_depth(3, gd), false, 1, 0.5, "model.15"); - - auto conv16 = convBlock(network, weightMap, *c3_15->getOutput(0), get_width(512, gw), 1, 1, 1, "model.16"); - auto upsample17 = network->addResize(*conv16->getOutput(0)); - assert(upsample17); - upsample17->setResizeMode(ResizeMode::kNEAREST); - upsample17->setOutputDimensions(c3_6->getOutput(0)->getDimensions()); - ITensor* inputTensors18[] = { upsample17->getOutput(0), c3_6->getOutput(0) }; - auto cat18 = network->addConcatenation(inputTensors18, 2); - auto c3_19 = C3(network, weightMap, *cat18->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.19"); - - auto conv20 = convBlock(network, weightMap, *c3_19->getOutput(0), get_width(256, gw), 1, 1, 1, "model.20"); - auto upsample21 = network->addResize(*conv20->getOutput(0)); - assert(upsample21); - upsample21->setResizeMode(ResizeMode::kNEAREST); - upsample21->setOutputDimensions(c3_4->getOutput(0)->getDimensions()); - ITensor* inputTensors21[] = { upsample21->getOutput(0), c3_4->getOutput(0) }; - auto cat22 = network->addConcatenation(inputTensors21, 2); - auto c3_23 = C3(network, weightMap, *cat22->getOutput(0), get_width(512, gw), get_width(256, gw), get_depth(3, gd), false, 1, 0.5, "model.23"); - - auto conv24 = convBlock(network, weightMap, *c3_23->getOutput(0), get_width(256, gw), 3, 2, 1, "model.24"); - ITensor* inputTensors25[] = { conv24->getOutput(0), conv20->getOutput(0) }; - auto cat25 = network->addConcatenation(inputTensors25, 2); - auto c3_26 = C3(network, weightMap, *cat25->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.26"); - - auto conv27 = convBlock(network, weightMap, *c3_26->getOutput(0), get_width(512, gw), 3, 2, 1, "model.27"); - ITensor* inputTensors28[] = { conv27->getOutput(0), conv16->getOutput(0) }; - auto cat28 = network->addConcatenation(inputTensors28, 2); - auto c3_29 = C3(network, weightMap, *cat28->getOutput(0), get_width(1536, gw), get_width(768, gw), get_depth(3, gd), false, 1, 0.5, "model.29"); - - auto conv30 = convBlock(network, weightMap, *c3_29->getOutput(0), get_width(768, gw), 3, 2, 1, "model.30"); - ITensor* inputTensors31[] = { conv30->getOutput(0), conv12->getOutput(0) }; - auto cat31 = network->addConcatenation(inputTensors31, 2); - auto c3_32 = C3(network, weightMap, *cat31->getOutput(0), get_width(2048, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.32"); - - /* ------ detect ------ */ - IConvolutionLayer* det0 = network->addConvolutionNd(*c3_23->getOutput(0), 3 * (kNumClass + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.0.weight"], weightMap["model.33.m.0.bias"]); - IConvolutionLayer* det1 = network->addConvolutionNd(*c3_26->getOutput(0), 3 * (kNumClass + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.1.weight"], weightMap["model.33.m.1.bias"]); - IConvolutionLayer* det2 = network->addConvolutionNd(*c3_29->getOutput(0), 3 * (kNumClass + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.2.weight"], weightMap["model.33.m.2.bias"]); - IConvolutionLayer* det3 = network->addConvolutionNd(*c3_32->getOutput(0), 3 * (kNumClass + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.3.weight"], weightMap["model.33.m.3.bias"]); - - auto yolo = addYoLoLayer(network, weightMap, "model.33", std::vector{det0, det1, det2, det3}); - yolo->getOutput(0)->setName(kOutputTensorName); - network->markOutput(*yolo->getOutput(0)); - - // Build engine - builder->setMaxBatchSize(maxBatchSize); - config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB -#if defined(USE_FP16) - config->setFlag(BuilderFlag::kFP16); -#elif defined(USE_INT8) - std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; - assert(builder->platformHasFastInt8()); - config->setFlag(BuilderFlag::kINT8); - Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, kInputW, kInputH, "./coco_calib/", "int8calib.table", kInputTensorName); - config->setInt8Calibrator(calibrator); -#endif - - std::cout << "Building engine, please wait for a while..." << std::endl; - ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); - std::cout << "Build engine successfully!" << std::endl; - - // Don't need the network any more - network->destroy(); - - // Release host memory - for (auto& mem : weightMap) { - free((void*)(mem.second.values)); - } - - return engine; -} +const static int kOutputSize = kMaxNumOutputBbox * sizeof(Detection) / sizeof(float) + 1; bool parse_args(int argc, char** argv, std::string& wts, std::string& engine, bool& is_p6, float& gd, float& gw, std::string& img_dir) { if (argc < 4) return false; @@ -282,9 +69,9 @@ void prepare_buffers(ICudaEngine* engine, float** gpu_input_buffer, float** gpu_ } void infer(IExecutionContext& context, cudaStream_t& stream, void** gpu_buffers, float* output, int batchsize) { - context.enqueue(batchsize, gpu_buffers, stream, nullptr); - CUDA_CHECK(cudaMemcpyAsync(output, gpu_buffers[1], batchsize * kOutputSize * sizeof(float), cudaMemcpyDeviceToHost, stream)); - cudaStreamSynchronize(stream); + context.enqueue(batchsize, gpu_buffers, stream, nullptr); + CUDA_CHECK(cudaMemcpyAsync(output, gpu_buffers[1], batchsize * kOutputSize * sizeof(float), cudaMemcpyDeviceToHost, stream)); + cudaStreamSynchronize(stream); } void serialize_engine(unsigned int max_batchsize, bool& is_p6, float& gd, float& gw, std::string& wts_name, std::string& engine_name) { @@ -295,9 +82,9 @@ void serialize_engine(unsigned int max_batchsize, bool& is_p6, float& gd, float& // Create model to populate the network, then set the outputs and create an engine ICudaEngine *engine = nullptr; if (is_p6) { - engine = build_engine_p6(max_batchsize, builder, config, DataType::kFLOAT, gd, gw, wts_name); + engine = build_det_p6_engine(max_batchsize, builder, config, DataType::kFLOAT, gd, gw, wts_name); } else { - engine = build_engine(max_batchsize, builder, config, DataType::kFLOAT, gd, gw, wts_name); + engine = build_det_engine(max_batchsize, builder, config, DataType::kFLOAT, gd, gw, wts_name); } assert(engine != nullptr); @@ -345,7 +132,7 @@ void deserialize_engine(std::string& engine_name, IRuntime** runtime, ICudaEngin } int main(int argc, char** argv) { - cudaSetDevice(DEVICE); + cudaSetDevice(kGpuId); std::string wts_name = ""; std::string engine_name = ""; @@ -356,7 +143,7 @@ int main(int argc, char** argv) { if (!parse_args(argc, argv, wts_name, engine_name, is_p6, gd, gw, img_dir)) { std::cerr << "arguments not right!" << std::endl; std::cerr << "./yolov5_det -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw] // serialize model to plan file" << std::endl; - std::cerr << "./yolov5_det -d [.engine] ../samples // deserialize plan file and run inference" << std::endl; + std::cerr << "./yolov5_det -d [.engine] ../images // deserialize plan file and run inference" << std::endl; return -1; } @@ -366,7 +153,6 @@ int main(int argc, char** argv) { return 0; } - // Deserialize the engine from file IRuntime* runtime = nullptr; ICudaEngine* engine = nullptr; @@ -375,89 +161,72 @@ int main(int argc, char** argv) { cudaStream_t stream; CUDA_CHECK(cudaStreamCreate(&stream)); - std::vector file_names; - if (read_files_in_dir(img_dir.c_str(), file_names) < 0) { - std::cerr << "read_files_in_dir failed." << std::endl; - return -1; - } + // Init CUDA preprocessing + cuda_preprocess_init(kMaxInputImageSize); // Prepare cpu and gpu buffers float* gpu_buffers[2]; float* cpu_output_buffer = nullptr; prepare_buffers(engine, &gpu_buffers[0], &gpu_buffers[1], &cpu_output_buffer); - uint8_t* img_host = nullptr; - uint8_t* img_device = nullptr; - // prepare input data cache in pinned memory - CUDA_CHECK(cudaMallocHost((void**)&img_host, MAX_IMAGE_INPUT_SIZE_THRESH * 3)); - // prepare input data cache in device memory - CUDA_CHECK(cudaMalloc((void**)&img_device, MAX_IMAGE_INPUT_SIZE_THRESH * 3)); + // Read images from directory + std::vector file_names; + if (read_files_in_dir(img_dir.c_str(), file_names) < 0) { + std::cerr << "read_files_in_dir failed." << std::endl; + return -1; + } - int fcount = 0; - std::vector imgs_buffer(kBatchSize); - for (int f = 0; f < (int)file_names.size(); f++) { - fcount++; - if (fcount < kBatchSize && f + 1 != (int)file_names.size()) continue; - //auto start = std::chrono::system_clock::now(); - float *buffer_idx = (float*)gpu_buffers[0]; - for (int b = 0; b < fcount; b++) { - cv::Mat img = cv::imread(img_dir + "/" + file_names[f - fcount + 1 + b]); - if (img.empty()) continue; - imgs_buffer[b] = img; - size_t size_image = img.cols * img.rows * 3; - size_t size_image_dst = kInputH * kInputW * 3; - //copy data to pinned memory - memcpy(img_host, img.data, size_image); - //copy data to device memory - CUDA_CHECK(cudaMemcpyAsync(img_device, img_host, size_image, cudaMemcpyHostToDevice, stream)); - preprocess_kernel_img(img_device, img.cols, img.rows, buffer_idx, kInputW, kInputH, stream); - buffer_idx += size_image_dst; - cudaStreamSynchronize(stream); + // batch predict + for (size_t i = 0; i < file_names.size(); i += kBatchSize) { + // Get a batch of images + std::vector img_batch; + std::vector img_name_batch; + for (size_t j = i; j < i + kBatchSize && j < file_names.size(); j++) { + cv::Mat img = cv::imread(img_dir + "/" + file_names[j]); + img_batch.push_back(img); + img_name_batch.push_back(file_names[j]); } + + // Preprocess + cuda_batch_preprocess(img_batch, gpu_buffers[0], kInputW, kInputH, stream); + // Run inference auto start = std::chrono::system_clock::now(); infer(*context, stream, (void**)gpu_buffers, cpu_output_buffer, kBatchSize); auto end = std::chrono::system_clock::now(); std::cout << "inference time: " << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; - std::vector> batch_res(fcount); - for (int b = 0; b < fcount; b++) { - auto& res = batch_res[b]; - nms(res, &cpu_output_buffer[b * kOutputSize], CONF_THRESH, NMS_THRESH); + + // NMS + std::vector> res_batch; + batch_nms(res_batch, cpu_output_buffer, img_batch.size(), kOutputSize, kConfThresh, kNmsThresh); + + // Draw bounding boxes + draw_bbox(img_batch, res_batch); + + // Save images + for (size_t j = 0; j < img_batch.size(); j++) { + cv::imwrite("_" + img_name_batch[j], img_batch[j]); } - for (int b = 0; b < fcount; b++) { - auto& res = batch_res[b]; - cv::Mat img = imgs_buffer[b]; - for (size_t j = 0; j < res.size(); j++) { - cv::Rect r = get_rect(img, res[j].bbox); - cv::rectangle(img, r, cv::Scalar(0x27, 0xC1, 0x36), 2); - cv::putText(img, std::to_string((int)res[j].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2); - } - cv::imwrite("_" + file_names[f - fcount + 1 + b], img); - } - fcount = 0; } // Release stream and buffers cudaStreamDestroy(stream); - CUDA_CHECK(cudaFree(img_device)); - CUDA_CHECK(cudaFreeHost(img_host)); CUDA_CHECK(cudaFree(gpu_buffers[0])); CUDA_CHECK(cudaFree(gpu_buffers[1])); delete[] cpu_output_buffer; + cuda_preprocess_destroy(); // Destroy the engine context->destroy(); engine->destroy(); runtime->destroy(); - // Print histogram of the output distribution - //std::cout << "\nOutput:\n\n"; - //for (unsigned int i = 0; i < kOutputSize; i++) - //{ - // std::cout << prob[i] << ", "; - // if (i % 10 == 0) std::cout << std::endl; - //} - //std::cout << std::endl; + // std::cout << "\nOutput:\n\n"; + // for (unsigned int i = 0; i < kOutputSize; i++) { + // std::cout << prob[i] << ", "; + // if (i % 10 == 0) std::cout << std::endl; + // } + // std::cout << std::endl; return 0; } diff --git a/yolov5/yolov5_det_cuda_python.py b/yolov5/yolov5_det_cuda_python.py index c2c4487..5db3c1d 100644 --- a/yolov5/yolov5_det_cuda_python.py +++ b/yolov5/yolov5_det_cuda_python.py @@ -435,7 +435,7 @@ if __name__ == "__main__": try: print('batch size is', yolov5_wrapper.batch_size) - image_dir = "samples/" + image_dir = "images/" image_path_batches = get_img_path_batches(yolov5_wrapper.batch_size, image_dir) for i in range(10): diff --git a/yolov5/yolov5_det_trt.py b/yolov5/yolov5_det_trt.py index bed0ba4..91cd16d 100644 --- a/yolov5/yolov5_det_trt.py +++ b/yolov5/yolov5_det_trt.py @@ -434,7 +434,7 @@ if __name__ == "__main__": try: print('batch size is', yolov5_wrapper.batch_size) - image_dir = "samples/" + image_dir = "images/" image_path_batches = get_img_path_batches(yolov5_wrapper.batch_size, image_dir) for i in range(10): diff --git a/yolov5/yolov5_seg.cpp b/yolov5/yolov5_seg.cpp index bfe424f..f3a60c4 100644 --- a/yolov5/yolov5_seg.cpp +++ b/yolov5/yolov5_seg.cpp @@ -1,159 +1,20 @@ +#include "config.h" +#include "cuda_utils.h" +#include "logging.h" +#include "utils.h" +#include "preprocess.h" +#include "postprocess.h" +#include "model.h" + #include #include #include -#include "cuda_utils.h" -#include "logging.h" -#include "common.hpp" -#include "utils.h" -#include "calibrator.h" -#include "preprocess.h" -#define USE_FP32 // set USE_INT8 or USE_FP16 or USE_FP32 -#define DEVICE 0 // GPU id -#define NMS_THRESH 0.4 -#define CONF_THRESH 0.5 -#define BATCH_SIZE 1 -#define MAX_IMAGE_INPUT_SIZE_THRESH 3000 * 3000 // ensure it exceed the maximum size in the input images ! +using namespace nvinfer1; -// stuff we know about the network and the input/output blobs -static const int INPUT_H = Yolo::INPUT_H; -static const int INPUT_W = Yolo::INPUT_W; -static const int CLASS_NUM = Yolo::CLASS_NUM; -static const int OUTPUT_SIZE1 = Yolo::MAX_OUTPUT_BBOX_COUNT * sizeof(Yolo::Detection) / sizeof(float) + 1; // we assume the yololayer outputs no more than MAX_OUTPUT_BBOX_COUNT boxes that conf >= 0.1 -static const int OUTPUT_SIZE2 = 32 * (INPUT_H / 4) * (INPUT_W / 4); -const char* INPUT_BLOB_NAME = "data"; -const char* OUTPUT_BLOB_NAME = "prob"; static Logger gLogger; - -static int get_width(int x, float gw, int divisor = 8) { - return int(ceil((x * gw) / divisor)) * divisor; -} - -static int get_depth(int x, float gd) { - if (x == 1) return 1; - int r = round(x * gd); - if (x * gd - int(x * gd) == 0.5 && (int(x * gd) % 2) == 0) { - --r; - } - return std::max(r, 1); -} - -ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) { - INetworkDefinition* network = builder->createNetworkV2(0U); - ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{ 3, INPUT_H, INPUT_W }); - assert(data); - std::map weightMap = loadWeights(wts_name); - - // Backbone - auto conv0 = convBlock(network, weightMap, *data, get_width(64, gw), 6, 2, 1, "model.0"); - assert(conv0); - auto conv1 = convBlock(network, weightMap, *conv0->getOutput(0), get_width(128, gw), 3, 2, 1, "model.1"); - auto bottleneck_CSP2 = C3(network, weightMap, *conv1->getOutput(0), get_width(128, gw), get_width(128, gw), get_depth(3, gd), true, 1, 0.5, "model.2"); - auto conv3 = convBlock(network, weightMap, *bottleneck_CSP2->getOutput(0), get_width(256, gw), 3, 2, 1, "model.3"); - auto bottleneck_csp4 = C3(network, weightMap, *conv3->getOutput(0), get_width(256, gw), get_width(256, gw), get_depth(6, gd), true, 1, 0.5, "model.4"); - auto conv5 = convBlock(network, weightMap, *bottleneck_csp4->getOutput(0), get_width(512, gw), 3, 2, 1, "model.5"); - auto bottleneck_csp6 = C3(network, weightMap, *conv5->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(9, gd), true, 1, 0.5, "model.6"); - auto conv7 = convBlock(network, weightMap, *bottleneck_csp6->getOutput(0), get_width(1024, gw), 3, 2, 1, "model.7"); - auto bottleneck_csp8 = C3(network, weightMap, *conv7->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), true, 1, 0.5, "model.8"); - auto spp9 = SPPF(network, weightMap, *bottleneck_csp8->getOutput(0), get_width(1024, gw), get_width(1024, gw), 5, "model.9"); - - // Head - auto conv10 = convBlock(network, weightMap, *spp9->getOutput(0), get_width(512, gw), 1, 1, 1, "model.10"); - - auto upsample11 = network->addResize(*conv10->getOutput(0)); - assert(upsample11); - upsample11->setResizeMode(ResizeMode::kNEAREST); - upsample11->setOutputDimensions(bottleneck_csp6->getOutput(0)->getDimensions()); - - ITensor* inputTensors12[] = { upsample11->getOutput(0), bottleneck_csp6->getOutput(0) }; - auto cat12 = network->addConcatenation(inputTensors12, 2); - auto bottleneck_csp13 = C3(network, weightMap, *cat12->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.13"); - auto conv14 = convBlock(network, weightMap, *bottleneck_csp13->getOutput(0), get_width(256, gw), 1, 1, 1, "model.14"); - - auto upsample15 = network->addResize(*conv14->getOutput(0)); - assert(upsample15); - upsample15->setResizeMode(ResizeMode::kNEAREST); - upsample15->setOutputDimensions(bottleneck_csp4->getOutput(0)->getDimensions()); - - ITensor* inputTensors16[] = { upsample15->getOutput(0), bottleneck_csp4->getOutput(0) }; - auto cat16 = network->addConcatenation(inputTensors16, 2); - - auto bottleneck_csp17 = C3(network, weightMap, *cat16->getOutput(0), get_width(512, gw), get_width(256, gw), get_depth(3, gd), false, 1, 0.5, "model.17"); - - // Segmentation - IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (32 + Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]); - auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), get_width(256, gw), 3, 2, 1, "model.18"); - ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) }; - auto cat19 = network->addConcatenation(inputTensors19, 2); - auto bottleneck_csp20 = C3(network, weightMap, *cat19->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.20"); - IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (32 + Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]); - auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), get_width(512, gw), 3, 2, 1, "model.21"); - ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) }; - auto cat22 = network->addConcatenation(inputTensors22, 2); - auto bottleneck_csp23 = C3(network, weightMap, *cat22->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.23"); - IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (32 + Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]); - - auto yolo = addYoLoLayer(network, weightMap, "model.24", std::vector{det0, det1, det2}, true); - yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME); - network->markOutput(*yolo->getOutput(0)); - - auto proto = Proto(network, weightMap, *bottleneck_csp17->getOutput(0), get_width(256, gw), 32, "model.24.proto"); - proto->getOutput(0)->setName("proto"); - network->markOutput(*proto->getOutput(0)); - - // Build engine - builder->setMaxBatchSize(maxBatchSize); - config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB -#if defined(USE_FP16) - config->setFlag(BuilderFlag::kFP16); -#elif defined(USE_INT8) - std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; - assert(builder->platformHasFastInt8()); - config->setFlag(BuilderFlag::kINT8); - Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, INPUT_W, INPUT_H, "./coco_calib/", "int8calib.table", INPUT_BLOB_NAME); - config->setInt8Calibrator(calibrator); -#endif - - std::cout << "Building engine, please wait for a while..." << std::endl; - ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); - std::cout << "Build engine successfully!" << std::endl; - - // Don't need the network any more - network->destroy(); - - // Release host memory - for (auto& mem : weightMap) { - free((void*)(mem.second.values)); - } - - return engine; -} - -void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream, float& gd, float& gw, std::string& wts_name) { - // Create builder - IBuilder* builder = createInferBuilder(gLogger); - IBuilderConfig* config = builder->createBuilderConfig(); - - // Create model to populate the network, then set the outputs and create an engine - ICudaEngine* engine = build_engine(maxBatchSize, builder, config, DataType::kFLOAT, gd, gw, wts_name); - assert(engine != nullptr); - - // Serialize the engine - (*modelStream) = engine->serialize(); - - // Close everything down - engine->destroy(); - builder->destroy(); - config->destroy(); -} - -void doInference(IExecutionContext& context, cudaStream_t& stream, void **buffers, float* output1, float* output2, int batchSize) { - // infer on the batch asynchronously, and DMA output back to host - context.enqueue(batchSize, buffers, stream, nullptr); - CUDA_CHECK(cudaMemcpyAsync(output1, buffers[1], batchSize * OUTPUT_SIZE1 * sizeof(float), cudaMemcpyDeviceToHost, stream)); - CUDA_CHECK(cudaMemcpyAsync(output2, buffers[2], batchSize * OUTPUT_SIZE2 * sizeof(float), cudaMemcpyDeviceToHost, stream)); - cudaStreamSynchronize(stream); -} +const static int kOutputSize1 = kMaxNumOutputBbox * sizeof(Detection) / sizeof(float) + 1; +const static int kOutputSize2 = 32 * (kInputH / 4) * (kInputW / 4); bool parse_args(int argc, char** argv, std::string& wts, std::string& engine, float& gd, float& gw, std::string& img_dir, std::string& labels_filename) { if (argc < 4) return false; @@ -192,264 +53,193 @@ bool parse_args(int argc, char** argv, std::string& wts, std::string& engine, fl return true; } -cv::Rect get_downscale_rect(float bbox[4], float scale) { - float left = bbox[0] - bbox[2] / 2; - float top = bbox[1] - bbox[3] / 2; - float right = bbox[0] + bbox[2] / 2; - float bottom = bbox[1] + bbox[3] / 2; - left /= scale; - top /= scale; - right /= scale; - bottom /= scale; - return cv::Rect(round(left), round(top), round(right - left), round(bottom - top)); +void prepare_buffers(ICudaEngine* engine, float** gpu_input_buffer, float** gpu_output_buffer1, float** gpu_output_buffer2, float** cpu_output_buffer1, float** cpu_output_buffer2) { + assert(engine->getNbBindings() == 3); + // In order to bind the buffers, we need to know the names of the input and output tensors. + // Note that indices are guaranteed to be less than IEngine::getNbBindings() + const int inputIndex = engine->getBindingIndex(kInputTensorName); + const int outputIndex1 = engine->getBindingIndex(kOutputTensorName); + const int outputIndex2 = engine->getBindingIndex("proto"); + assert(inputIndex == 0); + assert(outputIndex1 == 1); + assert(outputIndex2 == 2); + + // Create GPU buffers on device + CUDA_CHECK(cudaMalloc((void**)gpu_input_buffer, kBatchSize * 3 * kInputH * kInputW * sizeof(float))); + CUDA_CHECK(cudaMalloc((void**)gpu_output_buffer1, kBatchSize * kOutputSize1 * sizeof(float))); + CUDA_CHECK(cudaMalloc((void**)gpu_output_buffer2, kBatchSize * kOutputSize2 * sizeof(float))); + + // Alloc CPU buffers + *cpu_output_buffer1 = new float[kBatchSize * kOutputSize1]; + *cpu_output_buffer2 = new float[kBatchSize * kOutputSize2]; } -std::vector process_mask(const float* proto, std::vector& dets) { - std::vector masks; - for (size_t i = 0; i < dets.size(); i++) { - cv::Mat mask_mat = cv::Mat::zeros(INPUT_H / 4, INPUT_W / 4, CV_32FC1); - auto r = get_downscale_rect(dets[i].bbox, 4); - for (int x = r.x; x < r.x + r.width; x++) { - for (int y = r.y; y < r.y + r.height; y++) { - float e = 0.0f; - for (int j = 0; j < 32; j++) { - e += dets[i].mask[j] * proto[j * OUTPUT_SIZE2 / 32 + y * mask_mat.cols + x]; - } - e = 1.0f / (1.0f + expf(-e)); - mask_mat.at(y, x) = e; - } - } - cv::resize(mask_mat, mask_mat, cv::Size(INPUT_W, INPUT_H)); - masks.push_back(mask_mat); - } - return masks; +void infer(IExecutionContext& context, cudaStream_t& stream, void **buffers, float* output1, float* output2, int batchSize) { + context.enqueue(batchSize, buffers, stream, nullptr); + CUDA_CHECK(cudaMemcpyAsync(output1, buffers[1], batchSize * kOutputSize1 * sizeof(float), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaMemcpyAsync(output2, buffers[2], batchSize * kOutputSize2 * sizeof(float), cudaMemcpyDeviceToHost, stream)); + cudaStreamSynchronize(stream); } -cv::Mat scale_mask(cv::Mat mask, cv::Mat img) { - int x, y, w, h; - float r_w = INPUT_W / (img.cols * 1.0); - float r_h = INPUT_H / (img.rows * 1.0); - if (r_h > r_w) { - w = INPUT_W; - h = r_w * img.rows; - x = 0; - y = (INPUT_H - h) / 2; - } else { - w = r_h * img.cols; - h = INPUT_H; - x = (INPUT_W - w) / 2; - y = 0; +void serialize_engine(unsigned int max_batchsize, float& gd, float& gw, std::string& wts_name, std::string& engine_name) { + // Create builder + IBuilder* builder = createInferBuilder(gLogger); + IBuilderConfig* config = builder->createBuilderConfig(); + + // Create model to populate the network, then set the outputs and create an engine + ICudaEngine *engine = nullptr; + + engine = build_seg_engine(max_batchsize, builder, config, DataType::kFLOAT, gd, gw, wts_name); + + assert(engine != nullptr); + + // Serialize the engine + IHostMemory* serialized_engine = engine->serialize(); + assert(serialized_engine != nullptr); + + // Save engine to file + std::ofstream p(engine_name, std::ios::binary); + if (!p) { + std::cerr << "Could not open plan output file" << std::endl; + assert(false); } - cv::Rect r(x, y, w, h); - cv::Mat res; - cv::resize(mask(r), res, img.size()); - return res; + p.write(reinterpret_cast(serialized_engine->data()), serialized_engine->size()); + + // Close everything down + engine->destroy(); + builder->destroy(); + config->destroy(); + serialized_engine->destroy(); } -void draw_mask_bbox(cv::Mat& img, std::vector& dets, std::vector& masks, std::unordered_map& labels_map) { - static std::vector colors = {0xFF3838, 0xFF9D97, 0xFF701F, 0xFFB21D, 0xCFD231, 0x48F90A, - 0x92CC17, 0x3DDB86, 0x1A9334, 0x00D4BB, 0x2C99A8, 0x00C2FF, - 0x344593, 0x6473FF, 0x0018EC, 0x8438FF, 0x520085, 0xCB38FF, - 0xFF95C8, 0xFF37C7}; - for (size_t i = 0; i < dets.size(); i++) { - cv::Mat img_mask = scale_mask(masks[i], img); - auto color = colors[(int)dets[i].class_id % colors.size()]; - auto bgr = cv::Scalar(color & 0xFF, color >> 8 & 0xFF, color >> 16 & 0xFF); - - cv::Rect r = get_rect(img, dets[i].bbox); - for (int x = r.x; x < r.x + r.width; x++) { - for (int y = r.y; y < r.y + r.height; y++) { - float val = img_mask.at(y, x); - if (val <= 0.5) continue; - img.at(y, x)[0] = img.at(y, x)[0] / 2 + bgr[0] / 2; - img.at(y, x)[1] = img.at(y, x)[1] / 2 + bgr[1] / 2; - img.at(y, x)[2] = img.at(y, x)[2] / 2 + bgr[2] / 2; - } - } - - cv::rectangle(img, r, bgr, 2); - - // Get the size of the text - cv::Size textSize = cv::getTextSize(labels_map[(int)dets[i].class_id] + " " + to_string_with_precision(dets[i].conf), cv::FONT_HERSHEY_PLAIN, 1.2, 2, NULL); - // Set the top left corner of the rectangle - cv::Point topLeft(r.x, r.y - textSize.height); - - // Set the bottom right corner of the rectangle - cv::Point bottomRight(r.x + textSize.width, r.y + textSize.height); - - // Set the thickness of the rectangle lines - int lineThickness = 2; - - // Draw the rectangle on the image - cv::rectangle(img, topLeft, bottomRight, bgr, -1); - - cv::putText(img, labels_map[(int)dets[i].class_id] + " " + to_string_with_precision(dets[i].conf), cv::Point(r.x, r.y + 4), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar::all(0xFF), 2); - +void deserialize_engine(std::string& engine_name, IRuntime** runtime, ICudaEngine** engine, IExecutionContext** context) { + std::ifstream file(engine_name, std::ios::binary); + if (!file.good()) { + std::cerr << "read " << engine_name << " error!" << std::endl; + assert(false); } + size_t size = 0; + file.seekg(0, file.end); + size = file.tellg(); + file.seekg(0, file.beg); + char* serialized_engine = new char[size]; + assert(serialized_engine); + file.read(serialized_engine, size); + file.close(); + + *runtime = createInferRuntime(gLogger); + assert(*runtime); + *engine = (*runtime)->deserializeCudaEngine(serialized_engine, size); + assert(*engine); + *context = (*engine)->createExecutionContext(); + assert(*context); + delete[] serialized_engine; } int main(int argc, char** argv) { - cudaSetDevice(DEVICE); + cudaSetDevice(kGpuId); - std::string wts_name = ""; - std::string engine_name = ""; - std::string labels_filename = ""; - - float gd = 0.0f, gw = 0.0f; - std::string img_dir; - if (!parse_args(argc, argv, wts_name, engine_name, gd, gw, img_dir, labels_filename)) { - std::cerr << "arguments not right!" << std::endl; - std::cerr << "./yolov5_seg -s [.wts] [.engine] [n/s/m/l/x or c gd gw] // serialize model to plan file" << std::endl; - std::cerr << "./yolov5_seg -d [.engine] ../samples coco.txt // deserialize plan file, read the labels file and run inference" << std::endl; - return -1; - } + std::string wts_name = ""; + std::string engine_name = ""; + std::string labels_filename = ""; + float gd = 0.0f, gw = 0.0f; - // create a model using the API directly and serialize it to a stream - if (!wts_name.empty()) { - IHostMemory* modelStream{ nullptr }; - APIToModel(BATCH_SIZE, &modelStream, gd, gw, wts_name); - assert(modelStream != nullptr); - std::ofstream p(engine_name, std::ios::binary); - if (!p) { - std::cerr << "could not open plan output file" << std::endl; - return -1; - } - p.write(reinterpret_cast(modelStream->data()), modelStream->size()); - modelStream->destroy(); - return 0; - } - - // deserialize the .engine and run inference - std::ifstream file(engine_name, std::ios::binary); - if (!file.good()) { - std::cerr << "read " << engine_name << " error!" << std::endl; - return -1; - } - char *trtModelStream = nullptr; - size_t size = 0; - file.seekg(0, file.end); - size = file.tellg(); - file.seekg(0, file.beg); - trtModelStream = new char[size]; - assert(trtModelStream); - file.read(trtModelStream, size); - file.close(); - - std::vector file_names; - if (read_files_in_dir(img_dir.c_str(), file_names) < 0) { - std::cerr << "read_files_in_dir failed." << std::endl; - return -1; - } - - // read the txt file for classnames - std::ifstream labels_file(labels_filename, std::ios::binary); - if (!labels_file.good()) { - std::cerr << "read " << labels_filename << " error!" << std::endl; - return -1; - } - std::unordered_map labels_map; - read_labels(labels_filename, labels_map); - - assert(CLASS_NUM == labels_map.size()); - - - static float prob[BATCH_SIZE * OUTPUT_SIZE1]; - static float proto[BATCH_SIZE * OUTPUT_SIZE2]; - IRuntime* runtime = createInferRuntime(gLogger); - assert(runtime != nullptr); - ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size); - assert(engine != nullptr); - IExecutionContext* context = engine->createExecutionContext(); - assert(context != nullptr); - delete[] trtModelStream; - assert(engine->getNbBindings() == 3); - float* buffers[3]; - // In order to bind the buffers, we need to know the names of the input and output tensors. - // Note that indices are guaranteed to be less than IEngine::getNbBindings() - const int inputIndex = engine->getBindingIndex(INPUT_BLOB_NAME); - const int outputIndex1 = engine->getBindingIndex(OUTPUT_BLOB_NAME); - const int outputIndex2 = engine->getBindingIndex("proto"); - assert(inputIndex == 0); - assert(outputIndex1 == 1); - assert(outputIndex2 == 2); - // Create GPU buffers on device - CUDA_CHECK(cudaMalloc((void**)&buffers[inputIndex], BATCH_SIZE * 3 * INPUT_H * INPUT_W * sizeof(float))); - CUDA_CHECK(cudaMalloc((void**)&buffers[outputIndex1], BATCH_SIZE * OUTPUT_SIZE1 * sizeof(float))); - CUDA_CHECK(cudaMalloc((void**)&buffers[outputIndex2], BATCH_SIZE * OUTPUT_SIZE2 * sizeof(float))); - - // Create stream - cudaStream_t stream; - CUDA_CHECK(cudaStreamCreate(&stream)); - uint8_t* img_host = nullptr; - uint8_t* img_device = nullptr; - // prepare input data cache in pinned memory - CUDA_CHECK(cudaMallocHost((void**)&img_host, MAX_IMAGE_INPUT_SIZE_THRESH * 3)); - // prepare input data cache in device memory - CUDA_CHECK(cudaMalloc((void**)&img_device, MAX_IMAGE_INPUT_SIZE_THRESH * 3)); - int fcount = 0; - std::vector imgs_buffer(BATCH_SIZE); - for (int f = 0; f < (int)file_names.size(); f++) { - fcount++; - if (fcount < BATCH_SIZE && f + 1 != (int)file_names.size()) continue; - //auto start = std::chrono::system_clock::now(); - float *buffer_idx = (float*)buffers[inputIndex]; - for (int b = 0; b < fcount; b++) { - cv::Mat img = cv::imread(img_dir + "/" + file_names[f - fcount + 1 + b]); - if (img.empty()) continue; - imgs_buffer[b] = img; - size_t size_image = img.cols * img.rows * 3; - size_t size_image_dst = INPUT_H * INPUT_W * 3; - //copy data to pinned memory - memcpy(img_host, img.data, size_image); - //copy data to device memory - CUDA_CHECK(cudaMemcpyAsync(img_device, img_host, size_image, cudaMemcpyHostToDevice, stream)); - preprocess_kernel_img(img_device, img.cols, img.rows, buffer_idx, INPUT_W, INPUT_H, stream); - buffer_idx += size_image_dst; - cudaStreamSynchronize(stream); - } - // Run inference - auto start = std::chrono::system_clock::now(); - doInference(*context, stream, (void**)buffers, prob, proto, BATCH_SIZE); - auto end = std::chrono::system_clock::now(); - std::cout << "inference time: " << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; - std::vector> batch_res(fcount); - for (int b = 0; b < fcount; b++) { - auto& res = batch_res[b]; - nms(res, &prob[b * OUTPUT_SIZE1], CONF_THRESH, NMS_THRESH); - } - for (int b = 0; b < fcount; b++) { - auto& res = batch_res[b]; - cv::Mat img = imgs_buffer[b]; - - auto masks = process_mask(&proto[b * OUTPUT_SIZE2], res); - draw_mask_bbox(img, res, masks, labels_map); - cv::imwrite("_" + file_names[f - fcount + 1 + b], img); - } - fcount = 0; - } - - // Release stream and buffers - cudaStreamDestroy(stream); - CUDA_CHECK(cudaFree(img_device)); - CUDA_CHECK(cudaFreeHost(img_host)); - CUDA_CHECK(cudaFree(buffers[inputIndex])); - CUDA_CHECK(cudaFree(buffers[outputIndex1])); - CUDA_CHECK(cudaFree(buffers[outputIndex2])); - // Destroy the engine - context->destroy(); - engine->destroy(); - runtime->destroy(); - - - // Print histogram of the output distribution - //std::cout << "\nOutput:\n\n"; - //for (unsigned int i = 0; i < OUTPUT_SIZE; i++) - //{ - // std::cout << prob[i] << ", "; - // if (i % 10 == 0) std::cout << std::endl; - //} - //std::cout << std::endl; + std::string img_dir; + if (!parse_args(argc, argv, wts_name, engine_name, gd, gw, img_dir, labels_filename)) { + std::cerr << "arguments not right!" << std::endl; + std::cerr << "./yolov5_seg -s [.wts] [.engine] [n/s/m/l/x or c gd gw] // serialize model to plan file" << std::endl; + std::cerr << "./yolov5_seg -d [.engine] ../images coco.txt // deserialize plan file, read the labels file and run inference" << std::endl; + return -1; + } + // Create a model using the API directly and serialize it to a file + if (!wts_name.empty()) { + serialize_engine(kBatchSize, gd, gw, wts_name, engine_name); return 0; + } + + // Deserialize the engine from file + IRuntime* runtime = nullptr; + ICudaEngine* engine = nullptr; + IExecutionContext* context = nullptr; + deserialize_engine(engine_name, &runtime, &engine, &context); + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreate(&stream)); + + // Init CUDA preprocessing + cuda_preprocess_init(kMaxInputImageSize); + + // Prepare cpu and gpu buffers + float* gpu_buffers[3]; + float* cpu_output_buffer1 = nullptr; + float* cpu_output_buffer2 = nullptr; + prepare_buffers(engine, &gpu_buffers[0], &gpu_buffers[1], &gpu_buffers[2], &cpu_output_buffer1, &cpu_output_buffer2); + + // Read images from directory + std::vector file_names; + if (read_files_in_dir(img_dir.c_str(), file_names) < 0) { + std::cerr << "read_files_in_dir failed." << std::endl; + return -1; + } + + // Read the txt file for classnames + std::ifstream labels_file(labels_filename, std::ios::binary); + if (!labels_file.good()) { + std::cerr << "read " << labels_filename << " error!" << std::endl; + return -1; + } + std::unordered_map labels_map; + read_labels(labels_filename, labels_map); + assert(kNumClass == labels_map.size()); + + // batch predict + for (size_t i = 0; i < file_names.size(); i += kBatchSize) { + // Get a batch of images + std::vector img_batch; + std::vector img_name_batch; + for (size_t j = i; j < i + kBatchSize && j < file_names.size(); j++) { + cv::Mat img = cv::imread(img_dir + "/" + file_names[j]); + img_batch.push_back(img); + img_name_batch.push_back(file_names[j]); + } + + // Preprocess + cuda_batch_preprocess(img_batch, gpu_buffers[0], kInputW, kInputH, stream); + + // Run inference + auto start = std::chrono::system_clock::now(); + infer(*context, stream, (void**)gpu_buffers, cpu_output_buffer1, cpu_output_buffer2, kBatchSize); + auto end = std::chrono::system_clock::now(); + std::cout << "inference time: " << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; + + // NMS + std::vector> res_batch; + batch_nms(res_batch, cpu_output_buffer1, img_batch.size(), kOutputSize1, kConfThresh, kNmsThresh); + + // Draw result and save image + for (size_t b = 0; b < img_name_batch.size(); b++) { + auto& res = res_batch[b]; + cv::Mat img = img_batch[b]; + + auto masks = process_mask(&cpu_output_buffer2[b * kOutputSize2], kOutputSize2, res); + draw_mask_bbox(img, res, masks, labels_map); + cv::imwrite("_" + img_name_batch[b], img); + } + } + + // Release stream and buffers + cudaStreamDestroy(stream); + CUDA_CHECK(cudaFree(gpu_buffers[0])); + CUDA_CHECK(cudaFree(gpu_buffers[1])); + CUDA_CHECK(cudaFree(gpu_buffers[2])); + delete[] cpu_output_buffer1; + delete[] cpu_output_buffer2; + cuda_preprocess_destroy(); + // Destroy the engine + context->destroy(); + engine->destroy(); + runtime->destroy(); + + return 0; } diff --git a/yolov5/yolov5_seg_trt.py b/yolov5/yolov5_seg_trt.py index b0ac307..64ebd4a 100644 --- a/yolov5/yolov5_seg_trt.py +++ b/yolov5/yolov5_seg_trt.py @@ -543,7 +543,7 @@ if __name__ == "__main__": try: print('batch size is', yolov5_wrapper.batch_size) - image_dir = "samples/" + image_dir = "images/" image_path_batches = get_img_path_batches(yolov5_wrapper.batch_size, image_dir) for i in range(10):