ANSCENTER/ANSCORE
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
7e772f76bc534d47b23a30b45bdcff7dd3f99977
ANSCORE/modules/ANSFR/RetinaFaceTRT.cpp

366 lines
16 KiB
C++
Raw Blame History

#include "RetinaFaceTRT.h"
// This is standalone Retina Face detector using TensorRT (We will not use as it is not inherit the ANSFD class)
namespace ANSCENTER {
RetinaFaceTRT::RetinaFaceTRT()
{
m_outputBbox.clear();
}
bool RetinaFaceTRT::Initialize(const std::string engineFile,
int frameWidth,
int frameHeight,
std::string inputName,
std::vector<std::string> outputNames,
std::vector<int> inputShape,
int maxBatchSize,
int maxFacesPerScene,
float nms_threshold,
float bbox_threshold)
{
try
{
assert(inputShape.size() == 3);
m_INPUT_C = static_cast<const int>(inputShape[0]);
m_INPUT_H = static_cast<const int>(inputShape[1]);
m_INPUT_W = static_cast<const int>(inputShape[2]);
m_INPUT_SIZE = static_cast<const int>(m_INPUT_C * m_INPUT_H * m_INPUT_W * sizeof(float));
m_OUTPUT_SIZE_BASE = static_cast<const int>((m_INPUT_H / 8 * m_INPUT_W / 8 + m_INPUT_H / 16 * m_INPUT_W / 16 + m_INPUT_H / 32 * m_INPUT_W / 32) * 2);
m_output0 = new float[m_OUTPUT_SIZE_BASE * 4];
m_output1 = new float[m_OUTPUT_SIZE_BASE * 2];
m_maxBatchSize = static_cast<const int>(maxBatchSize);
m_maxFacesPerScene = static_cast<const int>(maxFacesPerScene);
m_nms_threshold = static_cast<const float>(nms_threshold);
m_bbox_threshold = static_cast<const float>(bbox_threshold);
// load engine from .engine file
LoadEngine(engineFile);
// create stream and pre-allocate GPU buffers memory
PreInference(inputName, outputNames);
return true;
}
catch (std::exception& e) {
this->_logger.LogFatal("RetinaFace::Initialize", e.what(), __FILE__, __LINE__);
return false;
}
}
void RetinaFaceTRT::LoadEngine(const std::string engineFile)
{
try {
if (FileExist(engineFile)) {
this->_logger.LogDebug("RetinaFace::LoadEngine", "Loading RetinaFace Engine...", __FILE__, __LINE__);
std::vector<char> trtModelStream_;
size_t size{ 0 };
std::ifstream file(engineFile, std::ios::binary);
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream_.resize(size);
file.read(trtModelStream_.data(), size);
file.close();
}
nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(m_logger);
assert(runtime != nullptr);
m_engine = runtime->deserializeCudaEngine(trtModelStream_.data(), size);
assert(m_engine != nullptr);
m_context = m_engine->createExecutionContext();
assert(m_context != nullptr);
}
else {
this->_logger.LogError("RetinaFace::LoadEngine", "Cant find engine file", __FILE__, __LINE__);
}
}
catch (std::exception& e) {
this->_logger.LogFatal("RetinaFace::Initialize", e.what(), __FILE__, __LINE__);
}
}
void RetinaFaceTRT::PreInference(std::string inputName, std::vector<std::string> outputNames) {
try {
/*
Does not make landmark head as we do not use face alignment.
*/
// Assert
assert(outputNames.size() == 2);
#if NV_TENSORRT_MAJOR >= 10
// TRT 10+: use named tensor API
assert(m_engine->getNbIOTensors() == 3);
// Look up tensor indices by name
for (int i = 0; i < m_engine->getNbIOTensors(); ++i) {
const char* name = m_engine->getIOTensorName(i);
if (name == inputName) inputIndex = i;
else if (name == outputNames[0]) outputIndex0 = i;
else if (name == outputNames[1]) outputIndex1 = i;
}
#else
// TRT 8.x: use binding API
assert(m_engine->getNbBindings() == 3);
inputIndex = m_engine->getBindingIndex(inputName.c_str());
outputIndex0 = m_engine->getBindingIndex(outputNames[0].c_str());
outputIndex1 = m_engine->getBindingIndex(outputNames[1].c_str());
#endif
// Create GPU buffers on device
ANSFRHelper::CheckCudaStatus(cudaMalloc(&buffers[inputIndex], m_maxBatchSize * m_INPUT_SIZE));
ANSFRHelper::CheckCudaStatus(cudaMalloc(&buffers[outputIndex0], m_maxBatchSize * m_OUTPUT_SIZE_BASE * 4 * sizeof(float)));
ANSFRHelper::CheckCudaStatus(cudaMalloc(&buffers[outputIndex1], m_maxBatchSize * m_OUTPUT_SIZE_BASE * 2 * sizeof(float)));
#if NV_TENSORRT_MAJOR >= 10
// TRT 10+: bind tensor addresses
m_context->setTensorAddress(inputName.c_str(), buffers[inputIndex]);
m_context->setTensorAddress(outputNames[0].c_str(), buffers[outputIndex0]);
m_context->setTensorAddress(outputNames[1].c_str(), buffers[outputIndex1]);
#endif
// Create stream
ANSFRHelper::CheckCudaStatus(cudaStreamCreate(&stream));
}
catch (std::exception& e) {
this->_logger.LogFatal("RetinaFace::PreInference", e.what(), __FILE__, __LINE__);
}
}
void RetinaFaceTRT::PreProcess(cv::Mat& img) {
try {
// Release input vector
m_input.release();
// Resize
float w, h, x, y;
if (m_scale_h > m_scale_w) {
w = float(m_INPUT_W);
h = float(m_scale_w * img.rows);
x = 0;
y = float((m_INPUT_H - h) / 2);
}
else {
w = float(m_scale_h * img.cols);
h = float(m_INPUT_H);
x = float((m_INPUT_W - w) / 2);
y = 0;
}
cv::Mat re((int)h, (int)w, CV_8UC3);
cv::resize(img, re, re.size(), 0, 0, cv::INTER_LINEAR);
cv::Mat out((int)m_INPUT_H, (int)m_INPUT_W, CV_8UC3, cv::Scalar(128, 128, 128));
re.copyTo(out(cv::Rect(int(x), int(y), re.cols, re.rows)));
// Normalize
out.convertTo(out, CV_32F);
out = out - cv::Scalar(104, 117, 123);
std::vector<cv::Mat> temp;
cv::split(out, temp);
for (int i = 0; i < temp.size(); i++) {
m_input.push_back(temp[i]);
}
}
catch (std::exception& e) {
this->_logger.LogFatal("RetinaFace::PreProcess", e.what(), __FILE__, __LINE__);
}
}
void RetinaFaceTRT::RunInference(float* input, float* output0, float* output1) {
try {
// DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
ANSCENTER::ANSFRHelper::CheckCudaStatus(cudaMemcpyAsync(buffers[inputIndex], input, m_maxBatchSize * m_INPUT_SIZE, cudaMemcpyHostToDevice, stream));
#if NV_TENSORRT_MAJOR >= 10
m_context->enqueueV3(stream);
#else
m_context->enqueueV2(buffers, stream, nullptr);
#endif
ANSFRHelper::CheckCudaStatus(cudaMemcpyAsync(output0, buffers[outputIndex0], m_maxBatchSize * m_OUTPUT_SIZE_BASE * 4 * sizeof(float), cudaMemcpyDeviceToHost, stream));
ANSFRHelper::CheckCudaStatus(cudaMemcpyAsync(output1, buffers[outputIndex1], m_maxBatchSize * m_OUTPUT_SIZE_BASE * 2 * sizeof(float), cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
}
catch (std::exception& e) {
this->_logger.LogFatal("RetinaFace::RunInference", e.what(), __FILE__, __LINE__);
}
}
std::vector<struct Bbox> RetinaFaceTRT::FindFace(cv::Mat& img) {
try {
m_outputBbox.clear();
int out_rows = img.rows;
int out_cols = img.cols;
m_frameWidth = static_cast<const int>(out_cols);
m_frameHeight = static_cast<const int>(out_rows);
m_scale_h = (float)(m_INPUT_H) / float(m_frameHeight);
m_scale_w = (float)(m_INPUT_W) / float(m_frameWidth);
PreProcess(img);
RunInference((float*)m_input.ptr<float>(0), m_output0, m_output1);
PostProcessing(m_output0, m_output1);
return m_outputBbox;
}
catch (std::exception& e) {
m_outputBbox.clear();
this->_logger.LogFatal("RetinaFace::FindFace", e.what(), __FILE__, __LINE__);
return m_outputBbox;
}
}
void RetinaFaceTRT::PostProcessing(float* bbox, float* conf) {
try {
m_outputBbox.clear();
std::vector<AnchorBox> anchor;
CreateAnchorRetinaface(anchor, m_INPUT_W, m_INPUT_H);
for (int i = 0; i < anchor.size(); ++i) {
if (*(conf + 1) > m_bbox_threshold) {
AnchorBox tmp = anchor[i];
AnchorBox tmp1;
Bbox result{};
// decode bbox (y - W; x - H)
tmp1.cx = float(tmp.cx + *bbox * 0.1 * tmp.sx);
tmp1.cy = float(tmp.cy + *(bbox + 1) * 0.1 * tmp.sy);
tmp1.sx = float(tmp.sx * exp(*(bbox + 2) * 0.2));
tmp1.sy = float(tmp.sy * exp(*(bbox + 3) * 0.2));
result.y1 = int((tmp1.cx - tmp1.sx / 2) * m_INPUT_W);
result.x1 = int((tmp1.cy - tmp1.sy / 2) * m_INPUT_H);
result.y2 = int((tmp1.cx + tmp1.sx / 2) * m_INPUT_W);
result.x2 = int((tmp1.cy + tmp1.sy / 2) * m_INPUT_H);
// rescale to original size
if (m_scale_h > m_scale_w) {
result.y1 = int(result.y1 / m_scale_w);
result.y2 = int(result.y2 / m_scale_w);
result.x1 = int((result.x1 - (m_INPUT_H - m_scale_w * m_frameHeight) / 2) / m_scale_w);
result.x2 = int((result.x2 - (m_INPUT_H - m_scale_w * m_frameHeight) / 2) / m_scale_w);
}
else {
result.y1 = int((result.y1 - (m_INPUT_W - m_scale_h * m_frameWidth) / 2) / m_scale_h);
result.y2 = int((result.y2 - (m_INPUT_W - m_scale_h * m_frameWidth) / 2) / m_scale_h);
result.x1 = int(result.x1 / m_scale_h);
result.x2 = int(result.x2 / m_scale_h);
}
// Clip object box coordinates to network resolution
result.y1 = CLIP(result.y1, 0, m_frameWidth - 1);
result.x1 = CLIP(result.x1, 0, m_frameHeight - 1);
result.y2 = CLIP(result.y2, 0, m_frameWidth - 1);
result.x2 = CLIP(result.x2, 0, m_frameHeight - 1);
// Get confidence
result.score = *(conf + 1);
// Push to result vector
m_outputBbox.push_back(result);
}
bbox += 4;
conf += 2;
}
std::sort(m_outputBbox.begin(), m_outputBbox.end(), MCMP);
NMS(m_outputBbox, m_nms_threshold);
if (m_outputBbox.size() > m_maxFacesPerScene)
m_outputBbox.resize(m_maxFacesPerScene);
}
catch (std::exception& e) {
m_outputBbox.clear();
this->_logger.LogFatal("RetinaFace::PostProcessing", e.what(), __FILE__, __LINE__);
}
}
void RetinaFaceTRT::CreateAnchorRetinaface(std::vector<AnchorBox>& anchor, int w, int h) {
try {
anchor.clear();
std::vector<std::vector<int>> feature_map(3), min_sizes(3);
float steps[] = { 8, 16, 32 };
for (int i = 0; i < feature_map.size(); ++i) {
feature_map[i].push_back(int(ceil(h / steps[i])));
feature_map[i].push_back(int(ceil(w / steps[i])));
}
std::vector<int> minsize1 = { 10, 20 };
min_sizes[0] = minsize1;
std::vector<int> minsize2 = { 32, 64 };
min_sizes[1] = minsize2;
std::vector<int> minsize3 = { 128, 256 };
min_sizes[2] = minsize3;
for (int k = 0; k < feature_map.size(); ++k) {
std::vector<int> min_size = min_sizes[k];
for (int i = 0; i < feature_map[k][0]; ++i) {
for (int j = 0; j < feature_map[k][1]; ++j) {
for (int l = 0; l < min_size.size(); ++l) {
float s_kx = float(min_size[l] * 1.0 / w);
float s_ky = float(min_size[l] * 1.0 / h);
float cx = float((j + 0.5) * steps[k] / w);
float cy = float((i + 0.5) * steps[k] / h);
AnchorBox axil = { cx, cy, s_kx, s_ky };
anchor.push_back(axil);
}
}
}
}
}
catch (std::exception& e) {
m_outputBbox.clear();
this->_logger.LogFatal("RetinaFace::CreateAnchorRetinaface", e.what(), __FILE__, __LINE__);
}
}
inline bool RetinaFaceTRT::MCMP(Bbox a, Bbox b) {
if (a.score > b.score)
return true;
return false;
}
void RetinaFaceTRT::NMS(std::vector<Bbox>& input_boxes, float NMS_THRESH) {
try {
std::vector<float> vArea(input_boxes.size());
for (int i = 0; i < int(input_boxes.size()); ++i) {
vArea[i] = float((input_boxes.at(i).x2 - input_boxes.at(i).x1 + 1) * (input_boxes.at(i).y2 - input_boxes.at(i).y1 + 1));
}
for (int i = 0; i < int(input_boxes.size()); ++i) {
for (int j = i + 1; j < int(input_boxes.size());) {
float xx1 = (float)(std::max(input_boxes[i].x1, input_boxes[j].x1));
float yy1 = (float)(std::max(input_boxes[i].y1, input_boxes[j].y1));
float xx2 = (float)(std::min(input_boxes[i].x2, input_boxes[j].x2));
float yy2 = (float)(std::min(input_boxes[i].y2, input_boxes[j].y2));
float w = (float)(std::max(float(0), xx2 - xx1 + 1));
float h = (float)(std::max(float(0), yy2 - yy1 + 1));
float inter = (float)(w * h);
float ovr = (float)(inter / (vArea[i] + vArea[j] - inter));
if (ovr >= NMS_THRESH) {
input_boxes.erase(input_boxes.begin() + j);
vArea.erase(vArea.begin() + j);
}
else {
j++;
}
}
}
}
catch (std::exception& e) {
m_outputBbox.clear();
this->_logger.LogFatal("RetinaFace::NMS", e.what(), __FILE__, __LINE__);
}
}
RetinaFaceTRT::~RetinaFaceTRT() {
try {
// Release stream and buffers
ANSFRHelper::CheckCudaStatus(cudaStreamDestroy(stream));
ANSFRHelper::CheckCudaStatus(cudaFree(buffers[inputIndex]));
ANSFRHelper::CheckCudaStatus(cudaFree(buffers[outputIndex0]));
ANSFRHelper::CheckCudaStatus(cudaFree(buffers[outputIndex1]));
// checkCudaStatus(cudaFree(buffers[outputIndex2]));
}
catch (std::exception& e) {
m_outputBbox.clear();
this->_logger.LogFatal("RetinaFace::~RetinaFace", e.what(), __FILE__, __LINE__);
}
}
}
Reference in New Issue View Git Blame Copy Permalink