From f94a6b6b53c0cf8bcccabc1c440af9fc3abebec4 Mon Sep 17 00:00:00 2001
From: Jithendar Reddy Chengalapattu <jithendar.rc@embedur.com>
Date: Fri, 10 Jul 2026 15:06:43 +0530
Subject: [PATCH] Update ML output SDS to publish all class confidence scores

---
 .../algorithm/algorithm_config.h              |  3 +-
 .../algorithm/arm_executor_runner.cc          | 33 +++++++++++++++++--
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/RockPaperScissors/AppKit-E8_USB/algorithm/algorithm_config.h b/RockPaperScissors/AppKit-E8_USB/algorithm/algorithm_config.h
index 8bcdf7b..8adc296 100644
--- a/RockPaperScissors/AppKit-E8_USB/algorithm/algorithm_config.h
+++ b/RockPaperScissors/AppKit-E8_USB/algorithm/algorithm_config.h
@@ -20,6 +20,7 @@
 #define ALGORITHM_CONFIG_H_
 
 #include "config_ml_model.h"            // ML model configuration
+#include "model_config.h"
 
 // Input Data block size, in bytes
 #ifndef ALGO_DATA_IN_BLOCK_SIZE
@@ -28,7 +29,7 @@
 
 // Output Data block size, in bytes
 #ifndef ALGO_DATA_OUT_BLOCK_SIZE
-#define ALGO_DATA_OUT_BLOCK_SIZE        (50U)
+#define ALGO_DATA_OUT_BLOCK_SIZE        (MODEL_NUM_CLASSES * sizeof(float))
 #endif
 
 #endif
diff --git a/RockPaperScissors/AppKit-E8_USB/algorithm/arm_executor_runner.cc b/RockPaperScissors/AppKit-E8_USB/algorithm/arm_executor_runner.cc
index 2f423b1..509e6ce 100644
--- a/RockPaperScissors/AppKit-E8_USB/algorithm/arm_executor_runner.cc
+++ b/RockPaperScissors/AppKit-E8_USB/algorithm/arm_executor_runner.cc
@@ -178,6 +178,8 @@ bool classify_object = false;
 
 output_label_t output_label;
 
+float class_probs[MODEL_NUM_CLASSES] = {0.0};
+
 constexpr int H = IMAGE_HEIGHT;
 
 constexpr int W = IMAGE_WIDTH;
@@ -490,6 +492,31 @@ static void FormatShortLabel(const char* label, char* short_label,
     short_label[copy_len] = '\0';
 }
 
+/**
+ * \brief Apply softmax function to convert logits to probabilities
+ * \param[in]  logits Pointer to logit values array
+ * \param[out] probs  Pointer to output probability array
+ * \param[in]  n      Number of elements
+ */
+void softmax(const float* logits, float* probs, int n) {
+    float max_val = logits[0];
+    for (int i = 1; i < n; i++) {
+        if (logits[i] > max_val) {
+            max_val = logits[i];
+        }
+    }
+
+    float sum = 0.0f;
+    for (int i = 0; i < n; i++) {
+        probs[i] = expf(logits[i] - max_val);
+        sum += probs[i];
+    }
+
+    for (int i = 0; i < n; i++) {
+        probs[i] /= sum;
+    }
+}
+
 /**
  * \brief Prepare image input tensor for SqueezeNet model
  * \param[in] image Pointer to input image data (HWC RGB format)
@@ -846,9 +873,9 @@ void print_outputs(RunnerContext& ctx)
             printf("Number of classes: %d, expected: %d\n", numel, MODEL_NUM_CLASSES);
         }
 
-        static float probs[MODEL_NUM_CLASSES];
         float confidence = 0.0f;
         uint16_t predicted_idx = -1;
+        softmax(logits, class_probs, numel);
         status_t postprocess_result = postprocess((float*)logits, &result);
 
         if(postprocess_result != STATUS_OK){
@@ -926,8 +953,8 @@ void postprocess(RunnerContext& ctx, uint8_t* img_buf,
     print_outputs(ctx);
 
     /* Copy shortened label plus confidence into caller's output buffer */
-    if (out_num >= sizeof(output_label_t)) {
-        memcpy(out_buf, &output_label, sizeof(output_label_t));
+    if (out_num >= sizeof(class_probs)) {
+        memcpy(out_buf, class_probs, sizeof(class_probs));
     }
 
     /* Only draw if label is valid */