fixed SetOverlayAlpha() issue

Media-4-Machine-Laboratory · Sep 14, 2020 · d54eb0c · d54eb0c
1 parent 610e2c5
commit d54eb0c
Show file tree

Hide file tree

Showing 7 changed files with 19 additions and 5 deletions.
diff --git a/c/segNet.cpp b/c/segNet.cpp
@@ -41,6 +41,7 @@ segNet::segNet() : tensorNet()
 	mLastInputHeight = 0;
 	mLastInputFormat = IMAGE_UNKNOWN;
 
+	mColorsAlphaSet = NULL;
 	mClassColors[0] = NULL;	// cpu ptr
 	mClassColors[1] = NULL;  // gpu ptr
 
@@ -382,6 +383,16 @@ segNet* segNet::Create( const char* prototxt, const char* model, const char* lab
 		net->mClassColors[0][n*4+3] = 255.0f;	// a
 	}
 
+	net->mColorsAlphaSet = (bool*)malloc(numClasses * sizeof(bool));
+
+	if( !net->mColorsAlphaSet )
+	{
+		printf(LOG_TRT "segNet -- failed to allocate class colors alpha flag array\n");
+		return NULL;
+	}
+
+	memset(net->mColorsAlphaSet, 0, numClasses * sizeof(bool));
+
 	// initialize array of classified argmax
 	const int s_w = DIMS_W(net->mOutputs[0].dims);
 	const int s_h = DIMS_H(net->mOutputs[0].dims);
@@ -620,6 +631,8 @@ void segNet::SetClassColor( uint32_t classIndex, float r, float g, float b, floa
 	mClassColors[0][i+1] = g;
 	mClassColors[0][i+2] = b;
 	mClassColors[0][i+3] = a;
+
+	mColorsAlphaSet[classIndex] = (a == 255) ? false : true;
 }
 
 
@@ -630,7 +643,7 @@ void segNet::SetOverlayAlpha( float alpha, bool explicit_exempt )
 
 	for( uint32_t n=0; n < numClasses; n++ )
 	{
-		if( !explicit_exempt || mClassColors[0][n*4+3] == 255 )
+		if( !explicit_exempt || !mColorsAlphaSet[n] /*mClassColors[0][n*4+3] == 255*/ )
 			mClassColors[0][n*4+3] = alpha;
 	}
 }

diff --git a/c/segNet.h b/c/segNet.h
@@ -43,7 +43,7 @@
  * Default alpha blending value used during overlay
  * @ingroup segNet
  */
-#define SEGNET_DEFAULT_ALPHA 120
+#define SEGNET_DEFAULT_ALPHA 150
 
 /**
  * Standard command-line options able to be passed to segNet::Create()
@@ -69,7 +69,7 @@
 		  "  --input-blob=INPUT   name of the input layer (default: '" SEGNET_DEFAULT_INPUT "')\n" 		\
 		  "  --output-blob=OUTPUT name of the output layer (default: '" SEGNET_DEFAULT_OUTPUT "')\n" 		\
 		  "  --batch-size=BATCH   maximum batch size (default is 1)\n"								\
-            "  --alpha=ALPHA        overlay alpha blending value, range 0-255 (default: 120)\n"			\
+            "  --alpha=ALPHA        overlay alpha blending value, range 0-255 (default: 150)\n"			\
 		  "  --visualize=VISUAL   visualization flags (e.g. --visualize=overlay,mask)\n"				\
 		  "                       valid combinations are:  'overlay', 'mask'\n"						\
 		  "  --profile            enable layer profiling in TensorRT\n\n"
@@ -363,6 +363,7 @@ class segNet : public tensorNet
 	std::vector<std::string> mClassLabels;
 	std::string mClassPath;
 
+	bool*    mColorsAlphaSet;	/**< true if class color had been explicitly set from file or user */
 	float*   mClassColors[2];	/**< array of overlay colors in shared CPU/GPU memory */
 	uint8_t* mClassMap[2];		/**< runtime buffer for the argmax-classified class index of each tile */
 

diff --git a/data/images/dog_3.jpg b/data/images/dog_3.jpg
diff --git a/data/images/dog_4.jpg b/data/images/dog_4.jpg
diff --git a/data/images/dog_5.jpg b/data/images/dog_5.jpg
diff --git a/examples/segnet/segnet.cpp b/examples/segnet/segnet.cpp
@@ -197,7 +197,7 @@ int main( int argc, char** argv )
 	}
 
 	// set alpha blending value for classes that don't explicitly already have an alpha	
-	net->SetOverlayAlpha(cmdLine.GetFloat("alpha", 180.0f));
+	net->SetOverlayAlpha(cmdLine.GetFloat("alpha", 150.0f));
 
 	// get the desired overlay/mask filtering mode
 	const segNet::FilterMode filterMode = segNet::FilterModeFromStr(cmdLine.GetString("filter-mode", "linear"));

diff --git a/python/examples/segnet.py b/python/examples/segnet.py
@@ -40,7 +40,7 @@
 parser.add_argument("--filter-mode", type=str, default="linear", choices=["point", "linear"], help="filtering mode used during visualization, options are:\n  'point' or 'linear' (default: 'linear')")
 parser.add_argument("--visualize", type=str, default="overlay,mask", help="Visualization options (can be 'overlay' 'mask' 'overlay,mask'")
 parser.add_argument("--ignore-class", type=str, default="void", help="optional name of class to ignore in the visualization results (default: 'void')")
-parser.add_argument("--alpha", type=float, default=175.0, help="alpha blending value to use during overlay, between 0.0 and 255.0 (default: 175.0)")
+parser.add_argument("--alpha", type=float, default=150.0, help="alpha blending value to use during overlay, between 0.0 and 255.0 (default: 150.0)")
 parser.add_argument("--stats", action="store_true", help="compute statistics about segmentation mask class output")
 
 is_headless = ["--headless"] if sys.argv[0].find('console.py') != -1 else [""]