50 lines
1.1 KiB
C++
50 lines
1.1 KiB
C++
#pragma once
|
|
|
|
#include "pipeline_types.hpp"
|
|
#include "trt_engine.hpp"
|
|
|
|
#include <opencv2/core.hpp>
|
|
|
|
#include <string>
|
|
|
|
namespace fangjia {
|
|
|
|
class RtmwEngine {
|
|
public:
|
|
RtmwEngine(
|
|
std::filesystem::path const& enginePath,
|
|
TrtLogger& logger,
|
|
int maxSourceWidth,
|
|
int maxSourceHeight,
|
|
int maxPeople = 3);
|
|
~RtmwEngine();
|
|
|
|
RtmwEngine(RtmwEngine const&) = delete;
|
|
RtmwEngine& operator=(RtmwEngine const&) = delete;
|
|
|
|
std::vector<PersonPose> infer(
|
|
cv::Mat const& frame,
|
|
std::vector<Box> const& boxes,
|
|
cudaStream_t stream,
|
|
float* preprocessMs = nullptr,
|
|
float* inferenceMs = nullptr,
|
|
float* postprocessMs = nullptr);
|
|
|
|
void warmup(cudaStream_t stream);
|
|
|
|
private:
|
|
TrtEngine engine_;
|
|
int maxSourceWidth_{};
|
|
int maxSourceHeight_{};
|
|
int maxPeople_{};
|
|
std::uint8_t* sourceDevice_{};
|
|
std::uint8_t* cropsDevice_{};
|
|
float* networkDevice_{};
|
|
std::vector<double> centersHost_;
|
|
std::vector<double> scalesHost_;
|
|
std::vector<float> simccXHost_;
|
|
std::vector<float> simccYHost_;
|
|
};
|
|
|
|
} // namespace fangjia
|