From 1ca14a9ad0ab00176f33e2116718cc4eea0b0c05 Mon Sep 17 00:00:00 2001
From: Myyura <zz940521@gmail.com>
Date: Sun, 31 Mar 2024 05:58:57 +0800
Subject: [PATCH] onnxruntime cpp yolo-cls fp16 fix (#9412)

---
 examples/YOLOv8-ONNXRuntime-CPP/inference.cpp | 14 +++++++++++++-
 examples/YOLOv8-ONNXRuntime-CPP/inference.h   |  1 +
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/examples/YOLOv8-ONNXRuntime-CPP/inference.cpp b/examples/YOLOv8-ONNXRuntime-CPP/inference.cpp
index 0bf07ded..88aeae37 100644
--- a/examples/YOLOv8-ONNXRuntime-CPP/inference.cpp
+++ b/examples/YOLOv8-ONNXRuntime-CPP/inference.cpp
@@ -301,12 +301,24 @@ char* YOLO_V8::TensorProcess(clock_t& starttime_1, cv::Mat& iImg, N& blob, std::
         break;
     }
     case YOLO_CLS:
+    case YOLO_CLS_HALF:
     {
+        cv::Mat rawData;
+        if (modelType == YOLO_CLS) {
+            // FP32
+            rawData = cv::Mat(1, this->classes.size(), CV_32F, output);
+        } else {
+            // FP16
+            rawData = cv::Mat(1, this->classes.size(), CV_16F, output);
+            rawData.convertTo(rawData, CV_32F);
+        }
+        float *data = (float *) rawData.data;
+
         DL_RESULT result;
         for (int i = 0; i < this->classes.size(); i++)
         {
             result.classId = i;
-            result.confidence = output[i];
+            result.confidence = data[i];
             oResult.push_back(result);
         }
         break;
diff --git a/examples/YOLOv8-ONNXRuntime-CPP/inference.h b/examples/YOLOv8-ONNXRuntime-CPP/inference.h
index 3174ae93..3a9d029c 100644
--- a/examples/YOLOv8-ONNXRuntime-CPP/inference.h
+++ b/examples/YOLOv8-ONNXRuntime-CPP/inference.h
@@ -29,6 +29,7 @@ enum MODEL_TYPE
     //FLOAT16 MODEL
     YOLO_DETECT_V8_HALF = 4,
     YOLO_POSE_V8_HALF = 5,
+    YOLO_CLS_HALF = 6
 };