19 lines
642 B
C++
19 lines
642 B
C++
#pragma once
|
|
#include "pose_engine.hpp"
|
|
#include "gpu_preprocess.hpp"
|
|
namespace cms {
|
|
class Detector {
|
|
public:
|
|
Detector(std::filesystem::path const& path, fangjia::TrtLogger& logger, float confidence, int maxPeople);
|
|
std::vector<fangjia::Box> infer(cv::Mat const& frame,cudaStream_t stream,int filterHeight=0);
|
|
private:
|
|
fangjia::TrtEngine engine_;
|
|
float confidence_;
|
|
int maxPeople_, width_, height_;
|
|
std::string input_,output_;
|
|
std::unique_ptr<YoloPreprocess> preprocess_;
|
|
std::vector<float> outputHost_;
|
|
};
|
|
cv::Mat maskOtherPeople(cv::Mat const& frame,std::vector<fangjia::Box> const& boxes,size_t target);
|
|
}
|