diff --git a/Main/DetectAlgorithm.cpp b/Main/DetectAlgorithm.cpp index 622db6c..a3ccbab 100644 --- a/Main/DetectAlgorithm.cpp +++ b/Main/DetectAlgorithm.cpp @@ -5,26 +5,51 @@ #include #include -DetectAlgorithm::DetectAlgorithm() { +class DetectAlgorithmPrivate { +public: + std::vector tracks; +}; + +DetectAlgorithm::DetectAlgorithm() : m_d(new DetectAlgorithmPrivate()) { } DetectAlgorithm::~DetectAlgorithm() { if (m_handle != nullptr) { ds_pedestrian_hisi_release(&m_handle); } + if (m_d != nullptr) { + delete m_d; + } } -void DetectAlgorithm::detect(const uint8_t *nv21ImageData, uint64_t frameIndex) { +DetectAlgorithm::Result DetectAlgorithm::detect(const uint8_t *nv21ImageData, uint64_t frameIndex) { + DetectAlgorithm::Result ret; if (m_handle == nullptr) { // 一定得在这里执行 initialize(); } ImageUtilities::NV21ToBGR24(nv21ImageData, m_rgbImageBuffer.data(), DetectWidth, DetectHeight); std::vector pedestrians; - std::vector tracks; + ds_pedestrian_det_hisi(m_handle, m_rgbImageBuffer.data(), pedestrians); - ds_pedestrian_track_hisi(m_handle, m_rgbImageBuffer.data(), frameIndex, pedestrians, tracks); - LOG(info) << "frame: " << frameIndex << ", pedestrians: " << pedestrians.size() << ", tracks: " << tracks.size(); + ds_pedestrian_track_hisi(m_handle, m_rgbImageBuffer.data(), frameIndex, pedestrians, m_d->tracks); + LOG(info) << "frame: " << frameIndex << ", pedestrians: " << pedestrians.size() + << ", tracks: " << m_d->tracks.size(); + for (auto iterator = m_d->tracks.cbegin(); iterator != m_d->tracks.cend();) { + if (iterator->track_state == TrackState::Remove) { + ret.leaveTrackers.push_back(iterator->track_id); + iterator = m_d->tracks.erase(iterator); + continue; + } else if (iterator->track_state == TrackState::Confirmed) { + LOG(info) << iterator->prediction.body.conf << " " << iterator->prediction.head.conf; + LOG(info) << iterator->prediction.body.state[0] << " " << iterator->prediction.body.state[1] << " " + << iterator->prediction.body.state[2] << " " << iterator->prediction.body.state[3]; + LOG(info) << iterator->prediction.head.state[0] << " " << iterator->prediction.head.state[1] << " " + << iterator->prediction.head.state[2] << " " << iterator->prediction.head.state[3]; + } + ++iterator; + } + return ret; } void DetectAlgorithm::initialize() { diff --git a/Main/DetectAlgorithm.h b/Main/DetectAlgorithm.h index cdc1158..f685634 100644 --- a/Main/DetectAlgorithm.h +++ b/Main/DetectAlgorithm.h @@ -3,17 +3,31 @@ #include #include +#include + +class DetectAlgorithmPrivate; class DetectAlgorithm { public: + enum TrackState : int { + Remove = -2, + Hibernate = -1, + PreTrack = 0, + Confirmed = 1, + }; + class Result { + public: + std::vector leaveTrackers; + }; constexpr static uint32_t DetectWidth = 576; constexpr static uint32_t DetectHeight = 320; DetectAlgorithm(); ~DetectAlgorithm(); - void detect(const uint8_t *nv21ImageData, uint64_t frameIndex); + Result detect(const uint8_t *nv21ImageData, uint64_t frameIndex); void initialize(); private: + DetectAlgorithmPrivate *m_d = nullptr; void *m_handle = nullptr; std::array m_rgbImageBuffer; }; diff --git a/Main/Geometry.cpp b/Main/Geometry.cpp new file mode 100644 index 0000000..954a1f8 --- /dev/null +++ b/Main/Geometry.cpp @@ -0,0 +1 @@ +#include "Geometry.h" \ No newline at end of file diff --git a/Main/Geometry.h b/Main/Geometry.h new file mode 100644 index 0000000..f546a3a --- /dev/null +++ b/Main/Geometry.h @@ -0,0 +1,14 @@ +#ifndef __GEOMETRY_H__ +#define __GEOMETRY_H__ + +#include + +class Rectangle { +public: + int32_t left = 0; + int32_t right = 0; + int32_t top = 0; + int32_t bottom = 0; +}; + +#endif // __GEOMETRY_H__ \ No newline at end of file diff --git a/Main/PedestrianInfo.cpp b/Main/PedestrianInfo.cpp new file mode 100644 index 0000000..1fcc83a --- /dev/null +++ b/Main/PedestrianInfo.cpp @@ -0,0 +1 @@ +#include "PedestrianInfo.h" \ No newline at end of file diff --git a/Main/PedestrianInfo.h b/Main/PedestrianInfo.h new file mode 100644 index 0000000..6df9e32 --- /dev/null +++ b/Main/PedestrianInfo.h @@ -0,0 +1,15 @@ +#ifndef __PEDESTRIANINFO_H__ +#define __PEDESTRIANINFO_H__ + +#include "Geometry.h" + +class PedestrianInfo { +public: + int32_t trackId = -1; + uint64_t frameIndex = 0; + float confidence = 0; + Rectangle head; + Rectangle body; +}; + +#endif // __PEDESTRIANINFO_H__ \ No newline at end of file