automatic model downloads for depthNet

dusty-nv · Dec 18, 2022 · 6b0d9c2 · 6b0d9c2
1 parent 1f56290
commit 6b0d9c2
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 104 deletions.
diff --git a/c/depthNet.cpp b/c/depthNet.cpp
@@ -22,20 +22,20 @@
 
 #include "depthNet.h"
 #include "tensorConvert.h"
+#include "modelDownloader.h"
 
 #include "commandLine.h"
+#include "filesystem.h"
 #include "cudaMappedMemory.h"
 
 #include "mat33.h"
 
-
 #define DEPTH_HISTOGRAM_CUDA
 
 
 // constructor
 depthNet::depthNet() : tensorNet()
 {
-	mNetworkType    = CUSTOM;
 	mDepthRange     = NULL;
 	mDepthEqualized = NULL;
 
@@ -112,64 +112,40 @@ uint32_t depthNet::VisualizationFlagsFromStr( const char* str_user, uint32_t def
 }
 
 
-// NetworkTypeFromStr
-depthNet::NetworkType depthNet::NetworkTypeFromStr( const char* modelName )
-{
-	if( !modelName )
-		return depthNet::CUSTOM;
-
-	depthNet::NetworkType type = depthNet::FCN_MOBILENET;
-
-	if( strcasecmp(modelName, "mobilenet") == 0 || strcasecmp(modelName, "fcn-mobilenet") == 0 || strcasecmp(modelName, "fcn_mobilenet") == 0 || strcasecmp(modelName, "monodepth-fcn-mobilenet") == 0 || strcasecmp(modelName, "monodepth_fcn_mobilenet") == 0 )
-		type = depthNet::FCN_MOBILENET;
-	else if( strcasecmp(modelName, "resnet18") == 0 || strcasecmp(modelName, "fcn-resnet18") == 0 || strcasecmp(modelName, "fcn_resnet18") == 0 || strcasecmp(modelName, "monodepth-fcn-resnet18") == 0 || strcasecmp(modelName, "monodepth_fcn_resnet18") == 0 )
-		type = depthNet::FCN_RESNET18;
-	else if( strcasecmp(modelName, "resnet50") == 0 || strcasecmp(modelName, "fcn-resnet50") == 0 || strcasecmp(modelName, "fcn_resnet50") == 0 || strcasecmp(modelName, "monodepth-fcn-resnet50") == 0 || strcasecmp(modelName, "monodepth_fcn_resnet50") == 0 )
-		type = depthNet::FCN_RESNET50;
-	else
-		type = depthNet::CUSTOM;
-
-	return type;
-}
-
-
-// NetworkTypeToStr
-const char* depthNet::NetworkTypeToStr( depthNet::NetworkType type )
-{
-	switch(type)
-	{
-		case FCN_MOBILENET:	return "MonoDepth-FCN-Mobilenet";
-		case FCN_RESNET18:	return "MonoDepth-FCN-ResNet18";
-		case FCN_RESNET50:	return "MonoDepth-FCN-ResNet50";
-		default:			return "Custom";
-	}
-}
-
-
 // Create
-depthNet* depthNet::Create( depthNet::NetworkType networkType, uint32_t maxBatchSize, 
+depthNet* depthNet::Create( const char* network, uint32_t maxBatchSize, 
 					   precisionType precision, deviceType device, bool allowGPUFallback )
 {
-	depthNet* net = NULL;
+	nlohmann::json model;
 
-	if( networkType == FCN_MOBILENET )
-		net = Create("networks/MonoDepth-FCN-Mobilenet/monodepth_fcn_mobilenet.onnx", DEPTHNET_DEFAULT_INPUT, DEPTHNET_DEFAULT_OUTPUT, maxBatchSize, precision, device, allowGPUFallback);
-	else if( networkType == FCN_RESNET18 )
-		net = Create("networks/MonoDepth-FCN-ResNet18/monodepth_fcn_resnet18.onnx", DEPTHNET_DEFAULT_INPUT, DEPTHNET_DEFAULT_OUTPUT, maxBatchSize, precision, device, allowGPUFallback);
-	else if( networkType == FCN_RESNET50 )
-		net = Create("networks/MonoDepth-FCN-ResNet50/monodepth_fcn_resnet50.onnx", DEPTHNET_DEFAULT_INPUT, DEPTHNET_DEFAULT_OUTPUT, maxBatchSize, precision, device, allowGPUFallback);
+	if( !DownloadModel(DEPTHNET_MODEL_TYPE, network, model) )
+		return NULL;
 
-	if( net != NULL )
-		net->mNetworkType = networkType;
+	std::string model_dir = "networks/" + model["dir"].get<std::string>() + "/";
+	std::string model_path = model_dir + JSON_STR(model["model"]);
+	std::string input = JSON_STR_DEFAULT(model["input"], DEPTHNET_DEFAULT_INPUT);
+	std::string output = JSON_STR_DEFAULT(model["output"], DEPTHNET_DEFAULT_OUTPUT);
 
-	return net;
+	return Create(model_path.c_str(), input.c_str(), output.c_str(), maxBatchSize, precision, device, allowGPUFallback);
 }
 
 
 // Create
 depthNet* depthNet::Create( const char* model_path, const char* input_blob, const char* output_blob,
 					   uint32_t maxBatchSize, precisionType precision, deviceType device, bool allowGPUFallback )
 {
+	// check for built-in model string
+	if( FindModel(DEPTHNET_MODEL_TYPE, model_path) )
+	{
+		return Create(model_path, maxBatchSize, precision, device, allowGPUFallback);
+	}
+	else if( fileExtension(model_path).length() == 0 )
+	{
+		LogError(LOG_TRT "couldn't find built-in mono-depth model '%s'\n", model_path);
+		return NULL;
+	}
+
+	// load custom model
 	depthNet* net = new depthNet();
 
 	if( !net )
@@ -272,9 +248,7 @@ depthNet* depthNet::Create( const commandLine& cmdLine )
 		modelName = cmdLine.GetString("model", "fcn-mobilenet");
 
 	// parse the network type
-	const depthNet::NetworkType type = NetworkTypeFromStr(modelName);
-
-	if( type == depthNet::CUSTOM )
+	if( !FindModel(DEPTHNET_MODEL_TYPE, modelName) )
 	{
 		const char* input  = cmdLine.GetString("input_blob");
 		const char* output = cmdLine.GetString("output_blob");
@@ -292,7 +266,7 @@ depthNet* depthNet::Create( const commandLine& cmdLine )
 	else
 	{
 		// create from pretrained model
-		net = depthNet::Create(type);
+		net = depthNet::Create(modelName, DEFAULT_MAX_BATCH_SIZE);
 	}
 
 	if( !net )

diff --git a/c/depthNet.h b/c/depthNet.h
@@ -39,6 +39,11 @@
  */
 #define DEPTHNET_DEFAULT_OUTPUT  "output_0"
 
+/**
+ * The model type for depthNet in data/networks/models.json
+ * @ingroup depthNet
+ */
+#define DEPTHNET_MODEL_TYPE "monodepth"
 
 /**
  * Command-line options able to be passed to depthNet::Create()
@@ -63,17 +68,6 @@
 class depthNet : public tensorNet
 {
 public:
-	/**
-	 * Network choice enumeration.
-	 */
-	enum NetworkType
-	{
-		CUSTOM,        /**< Custom model provided by the user */
-		FCN_MOBILENET,	/**< MobileNet backbone */
-		FCN_RESNET18,	/**< ResNet-18 backbone */
-		FCN_RESNET50,	/**< ResNet-50 backbone */
-	};
-
 	/**
 	 * Visualization flags.
 	 */
@@ -90,21 +84,10 @@ class depthNet : public tensorNet
 	static uint32_t VisualizationFlagsFromStr( const char* str, uint32_t default_value=VISUALIZE_INPUT|VISUALIZE_DEPTH );
 
 	/**
-	 * Parse a string to one of the built-in pretrained models.
-	 * Valid names are "mobilenet", "resnet-18", or "resnet-50", ect.
-	 * @returns one of the depthNet::NetworkType enums, or depthNet::CUSTOM on invalid string.
-	 */
-	static NetworkType NetworkTypeFromStr( const char* model_name );
-
-	/**
-	 * Convert a NetworkType enum to a string.
+	 * Load a pre-trained model.
+	 * @see DEPTHNET_USAGE_STRING for the available models.
 	 */
-	static const char* NetworkTypeToStr( NetworkType network );
-
-	/**
-	 * Load a new network instance
-	 */
-	static depthNet* Create( NetworkType networkType=FCN_MOBILENET, 
+	static depthNet* Create( const char* network="fcn-mobilenet", 
 						uint32_t maxBatchSize=DEFAULT_MAX_BATCH_SIZE, 
 						precisionType precision=TYPE_FASTEST,
 				   		deviceType device=DEVICE_GPU, bool allowGPUFallback=true );
@@ -241,16 +224,6 @@ class depthNet : public tensorNet
 	 */
 	inline uint32_t GetDepthFieldHeight() const					{ return DIMS_H(mOutputs[0].dims); }
 
-	/**
-	 * Retrieve the network type (alexnet or googlenet)
-	 */
-	inline NetworkType GetNetworkType() const					{ return mNetworkType; }
-
-	/**
- 	 * Retrieve a string describing the network name.
-	 */
-	inline const char* GetNetworkName() const					{ return NetworkTypeToStr(mNetworkType); }
-
 	/**
 	 * Extract and save the point cloud to a PCD file (depth only).
 	 * @note SavePointCloud() should only be called after Process()
@@ -291,8 +264,6 @@ class depthNet : public tensorNet
 	bool histogramEqualization();
 	bool histogramEqualizationCUDA();
 
-	NetworkType mNetworkType;
-
 	int2*     mDepthRange;
 	float*    mDepthEqualized;
 	uint32_t* mHistogram;

diff --git a/data/networks/models.json b/data/networks/models.json
@@ -1,6 +1,5 @@
 {
     "classification": {
-
         "alexnet": {
             "url": "https://nvidia.box.com/shared/static/2y1x8v1q7g6w5d1355dqaxd3snx6bpze.gz",
             "tar": "Alexnet.tar.gz",
@@ -110,7 +109,6 @@
     },
 
     "detection": {
-
         "ssd-mobilenet-v1": {
             "alias": ["ssd_mobilenet_v1", "ssd_mobilenet_v1_coco", "ssd-mobilenet-v1-coco"],
             "url": "https://nvidia.box.com/shared/static/0pg3xi9opwio65df14rdgrtw40ivbk1o.gz",
@@ -230,7 +228,6 @@
     },
 
     "segmentation": {
-
         "fcn-resnet18-cityscapes-512x256": {
             "alias": ["fcn_resnet18_cityscapes_512x256", "fcn_resnet18_cityscapes", "fcn-resnet18-cityscapes"],
             "url": "https://nvidia.box.com/shared/static/k7s7gdgi098309fndm2xbssj553vf71s.gz",
@@ -354,7 +351,6 @@
     },
 
     "pose": {
-
         "resnet18-body": {
             "alias": ["resnet18_body"],
             "url": "https://nvidia.box.com/shared/static/waf8bsu58v9qh9qj3sp3wsw1nyj61xm5.gz",
@@ -386,6 +382,35 @@
             "topology": "human_pose.json",
             "colors": "colors.txt",
             "description": "DenseNet121-based human body model with PAF attention"
+        }
+    },
+
+    "monodepth": {
+        "fcn-mobilenet": {
+            "alias": ["fcn_mobilenet", "monodepth-mobilenet", "monodepth_mobilenet", "mobilenet"],
+            "url": "https://nvidia.box.com/shared/static/frgbiqeieaja0o8b0eyb87fjbsqd4zup.gz",
+            "tar": "MonoDepth-FCN-Mobilenet.tar.gz",
+            "dir": "MonoDepth-FCN-Mobilenet",
+            "model": "monodepth_fcn_mobilenet.onnx",
+            "description": "FCN-Mobilenet monocular depth estimation model"
+        },
+
+        "fcn-resnet18": {
+            "alias": ["fcn_resnet18", "monodepth-resnet18", "monodepth_resnet18", "resnet18"],
+            "url": "https://nvidia.box.com/shared/static/ai2sxrp1tg8mk4j0jbrw3vthqjp8x0af.gz",
+            "tar": "MonoDepth-FCN-ResNet18.tar.gz",
+            "dir": "MonoDepth-FCN-ResNet18",
+            "model": "monodepth_fcn_resnet18.onnx",
+            "description": "FCN-ResNet18 monocular depth estimation model"
         },
+
+        "fcn-resnet50": {
+            "alias": ["fcn_resnet50", "monodepth-resnet50", "monodepth_resnet50", "resnet50"],
+            "url": "https://nvidia.box.com/shared/static/3umpq9yrv3nj3ltiwlooijx5of414gbh.gz",
+            "tar": "MonoDepth-FCN-ResNet50.tar.gz",
+            "dir": "MonoDepth-FCN-ResNet50",
+            "model": "monodepth_fcn_resnet50.onnx",
+            "description": "FCN-ResNet50 monocular depth estimation model"
+        }
     }
 }
diff --git a/python/bindings/PyDepthNet.cpp b/python/bindings/PyDepthNet.cpp
@@ -65,7 +65,7 @@ static int PyDepthNet_Init( PyDepthNet_Object* self, PyObject *args, PyObject *k
 	// determine whether to use argv or built-in network
 	if( argList != NULL && PyList_Check(argList) && PyList_Size(argList) > 0 )
 	{
-		LogVerbose(LOG_PY_INFERENCE "depthNet loading network using argv command line params\n");
+		LogDebug(LOG_PY_INFERENCE "depthNet loading network using argv command line params\n");
 
 		// parse the python list into char**
 		const size_t argc = PyList_Size(argList);
@@ -107,20 +107,11 @@ static int PyDepthNet_Init( PyDepthNet_Object* self, PyObject *args, PyObject *k
 	}
 	else
 	{
-		LogVerbose(LOG_PY_INFERENCE "depthNet loading build-in network '%s'\n", network);
-
-		// parse the selected built-in network
-		depthNet::NetworkType networkType = depthNet::NetworkTypeFromStr(network);
-
-		if( networkType == depthNet::CUSTOM )
-		{
-			PyErr_SetString(PyExc_Exception, LOG_PY_INFERENCE "depthNet invalid built-in network was requested");
-			return -1;
-		}
+		LogDebug(LOG_PY_INFERENCE "depthNet loading build-in network '%s'\n", network);
 
 		// load the built-in network
 		Py_BEGIN_ALLOW_THREADS
-		self->net = depthNet::Create(networkType);
+		self->net = depthNet::Create(network, DEFAULT_MAX_BATCH_SIZE);
 		Py_END_ALLOW_THREADS
 	}