Refactor yolov5 (#1211)
* refactor pre and post process * move model, refactor plugin * refactor yolov5 cls * refactor yolov5 seg * fix space * rename samples * update readme
This commit is contained in:
Wang Xinyu
GitHub
parent
57ba75f9a4
commit
208032970c
@ -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
|
||||
```
|
||||
|
||||
<p align="center">
|
||||
@ -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
|
||||
|
||||
|
||||
1
yolov5/images
Symbolic link
1
yolov5/images
Symbolic link
@ -0,0 +1 @@
|
||||
../yolov3-spp/samples
|
||||
@ -1,322 +1,280 @@
|
||||
#include <assert.h>
|
||||
#include <vector>
|
||||
#include <iostream>
|
||||
#include "yololayer.h"
|
||||
#include "cuda_utils.h"
|
||||
|
||||
namespace Tn
|
||||
{
|
||||
template<typename T>
|
||||
void write(char*& buffer, const T& val)
|
||||
{
|
||||
*reinterpret_cast<T*>(buffer) = val;
|
||||
buffer += sizeof(T);
|
||||
}
|
||||
#include <cassert>
|
||||
#include <vector>
|
||||
#include <iostream>
|
||||
|
||||
template<typename T>
|
||||
void read(const char*& buffer, T& val)
|
||||
{
|
||||
val = *reinterpret_cast<const T*>(buffer);
|
||||
buffer += sizeof(T);
|
||||
}
|
||||
namespace Tn {
|
||||
template<typename T>
|
||||
void write(char*& buffer, const T& val) {
|
||||
*reinterpret_cast<T*>(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<Yolo::YoloKernel>& 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<const char *>(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<char*>(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<PluginField> 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<Yolo::YoloKernel> 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<typename T>
|
||||
void read(const char*& buffer, T& val) {
|
||||
val = *reinterpret_cast<const T*>(buffer);
|
||||
buffer += sizeof(T);
|
||||
}
|
||||
}
|
||||
|
||||
namespace nvinfer1 {
|
||||
YoloLayerPlugin::YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, bool is_segmentation, const std::vector<YoloKernel>& 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<const char *>(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<char*>(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<PluginField> 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<YoloKernel> 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;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -1,140 +1,106 @@
|
||||
#ifndef _YOLO_LAYER_H
|
||||
#define _YOLO_LAYER_H
|
||||
#pragma once
|
||||
|
||||
#include "types.h"
|
||||
#include "macros.h"
|
||||
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <NvInfer.h>
|
||||
#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<YoloKernel>& 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<Yolo::YoloKernel>& 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<Yolo::YoloKernel> 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<PluginField> 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<YoloKernel> 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<PluginField> mPluginAttributes;
|
||||
};
|
||||
REGISTER_TENSORRT_PLUGIN(YoloPluginCreator);
|
||||
};
|
||||
|
||||
#endif // _YOLO_LAYER_H
|
||||
|
||||
@ -1 +0,0 @@
|
||||
../yolov3-spp/samples/
|
||||
@ -1,11 +1,35 @@
|
||||
#include <iostream>
|
||||
#include <iterator>
|
||||
#include <fstream>
|
||||
#include <opencv2/dnn/dnn.hpp>
|
||||
#include "calibrator.h"
|
||||
#include "cuda_utils.h"
|
||||
#include "utils.h"
|
||||
|
||||
#include <iostream>
|
||||
#include <iterator>
|
||||
#include <fstream>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <opencv2/dnn/dnn.hpp>
|
||||
|
||||
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<cv::Mat> 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<cv::Mat> 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<float>(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<float>(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<char>(input), std::istream_iterator<char>(), 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<char>(input), std::istream_iterator<char>(), 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<const char*>(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<const char*>(cache), length);
|
||||
}
|
||||
|
||||
|
||||
@ -1,10 +1,8 @@
|
||||
#ifndef ENTROPY_CALIBRATOR_H
|
||||
#define ENTROPY_CALIBRATOR_H
|
||||
#pragma once
|
||||
|
||||
#include <NvInfer.h>
|
||||
#include "macros.h"
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#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<std::string> img_files_;
|
||||
size_t input_count_;
|
||||
std::string calib_table_name_;
|
||||
const char* input_blob_name_;
|
||||
bool read_cache_;
|
||||
void* device_input_;
|
||||
std::vector<char> calib_cache_;
|
||||
private:
|
||||
int batchsize_;
|
||||
int input_w_;
|
||||
int input_h_;
|
||||
int img_idx_;
|
||||
std::string img_dir_;
|
||||
std::vector<std::string> img_files_;
|
||||
size_t input_count_;
|
||||
std::string calib_table_name_;
|
||||
const char* input_blob_name_;
|
||||
bool read_cache_;
|
||||
void* device_input_;
|
||||
std::vector<char> calib_cache_;
|
||||
};
|
||||
|
||||
#endif // ENTROPY_CALIBRATOR_H
|
||||
|
||||
@ -1,344 +0,0 @@
|
||||
#ifndef YOLOV5_COMMON_H_
|
||||
#define YOLOV5_COMMON_H_
|
||||
|
||||
#include <fstream>
|
||||
#include <map>
|
||||
#include <sstream>
|
||||
#include <vector>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#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<Yolo::Detection>& res, float *output, float conf_thresh, float nms_thresh = 0.5) {
|
||||
int det_size = sizeof(Yolo::Detection) / sizeof(float);
|
||||
std::map<float, std::vector<Yolo::Detection>> 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<Yolo::Detection>());
|
||||
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] <data x size in hex>
|
||||
std::map<std::string, Weights> loadWeights(const std::string file) {
|
||||
std::cout << "Loading weights: " << file << std::endl;
|
||||
std::map<std::string, Weights> 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<uint32_t*>(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<std::string, Weights>& 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<float*>(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<float*>(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<float*>(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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::vector<float>> getAnchors(std::map<std::string, Weights>& weightMap, std::string lname) {
|
||||
std::vector<std::vector<float>> 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<float> anchor(p, p + anchor_len);
|
||||
anchors.push_back(anchor);
|
||||
}
|
||||
return anchors;
|
||||
}
|
||||
|
||||
IPluginV2Layer* addYoLoLayer(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, std::string lname, std::vector<IConvolutionLayer*> 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<int> scales(p, p + strides.count);
|
||||
|
||||
std::vector<Yolo::YoloKernel> 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<ITensor*> 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_
|
||||
|
||||
55
yolov5/src/config.h
Normal file
55
yolov5/src/config.h
Normal file
@ -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;
|
||||
|
||||
@ -1,6 +1,8 @@
|
||||
#ifndef __MACROS_H
|
||||
#define __MACROS_H
|
||||
|
||||
#include <NvInfer.h>
|
||||
|
||||
#ifdef API_EXPORTS
|
||||
#if defined(_MSC_VER)
|
||||
#define API __declspec(dllexport)
|
||||
|
||||
628
yolov5/src/model.cpp
Normal file
628
yolov5/src/model.cpp
Normal file
@ -0,0 +1,628 @@
|
||||
#include "model.h"
|
||||
#include "calibrator.h"
|
||||
#include "config.h"
|
||||
#include "yololayer.h"
|
||||
|
||||
#include <iostream>
|
||||
#include <fstream>
|
||||
#include <map>
|
||||
#include <cassert>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
|
||||
using namespace nvinfer1;
|
||||
|
||||
// TensorRT weight files have a simple space delimited format:
|
||||
// [type] [size] <data x size in hex>
|
||||
static std::map<std::string, Weights> loadWeights(const std::string file) {
|
||||
std::cout << "Loading weights: " << file << std::endl;
|
||||
std::map<std::string, Weights> 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<uint32_t*>(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<int>(r, 1);
|
||||
}
|
||||
|
||||
static IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map<std::string, Weights>& 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<float*>(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<float*>(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<float*>(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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::vector<float>> getAnchors(std::map<std::string, Weights>& weightMap, std::string lname) {
|
||||
std::vector<std::vector<float>> 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<float> anchor(p, p + anchor_len);
|
||||
anchors.push_back(anchor);
|
||||
}
|
||||
return anchors;
|
||||
}
|
||||
|
||||
static IPluginV2Layer* addYoLoLayer(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, std::string lname, std::vector<IConvolutionLayer*> 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<int> scales(p, p + strides.count);
|
||||
|
||||
std::vector<YoloKernel> 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<ITensor*> 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<std::string, Weights> 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<IConvolutionLayer*>{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<std::string, Weights> 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<IConvolutionLayer*>{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<std::string, Weights> 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<std::string, Weights> 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<IConvolutionLayer*>{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;
|
||||
}
|
||||
|
||||
16
yolov5/src/model.h
Normal file
16
yolov5/src/model.h
Normal file
@ -0,0 +1,16 @@
|
||||
#pragma once
|
||||
|
||||
#include <NvInfer.h>
|
||||
#include <string>
|
||||
|
||||
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);
|
||||
189
yolov5/src/postprocess.cpp
Normal file
189
yolov5/src/postprocess.cpp
Normal file
@ -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<Detection>& res, float* output, float conf_thresh, float nms_thresh) {
|
||||
int det_size = sizeof(Detection) / sizeof(float);
|
||||
std::map<float, std::vector<Detection>> 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<Detection>());
|
||||
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<std::vector<Detection>>& 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<cv::Mat>& img_batch, std::vector<std::vector<Detection>>& 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<cv::Mat> process_mask(const float* proto, int proto_size, std::vector<Detection>& dets) {
|
||||
std::vector<cv::Mat> 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<float>(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<Detection>& dets, std::vector<cv::Mat>& masks, std::unordered_map<int, std::string>& labels_map) {
|
||||
static std::vector<uint32_t> 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<float>(y, x);
|
||||
if (val <= 0.5) continue;
|
||||
img.at<cv::Vec3b>(y, x)[0] = img.at<cv::Vec3b>(y, x)[0] / 2 + bgr[0] / 2;
|
||||
img.at<cv::Vec3b>(y, x)[1] = img.at<cv::Vec3b>(y, x)[1] / 2 + bgr[1] / 2;
|
||||
img.at<cv::Vec3b>(y, x)[2] = img.at<cv::Vec3b>(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);
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
16
yolov5/src/postprocess.h
Normal file
16
yolov5/src/postprocess.h
Normal file
@ -0,0 +1,16 @@
|
||||
#pragma once
|
||||
|
||||
#include "types.h"
|
||||
#include <opencv2/opencv.hpp>
|
||||
|
||||
cv::Rect get_rect(cv::Mat& img, float bbox[4]);
|
||||
|
||||
void nms(std::vector<Detection>& res, float *output, float conf_thresh, float nms_thresh = 0.5);
|
||||
|
||||
void batch_nms(std::vector<std::vector<Detection>>& batch_res, float *output, int batch_size, int output_size, float conf_thresh, float nms_thresh = 0.5);
|
||||
|
||||
void draw_bbox(std::vector<cv::Mat>& img_batch, std::vector<std::vector<Detection>>& res_batch);
|
||||
|
||||
std::vector<cv::Mat> process_mask(const float* proto, int proto_size, std::vector<Detection>& dets);
|
||||
|
||||
void draw_mask_bbox(cv::Mat& img, std::vector<Detection>& dets, std::vector<cv::Mat>& masks, std::unordered_map<int, std::string>& labels_map);
|
||||
@ -1,116 +1,153 @@
|
||||
#include "preprocess.h"
|
||||
#include <opencv2/opencv.hpp>
|
||||
#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<float>(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<<<blocks, threads, 0, stream>>>(
|
||||
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<float>(0), sizeof(d2s.value));
|
||||
|
||||
int jobs = dst_height * dst_width;
|
||||
int threads = 256;
|
||||
int blocks = ceil(jobs / (float)threads);
|
||||
|
||||
warpaffine_kernel<<<blocks, threads, 0, stream>>>(
|
||||
img_buffer_device, src_width * 3, src_width,
|
||||
src_height, dst, dst_width,
|
||||
dst_height, 128, d2s, jobs);
|
||||
}
|
||||
|
||||
void cuda_batch_preprocess(std::vector<cv::Mat>& 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));
|
||||
}
|
||||
|
||||
|
||||
@ -1,16 +1,15 @@
|
||||
#ifndef __PREPROCESS_H
|
||||
#define __PREPROCESS_H
|
||||
#pragma once
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cstdint>
|
||||
#include <opencv2/opencv.hpp>
|
||||
|
||||
|
||||
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<cv::Mat>& img_batch,
|
||||
float* dst, int dst_width, int dst_height,
|
||||
cudaStream_t stream);
|
||||
#endif // __PREPROCESS_H
|
||||
|
||||
|
||||
17
yolov5/src/types.h
Normal file
17
yolov5/src/types.h
Normal file
@ -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];
|
||||
};
|
||||
|
||||
@ -1,34 +1,12 @@
|
||||
#ifndef TRTX_YOLOV5_UTILS_H_
|
||||
#define TRTX_YOLOV5_UTILS_H_
|
||||
#pragma once
|
||||
|
||||
#include <dirent.h>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <fstream>
|
||||
#include <unordered_map>
|
||||
#include <string>
|
||||
#include <sstream>
|
||||
|
||||
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 <vector>
|
||||
#include <cstring>
|
||||
|
||||
static inline int read_files_in_dir(const char* p_dir_name, std::vector<std::string>& 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_
|
||||
|
||||
|
||||
@ -1,333 +1,286 @@
|
||||
#include "cuda_utils.h"
|
||||
#include "logging.h"
|
||||
#include "utils.h"
|
||||
#include "model.h"
|
||||
#include "config.h"
|
||||
|
||||
#include <iostream>
|
||||
#include <chrono>
|
||||
#include <cmath>
|
||||
#include <numeric>
|
||||
#include "cuda_utils.h"
|
||||
#include "logging.h"
|
||||
#include "common.hpp"
|
||||
#include "utils.h"
|
||||
#include "calibrator.h"
|
||||
#include <opencv2/opencv.hpp>
|
||||
|
||||
#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<cv::Mat>& 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<int>(r, 1);
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<float> softmax(float *prob, int n) {
|
||||
std::vector<float> 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<float> 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<int> topk(const std::vector<float>& vec, int k) {
|
||||
std::vector<int> topk_index;
|
||||
std::vector<size_t> vec_index(vec.size());
|
||||
std::iota(vec_index.begin(), vec_index.end(), 0);
|
||||
std::vector<int> topk_index;
|
||||
std::vector<size_t> 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<int>(vec.size(), k);
|
||||
int k_num = std::min<int>(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<std::string> read_classes(std::string file_name) {
|
||||
std::vector<std::string> 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<std::string, Weights> 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<std::string> 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<const char*>(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<const char*>(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<std::string> 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<std::chrono::milliseconds>(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<std::string> 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<cv::Mat> img_batch;
|
||||
std::vector<std::string> 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<std::chrono::milliseconds>(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;
|
||||
}
|
||||
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
@ -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 <iostream>
|
||||
#include <chrono>
|
||||
#include <cmath>
|
||||
|
||||
#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<int>(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<std::string, Weights> 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<IConvolutionLayer*>{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<std::string, Weights> 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<IConvolutionLayer*>{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<std::string> 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<std::string> 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<cv::Mat> 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<cv::Mat> img_batch;
|
||||
std::vector<std::string> 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<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
|
||||
std::vector<std::vector<Yolo::Detection>> 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<std::vector<Detection>> 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;
|
||||
}
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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 <iostream>
|
||||
#include <chrono>
|
||||
#include <cmath>
|
||||
#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<int>(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<std::string, Weights> 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<IConvolutionLayer*>{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<cv::Mat> process_mask(const float* proto, std::vector<Yolo::Detection>& dets) {
|
||||
std::vector<cv::Mat> 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<float>(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<const char*>(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<Yolo::Detection>& dets, std::vector<cv::Mat>& masks, std::unordered_map<int, std::string>& labels_map) {
|
||||
static std::vector<uint32_t> 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<float>(y, x);
|
||||
if (val <= 0.5) continue;
|
||||
img.at<cv::Vec3b>(y, x)[0] = img.at<cv::Vec3b>(y, x)[0] / 2 + bgr[0] / 2;
|
||||
img.at<cv::Vec3b>(y, x)[1] = img.at<cv::Vec3b>(y, x)[1] / 2 + bgr[1] / 2;
|
||||
img.at<cv::Vec3b>(y, x)[2] = img.at<cv::Vec3b>(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<const char*>(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<std::string> 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<int, std::string> 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<cv::Mat> 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<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
|
||||
std::vector<std::vector<Yolo::Detection>> 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<std::string> 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<int, std::string> 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<cv::Mat> img_batch;
|
||||
std::vector<std::string> 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<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
|
||||
|
||||
// NMS
|
||||
std::vector<std::vector<Detection>> 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;
|
||||
}
|
||||
|
||||
|
||||
@ -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):
|
||||
|
||||
Loading…
Reference in New Issue
Block a user