60 lines
3.6 KiB
C++
60 lines
3.6 KiB
C++
#include "detector.hpp"
|
|
#include <opencv2/imgproc.hpp>
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
namespace cms {
|
|
using fangjia::Box;
|
|
Detector::Detector(std::filesystem::path const& path,fangjia::TrtLogger& logger,float confidence,int maxPeople):engine_(path,logger),confidence_(confidence),maxPeople_(maxPeople) {
|
|
if(engine_.inputs().size()!=1||engine_.outputs().size()!=1) throw std::runtime_error("Expected YOLO raw ONNX output (one input/output)");
|
|
input_=engine_.inputs()[0];output_=engine_.outputs()[0];
|
|
auto d=engine_.shape(input_);
|
|
if(d.nbDims!=4||d.d[0]!=1||d.d[1]!=3||d.d[2]<=0||d.d[3]<=0) throw std::runtime_error("Expected fixed YOLO NCHW input");
|
|
height_=d.d[2];width_=d.d[3];
|
|
engine_.prepare();
|
|
d=engine_.shape(output_);
|
|
if(d.nbDims!=3||d.d[0]!=1||d.d[1]!=84) throw std::runtime_error("Expected YOLO output [1,84,N]");
|
|
if(engine_.dataType(input_)!=nvinfer1::DataType::kFLOAT || engine_.dataType(output_)!=nvinfer1::DataType::kFLOAT) throw std::runtime_error("Expected float32 engine I/O");
|
|
preprocess_=std::make_unique<YoloPreprocess>(width_,height_);outputHost_.resize(engine_.bytes(output_)/sizeof(float));
|
|
}
|
|
std::vector<Box> Detector::infer(cv::Mat const& frame,cudaStream_t stream,int filterHeight) {
|
|
double gain=std::min(double(width_)/frame.cols,double(height_)/frame.rows);
|
|
int rw=int(std::nearbyint(frame.cols*gain)),rh=int(std::nearbyint(frame.rows*gain));
|
|
int left=int(std::nearbyint((width_-rw)/2.0-0.1)),top=int(std::nearbyint((height_-rh)/2.0-0.1));
|
|
preprocess_->enqueue(frame,static_cast<float*>(engine_.device(input_)),stream);
|
|
if(!engine_.enqueue(stream)) throw std::runtime_error("YOLO inference failed");
|
|
CUDA_CHECK(cudaMemcpyAsync(outputHost_.data(),engine_.device(output_),engine_.bytes(output_),cudaMemcpyDeviceToHost,stream));
|
|
CUDA_CHECK(cudaStreamSynchronize(stream));
|
|
int n=engine_.shape(output_).d[2];std::vector<Box> candidates;
|
|
for(int i=0;i<n;++i) {
|
|
int cls=0; float score=outputHost_[4*n+i];
|
|
for(int c=1;c<80;++c) if(outputHost_[(4+c)*n+i]>score) {cls=c;score=outputHost_[(4+c)*n+i];}
|
|
if(cls!=0||score<=confidence_||!std::isfinite(score)) continue;
|
|
float x=outputHost_[i],y=outputHost_[n+i],w=outputHost_[2*n+i],h=outputHost_[3*n+i];
|
|
if(!std::isfinite(x)||!std::isfinite(y)||!std::isfinite(w)||!std::isfinite(h)||w<=0||h<=0) continue;
|
|
candidates.push_back({x-w/2,y-h/2,x+w/2,y+h/2,score,0});
|
|
}
|
|
std::stable_sort(candidates.begin(),candidates.end(),[](auto const&a,auto const&b){return a.score>b.score;});
|
|
std::vector<Box> kept;
|
|
for(auto b:candidates) {
|
|
bool suppressed=false;
|
|
for(auto const& a:kept) {
|
|
float inter=std::max(0.F,std::min(a.x2,b.x2)-std::max(a.x1,b.x1))*std::max(0.F,std::min(a.y2,b.y2)-std::max(a.y1,b.y1));
|
|
float uni=(a.x2-a.x1)*(a.y2-a.y1)+(b.x2-b.x1)*(b.y2-b.y1)-inter;
|
|
if(inter/uni>0.30F) {suppressed=true;break;}
|
|
}
|
|
if(!suppressed) kept.push_back(b);
|
|
if(kept.size()==300) break;
|
|
}
|
|
std::vector<Box> filtered;
|
|
for(auto b:kept) {
|
|
b.x1=std::clamp(float((b.x1-left)/gain),0.F,float(frame.cols));b.x2=std::clamp(float((b.x2-left)/gain),0.F,float(frame.cols));
|
|
b.y1=std::clamp(float((b.y1-top)/gain),0.F,float(frame.rows));b.y2=std::clamp(float((b.y2-top)/gain),0.F,float(frame.rows));
|
|
if((b.x1+b.x2)*0.5F>0.80F*frame.cols||b.y2-b.y1<0.20F*(filterHeight>0?filterHeight:frame.rows)) continue;
|
|
filtered.push_back(b);
|
|
}
|
|
if(int(filtered.size())>maxPeople_) filtered.resize(maxPeople_);
|
|
std::stable_sort(filtered.begin(),filtered.end(),[](auto const&a,auto const&b){return a.x1<b.x1;});
|
|
return filtered;
|
|
}
|
|
}
|