refine pose

triplemu/cpp-refine-240820
triplemu 3 months ago
parent 53da84d953
commit 8042b51a65
  1. 74
      csrc/pose/normal/CMakeLists.txt
  2. 138
      csrc/pose/normal/cmake/FindTensorRT.cmake
  3. 14
      csrc/pose/normal/cmake/Function.cmake
  4. 30
      csrc/pose/normal/include/common.hpp
  5. 6075
      csrc/pose/normal/include/filesystem.hpp
  6. 21
      csrc/pose/normal/main.cpp

@ -10,50 +10,78 @@ set(CMAKE_CXX_STANDARD 14)
set(CMAKE_BUILD_TYPE Release)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake")
include(Function)
# CUDA
find_package(CUDA REQUIRED)
message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n")
print_var(CUDA_LIBRARIES)
print_var(CUDA_INCLUDE_DIRS)
get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY)
message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n")
print_var(CUDA_LIB_DIR)
# OpenCV
find_package(OpenCV REQUIRED)
message(STATUS "OpenCV Libs: \n${OpenCV_LIBS}\n")
message(STATUS "OpenCV Libraries: \n${OpenCV_LIBRARIES}\n")
message(STATUS "OpenCV Headers: \n${OpenCV_INCLUDE_DIRS}\n")
print_var(OpenCV_LIBS)
print_var(OpenCV_LIBRARIES)
print_var(OpenCV_INCLUDE_DIRS)
# TensorRT
set(TensorRT_INCLUDE_DIRS /usr/include/x86_64-linux-gnu)
set(TensorRT_LIBRARIES /usr/lib/x86_64-linux-gnu)
message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n")
message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n")
find_package(TensorRT REQUIRED)
print_var(TensorRT_LIBRARIES)
print_var(TensorRT_INCLUDE_DIRS)
print_var(TensorRT_LIB_DIR)
if (TensorRT_VERSION_MAJOR GREATER_EQUAL 10)
message(STATUS "Build with -DTRT_10")
add_definitions(-DTRT_10)
endif ()
list(APPEND INCLUDE_DIRS
list(APPEND ALL_INCLUDE_DIRS
${CUDA_INCLUDE_DIRS}
${OpenCV_INCLUDE_DIRS}
${TensorRT_INCLUDE_DIRS}
./include
)
${CMAKE_CURRENT_SOURCE_DIR}/include
)
list(APPEND ALL_LIBS
${CUDA_LIBRARIES}
${CUDA_LIB_DIR}
${OpenCV_LIBRARIES}
${TensorRT_LIBRARIES}
)
)
list(APPEND ALL_LIB_DIRS
${CUDA_LIB_DIR}
${TensorRT_LIB_DIR}
)
print_var(ALL_INCLUDE_DIRS)
print_var(ALL_LIBS)
print_var(ALL_LIB_DIRS)
include_directories(${INCLUDE_DIRS})
add_executable(${PROJECT_NAME}
main.cpp
include/yolov8-pose.hpp
include/common.hpp
)
${CMAKE_CURRENT_SOURCE_DIR}/main.cpp
${CMAKE_CURRENT_SOURCE_DIR}/include/yolov8-pose.hpp
${CMAKE_CURRENT_SOURCE_DIR}/include/common.hpp
)
target_include_directories(
${PROJECT_NAME}
PUBLIC
${ALL_INCLUDE_DIRS}
)
target_link_directories(
${PROJECT_NAME}
PUBLIC
${ALL_LIB_DIRS}
)
target_link_directories(${PROJECT_NAME} PUBLIC ${ALL_LIBS})
target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin cudart ${OpenCV_LIBS})
target_link_libraries(
${PROJECT_NAME}
PRIVATE
${ALL_LIBS}
)
if (${OpenCV_VERSION} VERSION_GREATER_EQUAL 4.7.0)
message(STATUS "Build with -DBATCHED_NMS")

@ -0,0 +1,138 @@
# This module defines the following variables:
#
# ::
#
# TensorRT_INCLUDE_DIRS
# TensorRT_LIBRARIES
# TensorRT_FOUND
#
# ::
#
# TensorRT_VERSION_STRING - version (x.y.z)
# TensorRT_VERSION_MAJOR - major version (x)
# TensorRT_VERSION_MINOR - minor version (y)
# TensorRT_VERSION_PATCH - patch version (z)
#
# Hints
# ^^^^^
# A user may set ``TensorRT_ROOT`` to an installation root to tell this module where to look.
#
set(_TensorRT_SEARCHES)
if(TensorRT_ROOT)
set(_TensorRT_SEARCH_ROOT PATHS ${TensorRT_ROOT} NO_DEFAULT_PATH)
list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_ROOT)
endif()
# appends some common paths
set(_TensorRT_SEARCH_NORMAL
PATHS "/usr"
)
list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_NORMAL)
# Include dir
foreach(search ${_TensorRT_SEARCHES})
find_path(TensorRT_INCLUDE_DIR NAMES NvInfer.h ${${search}} PATH_SUFFIXES include)
endforeach()
if(NOT TensorRT_LIBRARY)
foreach(search ${_TensorRT_SEARCHES})
find_library(TensorRT_LIBRARY NAMES nvinfer ${${search}} PATH_SUFFIXES lib)
if(NOT TensorRT_LIB_DIR)
get_filename_component(TensorRT_LIB_DIR ${TensorRT_LIBRARY} DIRECTORY)
endif ()
endforeach()
endif()
if(NOT TensorRT_nvinfer_plugin_LIBRARY)
foreach(search ${_TensorRT_SEARCHES})
find_library(TensorRT_nvinfer_plugin_LIBRARY NAMES nvinfer_plugin ${${search}} PATH_SUFFIXES lib)
endforeach()
endif()
mark_as_advanced(TensorRT_INCLUDE_DIR)
if(TensorRT_INCLUDE_DIR AND EXISTS "${TensorRT_INCLUDE_DIR}/NvInfer.h")
if(EXISTS "${TensorRT_INCLUDE_DIR}/NvInferVersion.h")
set(_VersionSearchFile "${TensorRT_INCLUDE_DIR}/NvInferVersion.h")
else ()
set(_VersionSearchFile "${TensorRT_INCLUDE_DIR}/NvInfer.h")
endif ()
file(STRINGS "${_VersionSearchFile}" TensorRT_MAJOR REGEX "^#define NV_TENSORRT_MAJOR [0-9]+.*$")
file(STRINGS "${_VersionSearchFile}" TensorRT_MINOR REGEX "^#define NV_TENSORRT_MINOR [0-9]+.*$")
file(STRINGS "${_VersionSearchFile}" TensorRT_PATCH REGEX "^#define NV_TENSORRT_PATCH [0-9]+.*$")
string(REGEX REPLACE "^#define NV_TENSORRT_MAJOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MAJOR "${TensorRT_MAJOR}")
string(REGEX REPLACE "^#define NV_TENSORRT_MINOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MINOR "${TensorRT_MINOR}")
string(REGEX REPLACE "^#define NV_TENSORRT_PATCH ([0-9]+).*$" "\\1" TensorRT_VERSION_PATCH "${TensorRT_PATCH}")
set(TensorRT_VERSION_STRING "${TensorRT_VERSION_MAJOR}.${TensorRT_VERSION_MINOR}.${TensorRT_VERSION_PATCH}")
endif()
include(FindPackageHandleStandardArgs)
FIND_PACKAGE_HANDLE_STANDARD_ARGS(TensorRT REQUIRED_VARS TensorRT_LIBRARY TensorRT_INCLUDE_DIR VERSION_VAR TensorRT_VERSION_STRING)
if(TensorRT_FOUND)
set(TensorRT_INCLUDE_DIRS ${TensorRT_INCLUDE_DIR})
if(NOT TensorRT_LIBRARIES)
set(TensorRT_LIBRARIES ${TensorRT_LIBRARY})
if (TensorRT_nvinfer_plugin_LIBRARY)
list(APPEND TensorRT_LIBRARIES ${TensorRT_nvinfer_plugin_LIBRARY})
endif()
endif()
if(NOT TARGET TensorRT::TensorRT)
add_library(TensorRT INTERFACE IMPORTED)
add_library(TensorRT::TensorRT ALIAS TensorRT)
endif()
if(NOT TARGET TensorRT::nvinfer)
add_library(TensorRT::nvinfer SHARED IMPORTED)
if (WIN32)
foreach(search ${_TensorRT_SEARCHES})
find_file(TensorRT_LIBRARY_DLL
NAMES nvinfer.dll
PATHS ${${search}}
PATH_SUFFIXES bin
)
endforeach()
set_target_properties(TensorRT::nvinfer PROPERTIES
INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}"
IMPORTED_LOCATION "${TensorRT_LIBRARY_DLL}"
IMPORTED_IMPLIB "${TensorRT_LIBRARY}"
)
else()
set_target_properties(TensorRT::nvinfer PROPERTIES
INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}"
IMPORTED_LOCATION "${TensorRT_LIBRARY}"
)
endif()
target_link_libraries(TensorRT INTERFACE TensorRT::nvinfer)
endif()
if(NOT TARGET TensorRT::nvinfer_plugin AND TensorRT_nvinfer_plugin_LIBRARY)
add_library(TensorRT::nvinfer_plugin SHARED IMPORTED)
if (WIN32)
foreach(search ${_TensorRT_SEARCHES})
find_file(TensorRT_nvinfer_plugin_LIBRARY_DLL
NAMES nvinfer_plugin.dll
PATHS ${${search}}
PATH_SUFFIXES bin
)
endforeach()
set_target_properties(TensorRT::nvinfer_plugin PROPERTIES
INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}"
IMPORTED_LOCATION "${TensorRT_nvinfer_plugin_LIBRARY_DLL}"
IMPORTED_IMPLIB "${TensorRT_nvinfer_plugin_LIBRARY}"
)
else()
set_target_properties(TensorRT::nvinfer_plugin PROPERTIES
INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}"
IMPORTED_LOCATION "${TensorRT_nvinfer_plugin_LIBRARY}"
)
endif()
target_link_libraries(TensorRT INTERFACE TensorRT::nvinfer_plugin)
endif()
endif()

@ -0,0 +1,14 @@
function(print_var var)
set(value "${${var}}")
string(LENGTH "${value}" value_length)
if(value_length GREATER 0)
math(EXPR last_index "${value_length} - 1")
string(SUBSTRING "${value}" ${last_index} ${last_index} last_char)
endif()
if(NOT "${last_char}" STREQUAL "\n")
set(value "${value}\n")
endif()
message(STATUS "${var}:\n ${value}")
endfunction()

@ -5,9 +5,8 @@
#ifndef POSE_NORMAL_COMMON_HPP
#define POSE_NORMAL_COMMON_HPP
#include "NvInfer.h"
#include "filesystem.hpp"
#include "opencv2/opencv.hpp"
#include <sys/stat.h>
#include <unistd.h>
#define CHECK(call) \
do { \
@ -89,33 +88,6 @@ inline static float clamp(float val, float min, float max)
return val > min ? (val < max ? val : max) : min;
}
inline bool IsPathExist(const std::string& path)
{
if (access(path.c_str(), 0) == F_OK) {
return true;
}
return false;
}
inline bool IsFile(const std::string& path)
{
if (!IsPathExist(path)) {
printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
return false;
}
struct stat buffer;
return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
}
inline bool IsFolder(const std::string& path)
{
if (!IsPathExist(path)) {
return false;
}
struct stat buffer;
return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
}
namespace pose {
struct Binding {
size_t size = 1;

File diff suppressed because it is too large Load Diff

@ -5,6 +5,8 @@
#include "yolov8-pose.hpp"
#include <chrono>
namespace fs = ghc::filesystem;
const std::vector<std::vector<unsigned int>> KPS_COLORS = {{0, 255, 0},
{0, 255, 0},
{0, 255, 0},
@ -65,11 +67,16 @@ const std::vector<std::vector<unsigned int>> LIMB_COLORS = {{51, 153, 255},
int main(int argc, char** argv)
{
if (argc != 3) {
fprintf(stderr, "Usage: %s [engine_path] [image_path/image_dir/video_path]\n", argv[0]);
return -1;
}
// cuda:0
cudaSetDevice(0);
const std::string engine_file_path{argv[1]};
const std::string path{argv[2]};
const fs::path path{argv[2]};
std::vector<std::string> imagePathList;
bool isVideo{false};
@ -79,8 +86,8 @@ int main(int argc, char** argv)
auto yolov8_pose = new YOLOv8_pose(engine_file_path);
yolov8_pose->make_pipe(true);
if (IsFile(path)) {
std::string suffix = path.substr(path.find_last_of('.') + 1);
if (fs::exists(path)) {
std::string suffix = path.extension();
if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") {
imagePathList.push_back(path);
}
@ -93,8 +100,8 @@ int main(int argc, char** argv)
std::abort();
}
}
else if (IsFolder(path)) {
cv::glob(path + "/*.jpg", imagePathList);
else if (fs::is_directory(path)) {
cv::glob(path.string() + "/*.jpg", imagePathList);
}
cv::Mat res, image;
@ -131,9 +138,9 @@ int main(int argc, char** argv)
}
}
else {
for (auto& path : imagePathList) {
for (auto& p : imagePathList) {
objs.clear();
image = cv::imread(path);
image = cv::imread(p);
yolov8_pose->copy_from_Mat(image, size);
auto start = std::chrono::system_clock::now();
yolov8_pose->infer();

Loading…
Cancel
Save