cmake_minimum_required(VERSION 3.10)
project(trt_pipeline CUDA CXX)

# 设置 C++ 标准
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)

# 设置 CUDA 架构
set(CMAKE_CUDA_ARCHITECTURES 75)  # 根据实际GPU调整

# 启用测试
enable_testing()

# 查找依赖包
find_package(OpenCV REQUIRED)
find_package(yaml-cpp REQUIRED)
find_package(GTest REQUIRED)
find_package(CUDA REQUIRED)

# 设置 CUDA 包含目录和库目录
include_directories(${CUDA_INCLUDE_DIRS})
link_directories(${CUDA_TOOLKIT_ROOT_DIR}/lib64)

# 设置 TensorRT 路径
set(TENSORRT_ROOT "/usr" CACHE PATH "TensorRT root directory")
set(TENSORRT_INCLUDE_DIRS 
    "${TENSORRT_ROOT}/include"
    "${TENSORRT_ROOT}/include/x86_64-linux-gnu"
)
set(TENSORRT_LIB_DIRS 
    "${TENSORRT_ROOT}/lib"
    "${TENSORRT_ROOT}/lib/x86_64-linux-gnu"
)

# 查找 TensorRT 库
find_library(NVINFER_LIB nvinfer HINTS ${TENSORRT_LIB_DIRS})
find_library(NVONNXPARSER_LIB nvonnxparser HINTS ${TENSORRT_LIB_DIRS})
find_library(NVINFER_PLUGIN_LIB nvinfer_plugin HINTS ${TENSORRT_LIB_DIRS})

if(NOT NVINFER_LIB OR NOT NVONNXPARSER_LIB OR NOT NVINFER_PLUGIN_LIB)
    message(FATAL_ERROR "TensorRT libraries not found")
endif()

# 添加子目录
add_subdirectory(pipeline)
add_subdirectory(tests)

# 添加主程序
add_executable(trt_pipeline_demo main.cpp)

# 链接库
target_link_libraries(trt_pipeline_demo
    PRIVATE
        pipeline
        ${OpenCV_LIBS}
        yaml-cpp
        ${CUDA_LIBRARIES}
        ${NVINFER_LIB}
        ${NVONNXPARSER_LIB}
        ${NVINFER_PLUGIN_LIB}
        cudart
        avcodec
        avformat
        avutil
        swscale
        pthread
)