#include "trt_utils.h"
|
#include <NvInferRuntimeCommon.h>
|
#include <experimental/filesystem>
|
#include <fstream>
|
#include <iomanip>
|
using namespace nvinfer1;
|
REGISTER_TENSORRT_PLUGIN(MishPluginCreator);
|
REGISTER_TENSORRT_PLUGIN(ChunkPluginCreator);
|
REGISTER_TENSORRT_PLUGIN(HardswishPluginCreator);
|
|
cv::Mat blobFromDsImages(const std::vector<DsImage>& inputImages,
|
const int& inputH,
|
const int& inputW)
|
{
|
std::vector<cv::Mat> letterboxStack(inputImages.size());
|
for (uint32_t i = 0; i < inputImages.size(); ++i)
|
{
|
inputImages.at(i).getLetterBoxedImage().copyTo(letterboxStack.at(i));
|
}
|
return cv::dnn::blobFromImages(letterboxStack, 1.0, cv::Size(inputW, inputH),
|
cv::Scalar(0.0, 0.0, 0.0),true);
|
}
|
|
// static void leftTrim(std::string& s)
|
// {
|
// s.erase(s.begin(), find_if(s.begin(), s.end(), [](int ch) { return !isspace(ch); }));
|
// }
|
|
// static void rightTrim(std::string& s)
|
// {
|
// s.erase(find_if(s.rbegin(), s.rend(), [](int ch) { return !isspace(ch); }).base(), s.end());
|
// }
|
|
// std::string trim(std::string s)
|
// {
|
// leftTrim(s);
|
// rightTrim(s);
|
// return s;
|
// }
|
|
// std::string triml(std::string s,const char* t)
|
// {
|
// s.erase(0, s.find_first_not_of(t));
|
// return s;
|
// }
|
|
// std::string trimr(std::string s, const char* t)
|
// {
|
// s.erase(s.find_last_not_of(t) + 1);
|
// return s;
|
// }
|
|
float clamp(const float val, const float minVal, const float maxVal)
|
{
|
assert(minVal <= maxVal);
|
return std::min(maxVal, std::max(minVal, val));
|
}
|
|
bool fileExists(const std::string fileName, bool verbose)
|
{
|
if (!std::experimental::filesystem::exists(std::experimental::filesystem::path(fileName)))
|
{
|
if (verbose) std::cout << "File does not exist : " << fileName << std::endl;
|
return false;
|
}
|
return true;
|
}
|
|
// BBox convertBBoxNetRes(const float& bx, const float& by, const float& bw, const float& bh,
|
// const uint32_t& stride, const uint32_t& netW, const uint32_t& netH)
|
// {
|
// BBox b;
|
// // Restore coordinates to network input resolution
|
// float x = bx * stride;
|
// float y = by * stride;
|
|
// b.x1 = x - bw / 2;
|
// b.x2 = x + bw / 2;
|
|
// b.y1 = y - bh / 2;
|
// b.y2 = y + bh / 2;
|
|
// b.x1 = clamp(b.x1, 0, netW);
|
// b.x2 = clamp(b.x2, 0, netW);
|
// b.y1 = clamp(b.y1, 0, netH);
|
// b.y2 = clamp(b.y2, 0, netH);
|
|
// return b;
|
// }
|
|
// void convertBBoxImgRes(const float scalingFactor,
|
// const float xOffset,
|
// const float yOffset,
|
// BBox& bbox)
|
// {
|
// //// Undo Letterbox
|
// bbox.x1 -= xOffset;
|
// bbox.x2 -= xOffset;
|
// bbox.y1 -= yOffset;
|
// bbox.y2 -= yOffset;
|
// //// Restore to input resolution
|
// bbox.x1 /= scalingFactor;
|
// bbox.x2 /= scalingFactor;
|
// bbox.y1 /= scalingFactor;
|
// bbox.y2 /= scalingFactor;
|
// std::cout << "convertBBoxImgRes" << std::endl;
|
|
|
// }
|
|
// void printPredictions(const BBoxInfo& b, const std::string& className)
|
// {
|
// std::cout << " label:" << b.label << "(" << className << ")"
|
// << " confidence:" << b.prob << " xmin:" << b.box.x1 << " ymin:" << b.box.y1
|
// << " xmax:" << b.box.x2 << " ymax:" << b.box.y2 << std::endl;
|
// }
|
//
|
uint64_t get3DTensorVolume(nvinfer1::Dims inputDims)
|
{
|
assert(inputDims.nbDims == 3);
|
return inputDims.d[0] * inputDims.d[1] * inputDims.d[2];
|
}
|
|
|
nvinfer1::ICudaEngine* loadTRTEngine(const std::string planFilePath, PluginFactory* pluginFactory,
|
Logger& logger)
|
{
|
// reading the model in memory
|
std::cout << "Loading TRT Engine..." << std::endl;
|
assert(fileExists(planFilePath));
|
std::stringstream trtModelStream;
|
trtModelStream.seekg(0, trtModelStream.beg);
|
std::ifstream cache(planFilePath,std::ios::binary | std::ios::in);
|
assert(cache.good());
|
trtModelStream << cache.rdbuf();
|
cache.close();
|
|
// calculating model size
|
trtModelStream.seekg(0, std::ios::end);
|
const auto modelSize = trtModelStream.tellg();
|
trtModelStream.seekg(0, std::ios::beg);
|
void* modelMem = malloc(modelSize);
|
trtModelStream.read((char*) modelMem, modelSize);
|
|
nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);
|
std::cout << "test................................" << std::endl;
|
nvinfer1::ICudaEngine* engine
|
= runtime->deserializeCudaEngine(modelMem, modelSize, pluginFactory);
|
free(modelMem);
|
runtime->destroy();
|
std::cout << "Loading Complete!" << std::endl;
|
|
return engine;
|
}
|
|
std::vector<BBoxInfo> nmsAllClasses(const float nmsThresh,
|
std::vector<BBoxInfo>& binfo,
|
const uint32_t numClasses,
|
const std::string &model_type)
|
{
|
std::vector<BBoxInfo> result;
|
std::vector<std::vector<BBoxInfo>> splitBoxes(numClasses);
|
for (auto& box : binfo)
|
{
|
splitBoxes.at(box.label).push_back(box);
|
}
|
|
for (auto& boxes : splitBoxes)
|
{
|
boxes = nonMaximumSuppression(nmsThresh, boxes);
|
result.insert(result.end(), boxes.begin(), boxes.end());
|
}
|
|
return result;
|
}
|
|
std::vector<BBoxInfo> nonMaximumSuppression(const float nmsThresh, std::vector<BBoxInfo> binfo)
|
{
|
auto overlap1D = [](float x1min, float x1max, float x2min, float x2max) -> float
|
{
|
if (x1min > x2min)
|
{
|
std::swap(x1min, x2min);
|
std::swap(x1max, x2max);
|
}
|
return x1max < x2min ? 0 : std::min(x1max, x2max) - x2min;
|
};
|
auto computeIoU = [&overlap1D](BBox& bbox1, BBox& bbox2) -> float
|
{
|
float overlapX = overlap1D(bbox1.x1, bbox1.x2, bbox2.x1, bbox2.x2);
|
float overlapY = overlap1D(bbox1.y1, bbox1.y2, bbox2.y1, bbox2.y2);
|
float area1 = (bbox1.x2 - bbox1.x1) * (bbox1.y2 - bbox1.y1);
|
float area2 = (bbox2.x2 - bbox2.x1) * (bbox2.y2 - bbox2.y1);
|
float overlap2D = overlapX * overlapY;
|
float u = area1 + area2 - overlap2D;
|
return u == 0 ? 0 : overlap2D / u;
|
};
|
|
std::stable_sort(binfo.begin(), binfo.end(),
|
[](const BBoxInfo& b1, const BBoxInfo& b2) { return b1.prob > b2.prob; });
|
std::vector<BBoxInfo> out;
|
for (auto& i : binfo)
|
{
|
bool keep = true;
|
for (auto& j : out)
|
{
|
if (keep)
|
{
|
float overlap = computeIoU(i.box, j.box);
|
keep = overlap <= nmsThresh;
|
}
|
else
|
break;
|
}
|
if (keep) out.push_back(i);
|
}
|
return out;
|
}
|
