#include "../include/block.h"
|
#include <assert.h>
|
#include <math.h>
|
#include <fstream>
|
#include <iostream>
|
#include "../include/config.h"
|
#include "../plugin/yololayer.h"
|
|
std::map<std::string, nvinfer1::Weights> loadWeights(const std::string file) {
|
std::cout << "Loading weights: " << file << std::endl;
|
std::map<std::string, nvinfer1::Weights> WeightMap;
|
|
std::ifstream input(file);
|
assert(input.is_open() && "Unable to load weight file. please check if the .wts file path is right!!!!!!");
|
|
int32_t count;
|
input >> count;
|
assert(count > 0 && "Invalid weight map file.");
|
|
while (count--) {
|
nvinfer1::Weights wt{nvinfer1::DataType::kFLOAT, nullptr, 0};
|
uint32_t size;
|
|
std::string name;
|
input >> name >> std::dec >> size;
|
wt.type = nvinfer1::DataType::kFLOAT;
|
|
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 nvinfer1::IScaleLayer* addBatchNorm2d(nvinfer1::INetworkDefinition* network,
|
std::map<std::string, nvinfer1::Weights> weightMap,
|
nvinfer1::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);
|
}
|
nvinfer1::Weights scale{nvinfer1::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);
|
}
|
nvinfer1::Weights shift{nvinfer1::DataType::kFLOAT, shval, len};
|
|
float* pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
|
for (int i = 0; i < len; i++) {
|
pval[i] = 1.0;
|
}
|
nvinfer1::Weights power{nvinfer1::DataType::kFLOAT, pval, len};
|
weightMap[lname + ".scale"] = scale;
|
weightMap[lname + ".shift"] = shift;
|
weightMap[lname + ".power"] = power;
|
nvinfer1::IScaleLayer* output = network->addScale(input, nvinfer1::ScaleMode::kCHANNEL, shift, scale, power);
|
assert(output);
|
return output;
|
}
|
|
nvinfer1::IElementWiseLayer* convBnSiLU(nvinfer1::INetworkDefinition* network,
|
std::map<std::string, nvinfer1::Weights> weightMap, nvinfer1::ITensor& input,
|
int ch, int k, int s, int p, std::string lname) {
|
nvinfer1::Weights bias_empty{nvinfer1::DataType::kFLOAT, nullptr, 0};
|
nvinfer1::IConvolutionLayer* conv =
|
network->addConvolutionNd(input, ch, nvinfer1::DimsHW{k, k}, weightMap[lname + ".conv.weight"], bias_empty);
|
assert(conv);
|
conv->setStrideNd(nvinfer1::DimsHW{s, s});
|
conv->setPaddingNd(nvinfer1::DimsHW{p, p});
|
|
nvinfer1::IScaleLayer* bn = addBatchNorm2d(network, weightMap, *conv->getOutput(0), lname + ".bn", 1e-3);
|
|
nvinfer1::IActivationLayer* sigmoid = network->addActivation(*bn->getOutput(0), nvinfer1::ActivationType::kSIGMOID);
|
nvinfer1::IElementWiseLayer* ew =
|
network->addElementWise(*bn->getOutput(0), *sigmoid->getOutput(0), nvinfer1::ElementWiseOperation::kPROD);
|
assert(ew);
|
return ew;
|
}
|
|
nvinfer1::ILayer* bottleneck(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights> weightMap,
|
nvinfer1::ITensor& input, int c1, int c2, bool shortcut, float e, std::string lname) {
|
nvinfer1::IElementWiseLayer* conv1 = convBnSiLU(network, weightMap, input, c2, 3, 1, 1, lname + ".cv1");
|
nvinfer1::IElementWiseLayer* conv2 =
|
convBnSiLU(network, weightMap, *conv1->getOutput(0), c2, 3, 1, 1, lname + ".cv2");
|
|
if (shortcut && c1 == c2) {
|
nvinfer1::IElementWiseLayer* ew =
|
network->addElementWise(input, *conv2->getOutput(0), nvinfer1::ElementWiseOperation::kSUM);
|
return ew;
|
}
|
return conv2;
|
}
|
|
nvinfer1::IElementWiseLayer* C2F(nvinfer1::INetworkDefinition* network,
|
std::map<std::string, nvinfer1::Weights> weightMap, nvinfer1::ITensor& input, int c1,
|
int c2, int n, bool shortcut, float e, std::string lname) {
|
int c_ = (float)c2 * e;
|
|
nvinfer1::IElementWiseLayer* conv1 = convBnSiLU(network, weightMap, input, 2 * c_, 1, 1, 0, lname + ".cv1");
|
nvinfer1::Dims d = conv1->getOutput(0)->getDimensions();
|
|
nvinfer1::ISliceLayer* split1 =
|
network->addSlice(*conv1->getOutput(0), nvinfer1::Dims3{0, 0, 0},
|
nvinfer1::Dims3{d.d[0] / 2, d.d[1], d.d[2]}, nvinfer1::Dims3{1, 1, 1});
|
nvinfer1::ISliceLayer* split2 =
|
network->addSlice(*conv1->getOutput(0), nvinfer1::Dims3{d.d[0] / 2, 0, 0},
|
nvinfer1::Dims3{d.d[0] / 2, d.d[1], d.d[2]}, nvinfer1::Dims3{1, 1, 1});
|
nvinfer1::ITensor* inputTensor0[] = {split1->getOutput(0), split2->getOutput(0)};
|
nvinfer1::IConcatenationLayer* cat = network->addConcatenation(inputTensor0, 2);
|
nvinfer1::ITensor* y1 = split2->getOutput(0);
|
for (int i = 0; i < n; i++) {
|
auto* b = bottleneck(network, weightMap, *y1, c_, c_, shortcut, 1.0, lname + ".m." + std::to_string(i));
|
y1 = b->getOutput(0);
|
|
nvinfer1::ITensor* inputTensors[] = {cat->getOutput(0), b->getOutput(0)};
|
cat = network->addConcatenation(inputTensors, 2);
|
}
|
|
nvinfer1::IElementWiseLayer* conv2 =
|
convBnSiLU(network, weightMap, *cat->getOutput(0), c2, 1, 1, 0, lname + ".cv2");
|
|
return conv2;
|
}
|
|
nvinfer1::IElementWiseLayer* C2(nvinfer1::INetworkDefinition* network,
|
std::map<std::string, nvinfer1::Weights>& weightMap, nvinfer1::ITensor& input, int c1,
|
int c2, int n, bool shortcut, float e, std::string lname) {
|
assert(network != nullptr);
|
int hidden_channels = static_cast<int>(c2 * e);
|
|
// cv1 branch
|
nvinfer1::IElementWiseLayer* conv1 =
|
convBnSiLU(network, weightMap, input, 2 * hidden_channels, 1, 1, 0, lname + ".cv1");
|
nvinfer1::ITensor* cv1_out = conv1->getOutput(0);
|
|
// Split the output of cv1 into two tensors
|
nvinfer1::Dims dims = cv1_out->getDimensions();
|
nvinfer1::ISliceLayer* split1 =
|
network->addSlice(*cv1_out, nvinfer1::Dims3{0, 0, 0}, nvinfer1::Dims3{dims.d[0] / 2, dims.d[1], dims.d[2]},
|
nvinfer1::Dims3{1, 1, 1});
|
nvinfer1::ISliceLayer* split2 =
|
network->addSlice(*cv1_out, nvinfer1::Dims3{dims.d[0] / 2, 0, 0},
|
nvinfer1::Dims3{dims.d[0] / 2, dims.d[1], dims.d[2]}, nvinfer1::Dims3{1, 1, 1});
|
|
// Create y1 bottleneck sequence
|
nvinfer1::ITensor* y1 = split1->getOutput(0);
|
for (int i = 0; i < n; ++i) {
|
auto* bottleneck_layer = bottleneck(network, weightMap, *y1, hidden_channels, hidden_channels, shortcut, 1.0,
|
lname + ".m." + std::to_string(i));
|
y1 = bottleneck_layer->getOutput(0); // update 'y1' to be the output of the current bottleneck
|
}
|
|
// Concatenate y1 with the second split of cv1
|
nvinfer1::ITensor* concatInputs[2] = {y1, split2->getOutput(0)};
|
nvinfer1::IConcatenationLayer* cat = network->addConcatenation(concatInputs, 2);
|
|
// cv2 to produce the final output
|
nvinfer1::IElementWiseLayer* conv2 =
|
convBnSiLU(network, weightMap, *cat->getOutput(0), c2, 1, 1, 0, lname + ".cv2");
|
|
return conv2;
|
}
|
|
nvinfer1::IElementWiseLayer* SPPF(nvinfer1::INetworkDefinition* network,
|
std::map<std::string, nvinfer1::Weights> weightMap, nvinfer1::ITensor& input, int c1,
|
int c2, int k, std::string lname) {
|
int c_ = c1 / 2;
|
nvinfer1::IElementWiseLayer* conv1 = convBnSiLU(network, weightMap, input, c_, 1, 1, 0, lname + ".cv1");
|
nvinfer1::IPoolingLayer* pool1 =
|
network->addPoolingNd(*conv1->getOutput(0), nvinfer1::PoolingType::kMAX, nvinfer1::DimsHW{k, k});
|
pool1->setStrideNd(nvinfer1::DimsHW{1, 1});
|
pool1->setPaddingNd(nvinfer1::DimsHW{k / 2, k / 2});
|
nvinfer1::IPoolingLayer* pool2 =
|
network->addPoolingNd(*pool1->getOutput(0), nvinfer1::PoolingType::kMAX, nvinfer1::DimsHW{k, k});
|
pool2->setStrideNd(nvinfer1::DimsHW{1, 1});
|
pool2->setPaddingNd(nvinfer1::DimsHW{k / 2, k / 2});
|
nvinfer1::IPoolingLayer* pool3 =
|
network->addPoolingNd(*pool2->getOutput(0), nvinfer1::PoolingType::kMAX, nvinfer1::DimsHW{k, k});
|
pool3->setStrideNd(nvinfer1::DimsHW{1, 1});
|
pool3->setPaddingNd(nvinfer1::DimsHW{k / 2, k / 2});
|
nvinfer1::ITensor* inputTensors[] = {conv1->getOutput(0), pool1->getOutput(0), pool2->getOutput(0),
|
pool3->getOutput(0)};
|
nvinfer1::IConcatenationLayer* cat = network->addConcatenation(inputTensors, 4);
|
nvinfer1::IElementWiseLayer* conv2 =
|
convBnSiLU(network, weightMap, *cat->getOutput(0), c2, 1, 1, 0, lname + ".cv2");
|
return conv2;
|
}
|
|
nvinfer1::IShuffleLayer* DFL(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights> weightMap,
|
nvinfer1::ITensor& input, int ch, int grid, int k, int s, int p, std::string lname) {
|
|
nvinfer1::IShuffleLayer* shuffle1 = network->addShuffle(input);
|
shuffle1->setReshapeDimensions(nvinfer1::Dims3{4, 16, grid});
|
shuffle1->setSecondTranspose(nvinfer1::Permutation{1, 0, 2});
|
nvinfer1::ISoftMaxLayer* softmax = network->addSoftMax(*shuffle1->getOutput(0));
|
|
nvinfer1::Weights bias_empty{nvinfer1::DataType::kFLOAT, nullptr, 0};
|
nvinfer1::IConvolutionLayer* conv =
|
network->addConvolutionNd(*softmax->getOutput(0), 1, nvinfer1::DimsHW{1, 1}, weightMap[lname], bias_empty);
|
conv->setStrideNd(nvinfer1::DimsHW{s, s});
|
conv->setPaddingNd(nvinfer1::DimsHW{p, p});
|
|
nvinfer1::IShuffleLayer* shuffle2 = network->addShuffle(*conv->getOutput(0));
|
shuffle2->setReshapeDimensions(nvinfer1::Dims2{4, grid});
|
|
return shuffle2;
|
}
|
|
nvinfer1::IPluginV2Layer* addYoLoLayer(nvinfer1::INetworkDefinition* network,
|
std::vector<nvinfer1::IConcatenationLayer*> dets, const int* px_arry,
|
int px_arry_num, bool is_segmentation, bool is_pose) {
|
auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1");
|
const int netinfo_count = 8; // Assuming the first 5 elements are for netinfo as per existing code.
|
const int total_count = netinfo_count + px_arry_num; // Total number of elements for netinfo and px_arry combined.
|
|
std::vector<int> combinedInfo(total_count);
|
// Fill in the first 5 elements as per existing netinfo.
|
combinedInfo[0] = kNumClass;
|
combinedInfo[1] = kNumberOfPoints;
|
combinedInfo[2] = kConfThreshKeypoints;
|
combinedInfo[3] = kInputW;
|
combinedInfo[4] = kInputH;
|
combinedInfo[5] = kMaxNumOutputBbox;
|
combinedInfo[6] = is_segmentation;
|
combinedInfo[7] = is_pose;
|
|
// Copy the contents of px_arry into the combinedInfo vector after the initial 5 elements.
|
std::copy(px_arry, px_arry + px_arry_num, combinedInfo.begin() + netinfo_count);
|
|
// Now let's create the PluginField object to hold this combined information.
|
nvinfer1::PluginField pluginField;
|
pluginField.name = "combinedInfo"; // This can be any name that the plugin will recognize
|
pluginField.data = combinedInfo.data();
|
pluginField.type = nvinfer1::PluginFieldType::kINT32;
|
pluginField.length = combinedInfo.size();
|
|
// Create the PluginFieldCollection to hold the PluginField object.
|
nvinfer1::PluginFieldCollection pluginFieldCollection;
|
pluginFieldCollection.nbFields = 1; // We have just one field, but it's a combined array
|
pluginFieldCollection.fields = &pluginField;
|
|
// Create the plugin object using the PluginFieldCollection.
|
nvinfer1::IPluginV2* pluginObject = creator->createPlugin("yololayer", &pluginFieldCollection);
|
|
// We assume that the plugin is to be added onto the network.
|
// Prepare input tensors for the YOLO Layer.
|
std::vector<nvinfer1::ITensor*> inputTensors;
|
for (auto det : dets) {
|
inputTensors.push_back(det->getOutput(0)); // Assuming each IConcatenationLayer has one output tensor.
|
}
|
|
// Add the plugin to the network using the prepared input tensors.
|
nvinfer1::IPluginV2Layer* yoloLayer = network->addPluginV2(inputTensors.data(), inputTensors.size(), *pluginObject);
|
|
return yoloLayer; // Return the added YOLO layer.
|
}
|
