diff --git a/c/segNet.cpp b/c/segNet.cpp index 4943c2ed0..7670a29f5 100644 --- a/c/segNet.cpp +++ b/c/segNet.cpp @@ -735,6 +735,19 @@ bool segNet::Process( void* image, uint32_t width, uint32_t height, imageFormat 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 21f702202..126de5ff9 100644 --- a/c/segNet.h +++ b/c/segNet.h @@ -234,6 +234,16 @@ 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 colorized segmentation mask. */ template bool Mask( T* output, uint32_t width, uint32_t height, FilterMode filter=FILTER_LINEAR ) { return Mask((void*)output, width, height, imageFormatFromType(), filter); }