109 lines
2.7 KiB
C++
109 lines
2.7 KiB
C++
#pragma once
|
|
|
|
#include "pipeline_types.hpp"
|
|
#include "trt_engine.hpp"
|
|
|
|
#include <opencv2/core.hpp>
|
|
|
|
#include <deque>
|
|
#include <array>
|
|
#include <optional>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
|
|
namespace fangjia {
|
|
|
|
class BoxTracker {
|
|
public:
|
|
std::vector<int> update(std::vector<Box>& boxes, int frameIndex, int width, int height);
|
|
|
|
private:
|
|
struct Track { Box box; int lastFrame{}; };
|
|
int nextId_{1};
|
|
std::unordered_map<int, Track> tracks_;
|
|
};
|
|
|
|
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_;
|
|
};
|
|
|
|
class PoseFusionTracker {
|
|
public:
|
|
PoseFrameResult update(
|
|
std::vector<PersonPose> const& raw,
|
|
std::vector<PersonPose> const& isolated,
|
|
int frameIndex);
|
|
|
|
private:
|
|
struct Vec2 { double x{}; double y{}; };
|
|
struct SideFusion {
|
|
std::array<Vec2, 3> limbPrevious{};
|
|
std::array<Vec2, 3> limbVelocity{};
|
|
std::array<Vec2, 21> handPrevious{};
|
|
std::array<Vec2, 21> handVelocity{};
|
|
bool hasLimbPrevious{};
|
|
bool hasLimbVelocity{};
|
|
bool hasHandPrevious{};
|
|
bool hasHandVelocity{};
|
|
std::string limbSource;
|
|
std::string handSource;
|
|
};
|
|
struct Endpoint {
|
|
bool hasWrist{};
|
|
Vec2 wrist{};
|
|
Vec2 velocity{};
|
|
std::array<Vec2, 21> handShape{};
|
|
bool hasHandShape{};
|
|
int lastObservedFrame{-1};
|
|
int missedFrames{};
|
|
bool hasLastElbow{};
|
|
Vec2 lastElbow{};
|
|
bool hasLastDirection{};
|
|
Vec2 lastDirection{};
|
|
};
|
|
struct PersonState {
|
|
std::array<SideFusion, 2> fusion;
|
|
std::array<Endpoint, 2> endpoints;
|
|
std::array<std::deque<std::array<double, 2>>, 2> boneSamples;
|
|
std::array<Vec2, 133> smoothed{};
|
|
bool hasSmoothed{};
|
|
int lastFrame{-1};
|
|
};
|
|
std::unordered_map<int, PersonState> states_;
|
|
};
|
|
|
|
} // namespace fangjia
|