├── .gitignore
├── .idea
├── Solov2-TensorRT-CPP.iml
├── misc.xml
├── modules.xml
├── vcs.xml
└── workspace.xml
├── CMakeLists.txt
├── InstanceSegment
├── TensorRtSample
│ ├── ErrorRecorder.h
│ ├── common.h
│ ├── half.h
│ ├── logger.cpp
│ ├── logger.h
│ └── logging.h
├── buffer.cpp
├── buffer.h
├── infer.cpp
├── infer.h
├── parameters.cpp
├── parameters.h
├── pipeline.cpp
├── pipeline.h
├── solo.cpp
├── solo.h
├── utils.cpp
└── utils.h
├── LICENSE
├── README.md
├── build_model.cpp
├── common.py
├── config
├── config.yaml
├── kitti.png
└── solov2_cpp.png
├── demo.cpp
├── main.cpp
└── onnx_exporter.py
/.gitignore:
--------------------------------------------------------------------------------
1 | /cmake-build-debug/
--------------------------------------------------------------------------------
/.idea/Solov2-TensorRT-CPP.iml:
--------------------------------------------------------------------------------
1 |
2 |
--------------------------------------------------------------------------------
/.idea/misc.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
--------------------------------------------------------------------------------
/.idea/modules.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
--------------------------------------------------------------------------------
/.idea/vcs.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
--------------------------------------------------------------------------------
/.idea/workspace.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
25 |
26 |
27 |
28 |
29 |
30 |
31 |
32 |
33 |
34 |
35 |
36 |
37 |
38 |
39 |
40 |
41 |
42 |
43 |
44 |
45 |
46 |
47 |
48 |
49 |
50 |
51 |
52 |
53 |
54 |
55 |
56 |
57 |
58 |
59 |
60 |
61 |
62 |
63 |
64 |
65 |
66 |
67 |
68 |
69 |
70 |
71 |
72 |
73 |
74 |
75 |
76 |
77 |
78 |
79 |
80 |
81 |
82 |
83 |
84 |
85 |
86 |
87 |
88 |
89 |
90 |
91 |
95 |
96 |
97 |
98 |
99 |
100 | 1641028294137
101 |
102 |
103 | 1641028294137
104 |
105 |
106 |
107 |
108 |
109 |
110 |
111 |
112 |
113 |
114 |
115 |
116 |
117 |
118 |
--------------------------------------------------------------------------------
/CMakeLists.txt:
--------------------------------------------------------------------------------
1 | cmake_minimum_required(VERSION 3.20)
2 | project(Solov2-TensorRT-CPP)
3 |
4 | set(CMAKE_CXX_STANDARD 17)
5 |
6 |
7 | find_package(OpenCV 3.4 REQUIRED)
8 | include_directories(${OpenCV_INCLUDE_DIRS})
9 |
10 | #Libtorch
11 | set(CMAKE_PREFIX_PATH "/usr/local/lib;/home/chen/app/libtorch")
12 |
13 | find_package(Torch REQUIRED)
14 | include_directories(${TORCH_INCLUDE_DIRS})
15 |
16 | #CUDA
17 | find_package(CUDA 10.2 REQUIRED)
18 | include_directories(${CUDA_INCLUDE_DIRS})
19 |
20 | set(TensorRT_LIBS nvinfer nvonnxparser nvinfer_plugin)
21 | include_directories(${TensorRT_INCLUDE_DIRS})
22 |
23 |
24 | aux_source_directory(InstanceSegment Segment_SOURCES)
25 | aux_source_directory(InstanceSegment/TensorRtSample TensorRtSample_SOURCES)
26 |
27 | add_executable(segment main.cpp ${Segment_SOURCES} ${TensorRtSample_SOURCES})
28 | target_link_libraries(segment ${CUDA_LIBRARIES} ${OpenCV_LIBRARIES} ${TORCH_LIBRARIES} ${TensorRT_LIBS})
29 |
30 | add_executable(demo demo.cpp ${Segment_SOURCES} ${TensorRtSample_SOURCES})
31 | target_link_libraries(demo ${CUDA_LIBRARIES} ${OpenCV_LIBRARIES} ${TORCH_LIBRARIES} ${TensorRT_LIBS})
32 |
33 | add_executable(build_model build_model.cpp InstanceSegment/parameters.cpp InstanceSegment/TensorRtSample/logger.cpp)
34 | target_link_libraries(build_model ${TensorRT_LIBS} pthread ${OpenCV_LIBRARIES})
35 |
36 |
--------------------------------------------------------------------------------
/InstanceSegment/TensorRtSample/ErrorRecorder.h:
--------------------------------------------------------------------------------
1 | /*
2 | * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
3 | *
4 | * Licensed under the Apache License, Version 2.0 (the "License");
5 | * you may not use this file except in compliance with the License.
6 | * You may obtain a copy of the License at
7 | *
8 | * http://www.apache.org/licenses/LICENSE-2.0
9 | *
10 | * Unless required by applicable law or agreed to in writing, software
11 | * distributed under the License is distributed on an "AS IS" BASIS,
12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 | * See the License for the specific language governing permissions and
14 | * limitations under the License.
15 | */
16 |
17 | #ifndef ERROR_RECORDER_H
18 | #define ERROR_RECORDER_H
19 | #include "NvInferRuntimeCommon.h"
20 | #include "logger.h"
21 | #include "logging.h"
22 | #include
23 | #include
24 | #include
25 | #include
26 | #include
27 | using namespace nvinfer1;
28 | //!
29 | //! A simple implementation of the IErrorRecorder interface for
30 | //! use by samples. This interface also can be used as a reference
31 | //! implementation.
32 | //! The sample Error recorder is based on a vector that pairs the error
33 | //! code and the error string into a single element. It also uses
34 | //! standard mutex's and atomics in order to make sure that the code
35 | //! works in a multi-threaded environment.
36 | //!
37 | class SampleErrorRecorder : public IErrorRecorder
38 | {
39 | using errorPair = std::pair;
40 | using errorStack = std::vector;
41 |
42 | public:
43 | SampleErrorRecorder() = default;
44 |
45 | virtual ~SampleErrorRecorder() noexcept {}
46 | int32_t getNbErrors() const noexcept final
47 | {
48 | return mErrorStack.size();
49 | }
50 | ErrorCode getErrorCode(int32_t errorIdx) const noexcept final
51 | {
52 | return invalidIndexCheck(errorIdx) ? ErrorCode::kINVALID_ARGUMENT : (*this)[errorIdx].first;
53 | };
54 | IErrorRecorder::ErrorDesc getErrorDesc(int32_t errorIdx) const noexcept final
55 | {
56 | return invalidIndexCheck(errorIdx) ? "errorIdx out of range." : (*this)[errorIdx].second.c_str();
57 | }
58 | // This class can never overflow since we have dynamic resize via std::vector usage.
59 | bool hasOverflowed() const noexcept final
60 | {
61 | return false;
62 | }
63 |
64 | // Empty the errorStack.
65 | void clear() noexcept final
66 | {
67 | try
68 | {
69 | // grab a lock so that there is no addition while clearing.
70 | std::lock_guard guard(mStackLock);
71 | mErrorStack.clear();
72 | }
73 | catch (const std::exception& e)
74 | {
75 | sample::gLogFatal << "Internal Error: " << e.what() << std::endl;
76 | }
77 | };
78 |
79 | //! Simple helper function that
80 | bool empty() const noexcept
81 | {
82 | return mErrorStack.empty();
83 | }
84 |
85 | bool reportError(ErrorCode val, IErrorRecorder::ErrorDesc desc) noexcept final
86 | {
87 | try
88 | {
89 | std::lock_guard guard(mStackLock);
90 | sample::gLogError << "Error[" << static_cast(val) << "]: " << desc << std::endl;
91 | mErrorStack.push_back(errorPair(val, desc));
92 | }
93 | catch (const std::exception& e)
94 | {
95 | sample::gLogFatal << "Internal Error: " << e.what() << std::endl;
96 | }
97 | // All errors are considered fatal.
98 | return true;
99 | }
100 |
101 | // Atomically increment or decrement the ref counter.
102 | IErrorRecorder::RefCount incRefCount() noexcept final
103 | {
104 | return ++mRefCount;
105 | }
106 | IErrorRecorder::RefCount decRefCount() noexcept final
107 | {
108 | return --mRefCount;
109 | }
110 |
111 | private:
112 | // Simple helper functions.
113 | const errorPair& operator[](size_t index) const noexcept
114 | {
115 | return mErrorStack[index];
116 | }
117 |
118 | bool invalidIndexCheck(int32_t index) const noexcept
119 | {
120 | // By converting signed to unsigned, we only need a single check since
121 | // negative numbers turn into large positive greater than the size.
122 | size_t sIndex = index;
123 | return sIndex >= mErrorStack.size();
124 | }
125 |
126 | // Mutex to hold when locking mErrorStack.
127 | std::mutex mStackLock;
128 |
129 | // Reference count of the class. Destruction of the class when mRefCount
130 | // is not zero causes undefined behavior.
131 | std::atomic mRefCount{0};
132 |
133 | // The error stack that holds the errors recorded by TensorRT.
134 | errorStack mErrorStack;
135 |
136 | }; // class SampleErrorRecorder
137 | #endif // ERROR_RECORDER_H
138 |
--------------------------------------------------------------------------------
/InstanceSegment/TensorRtSample/common.h:
--------------------------------------------------------------------------------
1 | /*
2 | * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
3 | *
4 | * Licensed under the Apache License, Version 2.0 (the "License");
5 | * you may not use this file except in compliance with the License.
6 | * You may obtain a copy of the License at
7 | *
8 | * http://www.apache.org/licenses/LICENSE-2.0
9 | *
10 | * Unless required by applicable law or agreed to in writing, software
11 | * distributed under the License is distributed on an "AS IS" BASIS,
12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 | * See the License for the specific language governing permissions and
14 | * limitations under the License.
15 | */
16 |
17 | #ifndef TENSORRT_COMMON_H
18 | #define TENSORRT_COMMON_H
19 |
20 | // For loadLibrary
21 | #ifdef _MSC_VER
22 | // Needed so that the max/min definitions in windows.h do not conflict with std::max/min.
23 | #define NOMINMAX
24 | #include
25 | #undef NOMINMAX
26 | #else
27 | #include
28 | #endif
29 |
30 | #include "NvInfer.h"
31 | #include "NvInferPlugin.h"
32 | #include "logger.h"
33 | #include "logging.h"
34 | #include
35 | #include
36 | #include
37 | #include
38 | #include
39 | #include
40 | #include
41 | #include
42 | #include
43 | #include
44 | #include