From 95a85e25f1fbc7324a48d99ba3afa989fe8a79f6 Mon Sep 17 00:00:00 2001
From: kmfrick <kmfrick98@gmail.com>
Date: Sun, 17 May 2020 09:04:19 +0200
Subject: [PATCH] segNet: Add zero-copy class score getter

---
 c/segNet.cpp | 13 +++++++++++++
 c/segNet.h   |  9 +++++++++
 2 files changed, 22 insertions(+)

diff --git a/c/segNet.cpp b/c/segNet.cpp
index 369962528..deb0590ed 100644
--- a/c/segNet.cpp
+++ b/c/segNet.cpp
@@ -661,6 +661,19 @@ bool segNet::Process( float* rgba, uint32_t width, uint32_t height, const char*
 	return true;
 }
 
+bool segNet::GetClassScores( float** class_scores, uint32_t* width, uint32_t* height, uint32_t* num_classes )
+{
+	if ( mOutputs[0].CPU == NULL )
+	{
+		return false;
+	}
+	*class_scores = (float*) mOutputs[0].CPU;
+	*width = DIMS_W(mOutputs[0].dims);
+	*height = DIMS_H(mOutputs[0].dims);
+	*num_classes = DIMS_C(mOutputs[0].dims);
+	return true;
+}
+		
 
 // argmax classification
 bool segNet::classify( const char* ignore_class )
diff --git a/c/segNet.h b/c/segNet.h
index ce15f238e..96eccbba0 100644
--- a/c/segNet.h
+++ b/c/segNet.h
@@ -188,6 +188,15 @@ class segNet : public tensorNet
 	 */
 	bool Process( float* input, uint32_t width, uint32_t height, const char* ignore_class="void" );
 
+	/**
+	 * Return per-pixel class probabilities, as well as number of classes and size of the output layer.
+	 * Does not perform a memory copy.
+	 * @param class_scores float pointer to destination array
+	 * @param width pointer to the variable that will hold the output layer width
+	 * @param height pointer to the variable that will hold the output layer height
+	 * @param num_classes pointer to the variable that will hold the number of classes
+	 */
+	 bool GetClassScores( float** class_scores, uint32_t* width, uint32_t* height, uint32_t* num_classes );
 	/**
 	 * Produce a grayscale binary segmentation mask, where the pixel values
 	 * correspond to the class ID of the corresponding class type.