[update] add point number function parameter to YOLO-Pose

Author: Abandon-ht

projects/llm_framework/main_yolo/src/EngineWrapper.cpp

@@ -311,8 +311,8 @@ static const std::vector<std::vector<uint8_t>> SKELETON = {
     {8, 10},  {9, 11},  {2, 3},   {1, 2},   {1, 3},   {2, 4},  {3, 5},  {4, 6}, {5, 7}};
 
 void post_process(AX_ENGINE_IO_INFO_T* io_info, AX_ENGINE_IO_T* io_data, const cv::Mat& mat, int& input_w, int& input_h,
-                  int& cls_num, float& prob_threshold, float& nms_threshold, std::vector<detection::Object>& objects,
-                  std::string& model_type)
+                  int& cls_num, int& point_num, float& prob_threshold, float& nms_threshold,
+                  std::vector<detection::Object>& objects, std::string& model_type)
 {
     // std::vector<detection::Object> objects;
     std::vector<detection::Object> proposals;
@@ -352,7 +352,7 @@ void post_process(AX_ENGINE_IO_INFO_T* io_info, AX_ENGINE_IO_T* io_data, const c
             auto feat_kps_ptr = output_kps_ptr[i];
             int32_t stride    = (1 << i) * 8;
             detection::generate_proposals_yolov8_pose_native(stride, feat_ptr, feat_kps_ptr, prob_threshold, proposals,
-                                                             input_h, input_w, 17, cls_num);
+                                                             input_h, input_w, point_num, cls_num);
         }
         detection::get_out_bbox_kps(proposals, objects, nms_threshold, input_h, input_w, mat.rows, mat.cols);
         // detection::draw_keypoints(mat, objects, KPS_COLORS, LIMB_COLORS, SKELETON, "yolo11_pose_out");
@@ -368,10 +368,12 @@ void post_process(AX_ENGINE_IO_INFO_T* io_info, AX_ENGINE_IO_T* io_data, const c
     }
 }
 
-int EngineWrapper::Post_Process(cv::Mat& mat, int& input_w, int& input_h, int& cls_num, float& pron_threshold,
-                                float& nms_threshold, std::vector<detection::Object>& objects, std::string& model_type)
+int EngineWrapper::Post_Process(cv::Mat& mat, int& input_w, int& input_h, int& cls_num, int& point_num,
+                                float& pron_threshold, float& nms_threshold, std::vector<detection::Object>& objects,
+                                std::string& model_type)
 {
-    post_process(m_io_info, &m_io, mat, input_w, input_h, cls_num, pron_threshold, nms_threshold, objects, model_type);
+    post_process(m_io_info, &m_io, mat, input_w, input_h, cls_num, point_num, pron_threshold, nms_threshold, objects,
+                 model_type);
     return 0;
 }
 

projects/llm_framework/main_yolo/src/EngineWrapper.hpp

@@ -49,8 +49,8 @@ class EngineWrapper {
 
     int RunSync();
 
-    int Post_Process(cv::Mat& mat, int& input_w, int& input_, int& cls_num, float& pron_threshold, float& nms_threshold,
-                     std::vector<detection::Object>& objects, std::string& model_type);
+    int Post_Process(cv::Mat& mat, int& input_w, int& input_, int& cls_num, int& point_num, float& pron_threshold,
+                     float& nms_threshold, std::vector<detection::Object>& objects, std::string& model_type);
 
     int GetOutput(void* pOutput, int index);
 

projects/llm_framework/main_yolo/src/main.cpp

@@ -31,6 +31,7 @@ typedef struct {
     int img_h            = 640;
     int img_w            = 640;
     int cls_num          = 80;
+    int point_num        = 17;
     float pron_threshold = 0.45f;
     float nms_threshold  = 0.45;
 } yolo_config;
@@ -115,6 +116,7 @@ class llm_task {
             CONFIG_AUTO_SET(file_body["mode_param"], nms_threshold);
             CONFIG_AUTO_SET(file_body["mode_param"], cls_name);
             CONFIG_AUTO_SET(file_body["mode_param"], cls_num);
+            CONFIG_AUTO_SET(file_body["mode_param"], point_num);
             CONFIG_AUTO_SET(file_body["mode_param"], model_type);
             mode_config_.yolo_model = base_model + mode_config_.yolo_model;
             yolo_                   = std::make_unique<EngineWrapper>();
@@ -207,19 +209,19 @@ class llm_task {
             }
             std::vector<detection::Object> objects;
             yolo_->Post_Process(img_mat, mode_config_.img_w, mode_config_.img_h, mode_config_.cls_num,
-                                mode_config_.pron_threshold, mode_config_.nms_threshold, objects,
-                                mode_config_.model_type);
+                                mode_config_.point_num, mode_config_.pron_threshold, mode_config_.nms_threshold,
+                                objects, mode_config_.model_type);
             std::vector<nlohmann::json> yolo_output;
             for (size_t i = 0; i < objects.size(); i++) {
                 const detection::Object &obj = objects[i];
                 nlohmann::json output;
                 output["class"]      = mode_config_.cls_name[obj.label];
                 output["confidence"] = format_float(obj.prob, 2);
                 output["bbox"]       = nlohmann::json::array();
-                output["bbox"].push_back(format_float(obj.rect.x, 0));
-                output["bbox"].push_back(format_float(obj.rect.y, 0));
-                output["bbox"].push_back(format_float(obj.rect.x + obj.rect.width, 0));
-                output["bbox"].push_back(format_float(obj.rect.y + obj.rect.height, 0));
+                output["bbox"].push_back(format_float(obj.rect.x, 2));
+                output["bbox"].push_back(format_float(obj.rect.y, 2));
+                output["bbox"].push_back(format_float(obj.rect.x + obj.rect.width, 2));
+                output["bbox"].push_back(format_float(obj.rect.y + obj.rect.height, 2));
                 if (mode_config_.model_type == "segment") {
                     std::vector<std::string> formatted_mask_feat;
                     for (const auto &mask : obj.mask_feat) {