update cls code

csrc/cls/normal/CMakeLists.txt

@@ -1,4 +1,4 @@
-cmake_minimum_required(VERSION 3.1)
+cmake_minimum_required(VERSION 3.12)
 
 set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86 89 90)
 set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc)
@@ -10,47 +10,76 @@ 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)
+
+
+add_executable(
+        ${PROJECT_NAME}
+        ${CMAKE_CURRENT_SOURCE_DIR}/main.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/include/yolov8-cls.hpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/include/common.hpp
+)
 
-include_directories(${INCLUDE_DIRS})
+target_include_directories(
+        ${PROJECT_NAME}
+        PUBLIC
+        ${ALL_INCLUDE_DIRS}
+)
 
-add_executable(${PROJECT_NAME}
-        main.cpp
-        include/yolov8-cls.hpp
-        include/common.hpp
-        )
+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}
+)

csrc/cls/normal/include/common.hpp

@@ -5,9 +5,8 @@
 #ifndef CLS_NORMAL_COMMON_HPP
 #define CLS_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 cls {
 struct Binding {
     size_t         size  = 1;

csrc/cls/normal/include/yolov8-cls.hpp

@@ -6,7 +6,7 @@
 
 #include "NvInferPlugin.h"
 #include "common.hpp"
-#include "fstream"
+#include <fstream>
 
 using namespace cls;
 

csrc/cls/normal/main.cpp

@@ -1,9 +1,11 @@
 //
 // Created by ubuntu on 4/27/24.
 //
-#include "chrono"
 #include "opencv2/opencv.hpp"
 #include "yolov8-cls.hpp"
+#include <chrono>
+
+namespace fs = ghc::filesystem;
 
 const std::vector<std::string> CLASS_NAMES = {"tench",
                                               "goldfish",
@@ -1008,27 +1010,30 @@ const std::vector<std::string> CLASS_NAMES = {"tench",
 
 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};
 
-    assert(argc == 3);
-
     auto yolov8_cls = new YOLOv8_cls(engine_file_path);
     yolov8_cls->make_pipe(true);
 
-    if (IsFile(path)) {
-        std::string suffix = path.substr(path.find_last_of('.') + 1);
-        if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") {
+    if (fs::exists(path)) {
+        std::string suffix = path.extension();
+        if (suffix == ".jpg" || suffix == ".jpeg" || suffix == ".png") {
             imagePathList.push_back(path);
         }
-        else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov"
-                 || suffix == "mkv") {
+        else if (suffix == ".mp4" || suffix == ".avi" || suffix == ".m4v" || suffix == ".mpeg" || suffix == ".mov"
+                 || suffix == ".mkv") {
             isVideo = true;
         }
         else {
@@ -1036,8 +1041,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;
@@ -1071,9 +1076,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_cls->copy_from_Mat(image, size);
             auto start = std::chrono::system_clock::now();
             yolov8_cls->infer();

csrc/jetson/detect/include/yolov8.hpp

@@ -5,7 +5,7 @@
 #define JETSON_DETECT_YOLOV8_HPP
 #include "NvInferPlugin.h"
 #include "common.hpp"
-#include "fstream"
+#include <fstream>
 using namespace det;
 
 class YOLOv8 {

csrc/jetson/detect/main.cpp

@@ -1,9 +1,9 @@
 //
 // Created by ubuntu on 3/16/23.
 //
-#include "chrono"
 #include "opencv2/opencv.hpp"
 #include "yolov8.hpp"
+#include <chrono>
 
 const std::vector<std::string> CLASS_NAMES = {
     "person",         "bicycle",    "car",           "motorcycle",    "airplane",     "bus",           "train",

csrc/jetson/pose/include/yolov8-pose.hpp

@@ -6,7 +6,7 @@
 
 #include "NvInferPlugin.h"
 #include "common.hpp"
-#include "fstream"
+#include <fstream>
 
 using namespace pose;
 

csrc/jetson/pose/main.cpp

@@ -1,9 +1,9 @@
 //
 // Created by ubuntu on 4/7/23.
 //
-#include "chrono"
 #include "opencv2/opencv.hpp"
 #include "yolov8-pose.hpp"
+#include <chrono>
 
 const std::vector<std::vector<unsigned int>> KPS_COLORS = {{0, 255, 0},
                                                            {0, 255, 0},

csrc/jetson/segment/main.cpp

@@ -1,9 +1,9 @@
 //
 // Created by ubuntu on 3/16/23.
 //
-#include "chrono"
 #include "opencv2/opencv.hpp"
 #include "yolov8-seg.hpp"
+#include <chrono>
 
 const std::vector<std::string> CLASS_NAMES = {
     "person",         "bicycle",    "car",           "motorcycle",    "airplane",     "bus",           "train",

csrc/pose/normal/include/yolov8-pose.hpp

@@ -6,7 +6,7 @@
 
 #include "NvInferPlugin.h"
 #include "common.hpp"
-#include "fstream"
+#include <fstream>
 
 using namespace pose;
 

csrc/pose/normal/main.cpp

@@ -1,9 +1,9 @@
 //
 // Created by ubuntu on 4/7/23.
 //
-#include "chrono"
 #include "opencv2/opencv.hpp"
 #include "yolov8-pose.hpp"
+#include <chrono>
 
 const std::vector<std::vector<unsigned int>> KPS_COLORS = {{0, 255, 0},
                                                            {0, 255, 0},

csrc/segment/normal/main.cpp

@@ -1,9 +1,9 @@
 //
 // Created by ubuntu on 2/8/23.
 //
-#include "chrono"
 #include "opencv2/opencv.hpp"
 #include "yolov8-seg.hpp"
+#include <chrono>
 
 const std::vector<std::string> CLASS_NAMES = {
     "person",         "bicycle",    "car",           "motorcycle",    "airplane",     "bus",           "train",

csrc/segment/simple/main.cpp

@@ -1,9 +1,9 @@
 //
 // Created by ubuntu on 1/20/23.
 //
-#include "chrono"
 #include "opencv2/opencv.hpp"
 #include "yolov8-seg.hpp"
+#include <chrono>
 
 const std::vector<std::string> CLASS_NAMES = {
     "person",         "bicycle",    "car",           "motorcycle",    "airplane",     "bus",           "train",