Merge branch 'alexeyk/trt4' of https://github.com/NVIDIA-Jetson/redtail…

… into alexeyk/trt4
Droneswat · Sep 14, 2018 · 46a193c · 46a193c
2 parents e4446d2 + 42d68da
commit 46a193c
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 24 deletions.
diff --git a/stereoDNN/sample_app/main.cpp b/stereoDNN/sample_app/main.cpp
@@ -13,6 +13,7 @@
 #include <opencv2/opencv.hpp>
 
 #include "redtail_tensorrt_plugins.h"
+#include "networks.h"
 
 #define UNUSED(x) ((void)(x))
 
@@ -25,30 +26,6 @@
 using namespace nvinfer1;
 using namespace redtail::tensorrt;
 
-namespace redtail { namespace tensorrt
-{
-using weight_map = std::unordered_map<std::string, Weights>;
-
-// NVSmall DNN: 1025x321 input, 96 max disparity.
-INetworkDefinition* createNVSmall1025x321Network(IBuilder& builder, IPluginContainer& plugin_factory,
-                                                 DimsCHW img_dims, const weight_map& weights, DataType data_type, ILogger& log);
-
-// Tiny version of NVSmall DNN: 513x161 input, 48 max disparity.
-INetworkDefinition* createNVTiny513x161Network(IBuilder& builder, IPluginContainer& plugin_factory,
-                                               DimsCHW img_dims, const weight_map& weights, DataType data_type,
-                                               ILogger& log);
-
-// Baseline ResNet-18 DNN: 1025x321 input, 136 max disparity.
-INetworkDefinition* createResNet18_1025x321Network(IBuilder& builder, IPluginContainer& plugin_factory,
-                                                   DimsCHW img_dims, const weight_map& weights, DataType data_type,
-                                                   ILogger& log);
-
-// ResNet18_2D DNN: 513x256 input, 96 max disparity.
-INetworkDefinition* createResNet18_2D_513x257Network(IBuilder& builder, IPluginContainer& plugin_factory,
-                                                     DimsCHW img_dims, const weight_map& weights, DataType data_type, ILogger& log);
-
-} }
-
 class Logger : public nvinfer1::ILogger
 {
 public:

diff --git a/stereoDNN/sample_app/networks.h b/stereoDNN/sample_app/networks.h
@@ -0,0 +1,40 @@
+
+// Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+// Full license terms provided in LICENSE.md file.
+
+#ifndef NETWORKS_H
+#define NETWORKS_H
+
+#include <NvInfer.h>
+#include <string>
+#include <unordered_map>
+
+namespace redtail { namespace tensorrt
+{
+
+using namespace nvinfer1;
+
+using weight_map = std::unordered_map<std::string, Weights>;
+
+class IPluginContainer;
+
+// NVSmall DNN: 1025x321 input, 96 max disparity.
+INetworkDefinition* createNVSmall1025x321Network(IBuilder& builder, IPluginContainer& plugin_factory,
+                                                 DimsCHW img_dims, const weight_map& weights, DataType data_type, ILogger& log);
+
+// Tiny version of NVSmall DNN: 513x161 input, 48 max disparity.
+INetworkDefinition* createNVTiny513x161Network(IBuilder& builder, IPluginContainer& plugin_factory,
+                                               DimsCHW img_dims, const weight_map& weights, DataType data_type,
+                                               ILogger& log);
+
+// Baseline ResNet-18 DNN: 1025x321 input, 136 max disparity.
+INetworkDefinition* createResNet18_1025x321Network(IBuilder& builder, IPluginContainer& plugin_factory,
+                                                   DimsCHW img_dims, const weight_map& weights, DataType data_type,
+                                                   ILogger& log);
+
+// ResNet18_2D DNN: 513x256 input, 96 max disparity.
+INetworkDefinition* createResNet18_2D_513x257Network(IBuilder& builder, IPluginContainer& plugin_factory,
+                                                     DimsCHW img_dims, const weight_map& weights, DataType data_type, ILogger& log);
+}}
+
+#endif