added 'lines' mode to detectNet overlay

shoaib42 · Feb 1, 2022 · cf9376f · cf9376f
1 parent b75ea72
commit cf9376f
Show file tree

Hide file tree

Showing 4 changed files with 75 additions and 13 deletions.
diff --git a/c/detectNet.cpp b/c/detectNet.cpp
@@ -25,6 +25,7 @@
 
 #include "cudaMappedMemory.h"
 #include "cudaFont.h"
+#include "cudaDraw.h"
 
 #include "commandLine.h"
 #include "filesystem.h"
@@ -48,6 +49,7 @@ detectNet::detectNet( float meanPixel ) : tensorNet()
 {
 	mCoverageThreshold = DETECTNET_DEFAULT_THRESHOLD;
 	mMeanPixel         = meanPixel;
+	mLineWidth         = 2.0f;
 	mNumClasses        = 0;
 
 	mClassColors[0]   = NULL; // cpu ptr
@@ -1072,7 +1074,22 @@ bool detectNet::Overlay( void* input, void* output, uint32_t width, uint32_t hei
 		if( CUDA_FAILED(cudaDetectionOverlay(input, output, width, height, format, detections, numDetections, (float4*)mClassColors[1])) )
 			return false;
 	}
+
+	// bounding box lines
+	if( flags & OVERLAY_LINES )
+	{
+		for( uint32_t n=0; n < numDetections; n++ )
+		{
+			const Detection* d = detections + n;
+			const float4& color = ((float4*)mClassColors[0])[d->ClassID];
 
+			CUDA(cudaDrawLine(input, output, width, height, format, d->Left, d->Top, d->Right, d->Top, color, mLineWidth));
+			CUDA(cudaDrawLine(input, output, width, height, format, d->Right, d->Top, d->Right, d->Bottom, color, mLineWidth));
+			CUDA(cudaDrawLine(input, output, width, height, format, d->Left, d->Bottom, d->Right, d->Bottom, color, mLineWidth));
+			CUDA(cudaDrawLine(input, output, width, height, format, d->Left, d->Top, d->Left, d->Bottom, color, mLineWidth));
+		}
+	}
+
 	// class label overlay
 	if( (flags & OVERLAY_LABEL) || (flags & OVERLAY_CONFIDENCE) )
 	{
@@ -1166,6 +1183,8 @@ uint32_t detectNet::OverlayFlagsFromStr( const char* str_user )
 			flags |= OVERLAY_LABEL;
 		else if( strcasecmp(token, "conf") == 0 || strcasecmp(token, "confidence") == 0 )
 			flags |= OVERLAY_CONFIDENCE;
+		else if( strcasecmp(token, "line") == 0 || strcasecmp(token, "lines") == 0 )
+			flags |= OVERLAY_LINES;
 		else if( strcasecmp(token, "default") == 0 )
 			flags |= OVERLAY_DEFAULT;
 

diff --git a/c/detectNet.h b/c/detectNet.h
@@ -84,7 +84,7 @@
 		  "  --threshold=THRESHOLD minimum threshold for detection (default is 0.5)\n"							\
             "  --alpha=ALPHA         overlay alpha blending value, range 0-255 (default: 120)\n"					\
 		  "  --overlay=OVERLAY     detection overlay flags (e.g. --overlay=box,labels,conf)\n"					\
-		  "                        valid combinations are:  'box', 'labels', 'conf', 'none'\n"					\
+		  "                        valid combinations are:  'box', 'lines', 'labels', 'conf', 'none'\n"			\
 		  "  --profile             enable layer profiling in TensorRT\n\n"
 
 
@@ -172,10 +172,11 @@ class detectNet : public tensorNet
 	enum OverlayFlags
 	{
 		OVERLAY_NONE       = 0,			/**< No overlay. */
-		OVERLAY_BOX        = (1 << 0),	/**< Overlay the object bounding boxes */
+		OVERLAY_BOX        = (1 << 0),	/**< Overlay the object bounding boxes (filled) */
 		OVERLAY_LABEL 	    = (1 << 1),	/**< Overlay the class description labels */
 		OVERLAY_CONFIDENCE = (1 << 2),	/**< Overlay the detection confidence values */
-		OVERLAY_DEFAULT    = OVERLAY_BOX,	/**< The default choice of overlay */
+		OVERLAY_LINES      = (1 << 3),     /**< Overlay the bounding box lines (unfilled) */
+		OVERLAY_DEFAULT    = OVERLAY_BOX|OVERLAY_LABEL|OVERLAY_CONFIDENCE, /**< The default choice of overlay */
 	};
 
 	/**
@@ -319,7 +320,7 @@ class detectNet : public tensorNet
 	 * @param[in]  overlay bitwise OR combination of overlay flags (@see OverlayFlags and @see Overlay()), or OVERLAY_NONE.
 	 * @returns    The number of detected objects, 0 if there were no detected objects, and -1 if an error was encountered.
 	 */
-	template<typename T> int Detect( T* image, uint32_t width, uint32_t height, Detection** detections, uint32_t overlay=OVERLAY_BOX )		{ return Detect((void*)image, width, height, imageFormatFromType<T>(), detections, overlay); }
+	template<typename T> int Detect( T* image, uint32_t width, uint32_t height, Detection** detections, uint32_t overlay=OVERLAY_DEFAULT )		{ return Detect((void*)image, width, height, imageFormatFromType<T>(), detections, overlay); }
 
 	/**
 	 * Detect object locations in an image, into an array of the results allocated by the user.
@@ -331,7 +332,7 @@ class detectNet : public tensorNet
 	 * @param[in]  overlay bitwise OR combination of overlay flags (@see OverlayFlags and @see Overlay()), or OVERLAY_NONE.
 	 * @returns    The number of detected objects, 0 if there were no detected objects, and -1 if an error was encountered.
 	 */
-	template<typename T> int Detect( T* image, uint32_t width, uint32_t height, Detection* detections, uint32_t overlay=OVERLAY_BOX )			{ return Detect((void*)image, width, height, imageFormatFromType<T>(), detections, overlay); }
+	template<typename T> int Detect( T* image, uint32_t width, uint32_t height, Detection* detections, uint32_t overlay=OVERLAY_DEFAULT )		{ return Detect((void*)image, width, height, imageFormatFromType<T>(), detections, overlay); }
 
 	/**
 	 * Detect object locations from an image, returning an array containing the detection results.
@@ -342,7 +343,7 @@ class detectNet : public tensorNet
 	 * @param[in]  overlay bitwise OR combination of overlay flags (@see OverlayFlags and @see Overlay()), or OVERLAY_NONE.
 	 * @returns    The number of detected objects, 0 if there were no detected objects, and -1 if an error was encountered.
 	 */
-	int Detect( void* input, uint32_t width, uint32_t height, imageFormat format, Detection** detections, uint32_t overlay=OVERLAY_BOX );
+	int Detect( void* input, uint32_t width, uint32_t height, imageFormat format, Detection** detections, uint32_t overlay=OVERLAY_DEFAULT );
 
 	/**
 	 * Detect object locations from an image, into an array of the results allocated by the user.
@@ -354,7 +355,7 @@ class detectNet : public tensorNet
 	 * @param[in]  overlay bitwise OR combination of overlay flags (@see OverlayFlags and @see Overlay()), or OVERLAY_NONE.
 	 * @returns    The number of detected objects, 0 if there were no detected objects, and -1 if an error was encountered.
 	 */
-	int Detect( void* input, uint32_t width, uint32_t height, imageFormat format, Detection* detections, uint32_t overlay=OVERLAY_BOX );
+	int Detect( void* input, uint32_t width, uint32_t height, imageFormat format, Detection* detections, uint32_t overlay=OVERLAY_DEFAULT );
 
 	/**
 	 * Detect object locations from an RGBA image, returning an array containing the detection results.
@@ -366,7 +367,7 @@ class detectNet : public tensorNet
 	 * @param[in]  overlay bitwise OR combination of overlay flags (@see OverlayFlags and @see Overlay()), or OVERLAY_NONE.
 	 * @returns    The number of detected objects, 0 if there were no detected objects, and -1 if an error was encountered.
 	 */
-	int Detect( float* input, uint32_t width, uint32_t height, Detection** detections, uint32_t overlay=OVERLAY_BOX );
+	int Detect( float* input, uint32_t width, uint32_t height, Detection** detections, uint32_t overlay=OVERLAY_DEFAULT );
 
 	/**
 	 * Detect object locations in an RGBA image, into an array of the results allocated by the user.
@@ -379,7 +380,7 @@ class detectNet : public tensorNet
 	 * @param[in]  overlay bitwise OR combination of overlay flags (@see OverlayFlags and @see Overlay()), or OVERLAY_NONE.
 	 * @returns    The number of detected objects, 0 if there were no detected objects, and -1 if an error was encountered.
 	 */
-	int Detect( float* input, uint32_t width, uint32_t height, Detection* detections, uint32_t overlay=OVERLAY_BOX );
+	int Detect( float* input, uint32_t width, uint32_t height, Detection* detections, uint32_t overlay=OVERLAY_DEFAULT );
 
 	/**
 	 * Draw the detected bounding boxes overlayed on an RGBA image.
@@ -450,6 +451,11 @@ class detectNet : public tensorNet
  	 * Set overlay alpha blending value for all classes (between 0-255).
 	 */
 	void SetOverlayAlpha( float alpha );
+
+	/**
+	 * Set the line width used during overlay when OVERLAY_LINES is used.
+	 */
+	inline void SetLineWidth( float width )						{ mLineWidth = width; }
 
 	/**
 	 * Load class descriptions from a label file.
@@ -487,7 +493,8 @@ class detectNet : public tensorNet
 	float  mCoverageThreshold;
 	float* mClassColors[2];
 	float  mMeanPixel;
-
+	float  mLineWidth;
+
 	std::vector<std::string> mClassDesc;
 	std::vector<std::string> mClassSynset;
 

diff --git a/python/bindings/PyDetectNet.cpp b/python/bindings/PyDetectNet.cpp
@@ -584,7 +584,7 @@ static void PyDetectNet_Dealloc( PyDetectNet_Object* self )
 				 "  image   (capsule) -- CUDA memory capsule\n" \
 				 "  width   (int)  -- width of the image (in pixels)\n" \
 				 "  height  (int)  -- height of the image (in pixels)\n" \
-				 "  overlay (str)  -- combination of box,labels,none flags (default is 'box')\n\n" \
+				 "  overlay (str)  -- combination of box,lines,labels,conf,none flags (default is 'box,labels,conf')\n\n" \
 				 "Returns:\n" \
 				 "  [Detections] -- list containing the detected objects (see detectNet.Detection)"
 
@@ -659,7 +659,7 @@ static PyObject* PyDetectNet_Detect( PyDetectNet_Object* self, PyObject* args, P
 				 "  [Detections]   -- list containing the detected objects (see detectNet.Detection)" \
 				 "  width   (int)  -- width of the image (in pixels)\n" \
 				 "  height  (int)  -- height of the image (in pixels)\n" \
-				 "  overlay (str)  -- combination of box,labels,none flags (default is 'box')\n\n" \
+				 "  overlay (str)  -- combination of box,lines,labels,conf,none flags (default is 'box,labels,conf')\n\n" \
 				 "Returns:\n" \
 				 "  None"
 
@@ -907,6 +907,41 @@ PyObject* PyDetectNet_SetOverlayAlpha( PyDetectNet_Object* self, PyObject* args,
 }
 
 
+#define DOC_SET_LINE_WIDTH "Set the line width used during overlay when 'lines' mode is used\n\n" \
+				 	  "Parameters:\n" \
+					  "  width (float) -- desired line width, in pixels\n" \
+					  "Returns:  (none)"
+
+// SetOverlayAlpha
+PyObject* PyDetectNet_SetLineWidth( PyDetectNet_Object* self, PyObject* args, PyObject *kwds )
+{
+	if( !self || !self->net )
+	{
+		PyErr_SetString(PyExc_Exception, LOG_PY_INFERENCE "detectNet invalid object instance");
+		return NULL;
+	}
+
+	float width = 0.0f;
+	static char* kwlist[] = {"width", NULL};
+
+	if( !PyArg_ParseTupleAndKeywords(args, kwds, "f", kwlist, &width) )
+	{
+		PyErr_SetString(PyExc_Exception, LOG_PY_INFERENCE "detectNet.SetLineWidth() failed to parse arguments");
+		return NULL;
+	}
+
+	if( width <= 0.0f )
+	{
+		PyErr_SetString(PyExc_Exception, LOG_PY_INFERENCE "detectNet.SetLineWidth() -- provided value is out-of-range");
+		return NULL;
+	}
+
+	self->net->SetLineWidth(width);
+
+	Py_RETURN_NONE;
+}
+
+
 #define DOC_USAGE_STRING     "Return the command line parameters accepted by __init__()\n\n" \
 					    "Parameters:  (none)\n\n" \
 					    "Returns:\n" \
@@ -934,6 +969,7 @@ static PyMethodDef pyDetectNet_Methods[] =
 	{ "GetClassDesc", (PyCFunction)PyDetectNet_GetClassDesc, METH_VARARGS, DOC_GET_CLASS_DESC},
 	{ "GetClassSynset", (PyCFunction)PyDetectNet_GetClassSynset, METH_VARARGS, DOC_GET_CLASS_SYNSET},
 	{ "SetOverlayAlpha", (PyCFunction)PyDetectNet_SetOverlayAlpha, METH_VARARGS|METH_KEYWORDS, DOC_SET_OVERLAY_ALPHA},
+	{ "SetLineWidth", (PyCFunction)PyDetectNet_SetLineWidth, METH_VARARGS|METH_KEYWORDS, DOC_SET_LINE_WIDTH},
 	{ "Usage", (PyCFunction)PyDetectNet_Usage, METH_NOARGS|METH_STATIC, DOC_USAGE_STRING},	
 	{NULL}  /* Sentinel */
 };

diff --git a/utils b/utils