-
Notifications
You must be signed in to change notification settings - Fork 0
/
CMakeLists.txt
72 lines (65 loc) · 3.26 KB
/
CMakeLists.txt
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
cmake_minimum_required(VERSION 3.1)
project(tensorrt-engine)
#================================================================================
# Set compiler
#================================================================================
set(CMAKE_CXX_FLAGS "-std=c++17 -g -O3 -fPIC")
#================================================================================
# Build options
#================================================================================
option(BUILD_RESNET50_TEST "Build test binary for ResNet50 model" OFF)
option(BUILD_ULTRAFACE_TEST "Build test binary for Ultraface model" OFF)
#================================================================================
# Find packages
#================================================================================
find_package(CUDA)
find_package(OpenCV)
#================================================================================
# Build tensorrt-engine lib
#================================================================================
set(LIB "tensorrt-engine")
# Cuda link libs dir
link_directories(${CUDA_TOOLKIT_ROOT_DIR}/targets/x86_64-linux/lib)
# Store all cpp files into SRC variable
file(GLOB SRC ${CMAKE_CURRENT_SOURCE_DIR}/src/*.cpp)
# Lib
add_library(${LIB} SHARED ${SRC})
# Include header files
target_include_directories(${LIB} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include
PUBLIC ${CUDA_INCLUDE_DIRS}
PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/ext/pkg-json)
# Link libs
target_link_libraries(${LIB} PRIVATE cudart
PRIVATE nvinfer
PRIVATE nvonnxparser
PRIVATE ${OpenCV_LIBS})
#================================================================================
# Build engine tests
#================================================================================
if(BUILD_RESNET50_TEST)
set(BIN "ResNet50_test")
# Create executable
add_executable(${BIN} ${CMAKE_CURRENT_SOURCE_DIR}/test/ResNet50_test.cpp)
# Link libs
target_link_libraries(${BIN} PRIVATE ${LIB})
# Include header files
target_include_directories(${BIN} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)
# Install test dataset
install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/test/cat.rgb DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
endif()
if(BUILD_ULTRAFACE_TEST)
set(BIN "Ultraface_test")
# Create executable
add_executable(${BIN} ${CMAKE_CURRENT_SOURCE_DIR}/test/Ultraface_test.cpp)
# Link libs
target_link_libraries(${BIN} PRIVATE ${LIB}
PRIVATE ${OpenCV_LIBS})
# Include header files
target_include_directories(${BIN} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)
endif()
#================================================================================
# Install models and necesery extra files
#================================================================================
install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/models/Ultraface/ultraface-RFB-320.onnx DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/models/ResNet50/resnet50-v1-7.onnx DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/models/ResNet50/resnet_classes.json DESTINATION ${CMAKE_CURRENT_BINARY_DIR})