Adding class_mask scores output to SegNet

[pq53ui]

I would really like to get the scores for the segmentation mask. I don't care about the entire ndarray, just the scores for the most probable class.
Looking at this issue I understand there is no way to get that at the moment in python.

So.. I thought I'd add it myself. But since it's been a while since I did any C++ coding I ran into some issues.

I made some minor changes to segNet.cpp anad segNet.h
I want to add ScoreMask to SegNet and I thought I could just save the scores alongside argmax calculation in the classify method and then expose it through a new function. Compiles fine, but can't get the attribute to show in the python SegNet object.

Changes to segNet.cpp

// constructor
segNet::segNet() : tensorNet()
{
	mLastInputImg    = NULL;
	mLastInputWidth  = 0;
	mLastInputHeight = 0;
	mLastInputFormat = IMAGE_UNKNOWN;

	mColorsAlphaSet = NULL;
	mClassColors    = NULL;
	mClassMap       = NULL;
	mClassMapScore  = NULL;  // **New variable** 
}

// destructor
segNet::~segNet()
{
	CUDA_FREE_HOST(mClassColors);
	CUDA_FREE_HOST(mClassMap);
	CUDA_FREE_HOST(mClassMapScore);
	
	if( mColorsAlphaSet != NULL )
	{
		free(mColorsAlphaSet);
		mColorsAlphaSet = NULL;
	}
}
....

// argmax classification
bool segNet::classify( const char* ignore_class )
{
	// retrieve scores
	float* scores = mOutputs[0].CPU;

	...

	// find the argmax-classified class of each tile
	uint8_t* classMap = mClassMap;
	float* classMapScore = mClassMapScore;   // **Retrive this new variable**
        ...
			classMap[y * s_w + x] = c_max;
			classMapScore[y * s_w + x] = p_max;  // **assign score**
			//printf("(%u, %u) -> class %i\n", x, y, (uint32_t)c_max);
		}
	}

	return true;
}

// Create
segNet* segNet::Create( const char* prototxt, const char* model, const char* labels_path, const char* colors_path, 
				    const char* input_blob, const char* output_blob, uint32_t maxBatchSize,
				    precisionType precision, deviceType device, bool allowGPUFallback )
{
     ....
     if( !cudaAllocMapped((void**)&net->mClassMapScore, s_w * s_h * sizeof(uint8_t)) )
		return NULL;

    ....
}

// Mostly copied from ::Mask
bool segNet::ScoreMask(float* output, uint32_t out_width, uint32_t out_height)
{
	if( !output || out_width == 0 || out_height == 0 )
	{
		LogError(LOG_TRT "segNet::ScoreMask( 0x%p, %u, %u ) -> invalid parameters\n", output, out_width, out_height); 
		return false;
	}

	PROFILER_BEGIN(PROFILER_VISUALIZE);

	//retrive score map
	float* classMapScore = mClassMapScore;

	const int s_w = DIMS_W(mOutputs[0].dims);
	const int s_h = DIMS_H(mOutputs[0].dims);

	if( out_width == s_w && out_height == s_h )
	{
		memcpy(output, classMapScore, s_w * s_h * sizeof(float));
	}
	else
	{
		const float s_x = float(s_w) / float(out_width);
		const float s_y = float(s_h) / float(out_height);

		// overlay pixels onto original
		for( uint32_t y=0; y < out_height; y++ )
		{
			for( uint32_t x=0; x < out_width; x++ )
			{
				const int cx = float(x) * s_x;
				const int cy = float(y) * s_y;

				// get the score of this cell
				const float score = classMapScore[cy * s_w + cx];

				// output the pixel
				output[y * out_width + x] = score;
			}
		}
	}

	PROFILER_END(PROFILER_VISUALIZE);
	return true;
}

and then segNet.h

	/**
	 * Produce a score mask.
	 */
	template<typename T> bool ScoreMask( T* output, uint32_t width, uint32_t height)				{ return ScoreMask((void*)output, width, height); }
	
	/**
	 * Produce a grayscale float mask, where the pixel values
	 * correspond to the score.
	 */
	bool ScoreMask(float* output, uint32_t out_width, uint32_t out_height);

PySegNet.cpp

\\ added to the PySegNet_Methods[]
...
{ "ScoreMask", (PyCFunction)PySegNet_ScoreMask, METH_VARARGS|METH_KEYWORDS, DOC_SCOREMASK},
...

static PyObject* PySegNet_ScoreMask( PySegNet_Object* self, PyObject* args, PyObject *kwds )
{
	if( !self || !self->net )
	{
		PyErr_SetString(PyExc_Exception, LOG_PY_INFERENCE "segNet invalid object instance");
		return NULL;
	}
	
	// parse arguments
	PyObject* capsule = NULL;

	int width = 0;
	int height = 0;

	const char* filter_str = "linear";
	const char* format_str = "rgba32f";

	static char* kwlist[] = {"image", "width", "height", "filter_mode", "format", NULL};

	if( !PyArg_ParseTupleAndKeywords(args, kwds, "O|iiss", kwlist, &capsule, &width, &height, &filter_str, &format_str))
		return NULL;


	// parse format string
	imageFormat format = imageFormatFromStr(format_str);

	// get pointer to image data
	void* ptr = PyCUDA_GetImage(capsule, &width, &height, &format);

	if( !ptr )
		return NULL;

	bool result = false;
	Py_BEGIN_ALLOW_THREADS

	result = self->net->ScoreMask((float*)ptr, width, height);
	
	Py_END_ALLOW_THREADS
	
	if( !result )
	{
		PyErr_SetString(PyExc_Exception, LOG_PY_INFERENCE "segNet.ScoreMask() encountered an error generating the segmentation mask");
		return NULL;
	}

	Py_RETURN_NONE;
}

And then did the cmake and make. Everything compiles, no errors.
But when I try to call the newly created ScoreMask I get
*** AttributeError: 'jetson.inference.segNet' object has no attribute 'ScoreMask'
error.

What am I missing? Any help would be much appreciated.

[pq53ui]

I am an idiot.

Didn't run make install.
Closing the issue and leaving this here. Might help someone. Either to get the score mask or to run make install :)

[pq53ui]

But, I'd still like someone to tell me if that is the right way to do it.
Especially with the float conversion from C++ to Python - should this be fine?