From b1396e75ecf230b816d1494bd42886aad5848692 Mon Sep 17 00:00:00 2001
From: luocai <luocai@reconova.com>
Date: Thu, 14 Mar 2024 21:41:54 +0800
Subject: [PATCH] Add geo.

---
 Main/DetectAlgorithm.cpp | 35 ++++++++++++++++++++++++++++++-----
 Main/DetectAlgorithm.h   | 16 +++++++++++++++-
 Main/Geometry.cpp        |  1 +
 Main/Geometry.h          | 14 ++++++++++++++
 Main/PedestrianInfo.cpp  |  1 +
 Main/PedestrianInfo.h    | 15 +++++++++++++++
 6 files changed, 76 insertions(+), 6 deletions(-)
 create mode 100644 Main/Geometry.cpp
 create mode 100644 Main/Geometry.h
 create mode 100644 Main/PedestrianInfo.cpp
 create mode 100644 Main/PedestrianInfo.h

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 <filesystem>
 #include <thread>
 
-DetectAlgorithm::DetectAlgorithm() {
+class DetectAlgorithmPrivate {
+public:
+    std::vector<io_TrackData> 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<PedestrianRect> pedestrians;
-    std::vector<io_TrackData> 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 <array>
 #include <cstdint>
+#include <vector>
+
+class DetectAlgorithmPrivate;
 
 class DetectAlgorithm {
 public:
+    enum TrackState : int {
+        Remove = -2,
+        Hibernate = -1,
+        PreTrack = 0,
+        Confirmed = 1,
+    };
+    class Result {
+    public:
+        std::vector<int> 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<uint8_t, DetectWidth * DetectHeight * 3> 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 <cstdint>
+
+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