diff --git a/CMakeLists.txt b/CMakeLists.txt index f31cbee..2e429c3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,13 +1,16 @@ cmake_minimum_required(VERSION 3.10) PROJECT(Linfer VERSION 1.0.0 LANGUAGES C CXX CUDA) -set(CMAKE_CXX_STANDARD 11) +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_BUILD_TYPE Debug) -#set(CMAKE_BUILD_TYPE Release) set(EXECUTABLE_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/workspace) set(LIBRARY_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/workspace) set(CMAKE_SKIP_BUILD_RPATH False) set(CMAKE_SKIP_RPATH False) -set(CMAKE_BUILD_RPATH "/home/lsf/Third_party/TensorRT-8.6.1.6/lib") +set(CMAKE_BUILD_RPATH "/usr/local/TensorRT-8.5.1.7/lib") # TODO: Modify this +find_package(PkgConfig) +pkg_check_modules(YAMLCPP REQUIRED yaml-cpp>=0.5) +include_directories(${YAMLCPP_INCLUDE_DIRS}) file(GLOB_RECURSE CPPS ${PROJECT_SOURCE_DIR}/apps/*.cpp @@ -15,9 +18,8 @@ file(GLOB_RECURSE CPPS ${PROJECT_SOURCE_DIR}/trt_common/*.cpp ${PROJECT_SOURCE_DIR}/trt_common/*.cu ) - -set(CUDA_DIR "/usr/local/cuda") -set(TENSORRT_DIR "/home/lsf/Third_party/TensorRT-8.6.1.6") +set(CUDA_DIR "/usr/local/cuda-11.8") # TODO: Modify this +set(TENSORRT_DIR "/usr/local/TensorRT-8.5.1.7") # TODO: Modify this find_package(OpenCV REQUIRED) if(POLICY CMP0146) cmake_policy(SET CMP0146 OLD) @@ -28,7 +30,7 @@ include_directories( ${OpenCV_INCLUDE_DIRS} ${CUDA_DIR}/include ${TENSORRT_DIR}/include - "/usr/include/eigen3" + "/usr/include/eigen3" # TODO: Modify this (maybe) ) link_directories( ${CUDA_DIR}/lib64 @@ -41,15 +43,16 @@ list(APPEND ALL_LIBS ) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-deprecated-declarations -Wfatal-errors -pthread -w") -set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -std=c++11 -g -O0 -Xcompiler -fPIC") +set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -std=c++17 -g -O0 -Xcompiler -fPIC") add_library(${PROJECT_NAME} SHARED ${CPPS}) target_link_libraries(${PROJECT_NAME} ${ALL_LIBS}) # reference：https://developer.nvidia.com/cuda-gpus#compute -set_property(TARGET ${PROJECT_NAME} PROPERTY CUDA_ARCHITECTURES 89) +# set_property(TARGET ${PROJECT_NAME} PROPERTY CUDA_ARCHITECTURES 89) +set_property(TARGET ${PROJECT_NAME} PROPERTY CUDA_ARCHITECTURES 86) # TODO: Modify this if necessary target_compile_options(${PROJECT_NAME} PUBLIC $<$:--default-stream per-thread -lineinfo --use_fast_math --disable-warnings>) add_executable(pro main.cpp) -target_link_libraries(pro ${PROJECT_NAME} ${ALL_LIBS}) +target_link_libraries(pro ${PROJECT_NAME} ${ALL_LIBS} ${YAMLCPP_LIBRARIES}) \ No newline at end of file diff --git a/README.md b/README.md index d1f7941..ac12164 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,8 @@ ![Language](https://img.shields.io/badge/language-c++-brightgreen) ![Language](https://img.shields.io/badge/CUDA-12.1-brightgreen) ![Language](https://img.shields.io/badge/TensorRT-8.6.1.6-brightgreen) ![Language](https://img.shields.io/badge/OpenCV-4.5.5-brightgreen) ![Language](https://img.shields.io/badge/ubuntu-20.04-brightorigin) +[English](README_EN.md) | 简体中文 + ## Introduction 基于 TensorRT 的 C++ 高性能推理库。 diff --git a/README_EN.md b/README_EN.md new file mode 100644 index 0000000..c9984b0 --- /dev/null +++ b/README_EN.md @@ -0,0 +1,154 @@ +# Linfer + +![Language](https://img.shields.io/badge/language-c++-brightgreen) ![Language](https://img.shields.io/badge/CUDA-12.1-brightgreen) ![Language](https://img.shields.io/badge/TensorRT-8.6.1.6-brightgreen) ![Language](https://img.shields.io/badge/OpenCV-4.5.5-brightgreen) ![Language](https://img.shields.io/badge/ubuntu-20.04-brightorigin) + +English | [简体中文](README.md) + +## Introduction +A high-performance inference library for C++ based on TensorRT. + + + +## Update News +🚀(2024.12.24) Supports YAML Configuration + +🚀(2024.06.06) Supports target detection algorithm Yolov10! + +🚀(2024.05.23) Supports semantic segmentation algorithm: PP-LiteSeg and MobileSeg in PaddleSeg, which are lightweight and efficient and suitable for deployment! + +🚀(2023.12.03) Supports Panoramic driving perception algorithm YOLOPv2, Better, Faster, Stronger! + +🚀 (2023.11.06) Support panoramic driving perception algorithm YOLOP! + +🚀 (2023.10.19) Support single target tracking OSTrack, LightTrack! The separate single target tracking repository is [github]( https://github.com/l-sf/Track-trt) + +🚀(2023.10.09) Support target detection algorithm RT-DETR! + +🚀(2023.08.26) Support PTQ quantization, Yolov5/7 QAT quantization! + +🚀(2023.07.19) Support target detection Yolo series 5/X/7/8, multi-target tracking Bytetrack. + +## Highlights + +- Support panoramic driving perception YOLOPv2, Target detection RT-DETR, Yolov5/X/7/8/10, multi-target tracking Bytetrack, single target tracking OSTrack, LightTrack; +- Pre-processing and post-processing implement CUDA kernel functions, and high-performance reasoning can also be achieved on the Jetson edge; +- Encapsulate Tensor and Infer to achieve memory reuse, automatic CPU/GPU memory copying, engine context management, input and output binding, etc.; +- The inference process implements the producer-consumer model, realizes the parallelization of preprocessing and inference, and further improves performance; - Use RAII concept + interface mode to encapsulate applications, which is safe and convenient to use. + +## Easy Using + +The code structure of this project is as follows: The implementation code of each algorithm is stored in the `apps` folder, where `app_xxx.cpp` is the call demo function corresponding to the `xxx` algorithm. Each algorithm has no dependency on each other. If you only need to use yolopv2, you can delete all other algorithms in this folder without any impact; the `trt_common` folder includes the commonly used cuda_tools, which encapsulates TensorRT's Tensor and Infer, and the producer-consumer model; The `quant-tools` folder contains quantitative scripts, mainly yolov5/7. + +Which algorithm to use is called in `main.cpp` demo function. + +```bash +. +├── apps +│   ├── yolo +│   └── yolop +│   ├── app_yolo.cpp +│   ├── app_yolop.cpp +│   ├── ... +├── trt_common +│   ├── cuda_tools. hpp +│   ├── trt_infer.hpp +│   ├── trt_tensor.hpp +│   └── ... +├── quant-tools +│   └── ... +├── workspace +│ └── ... +├── CMakeLists .txt +└── main.cpp +``` + +If you want to deploy your own algorithm, just create a new folder for your algorithm in the `apps` folder, and imitate the `trt_infer/trt_tensor` in other algorithms. You can use it as you like. I will update more detailed instructions later when I have more free time. + + +## Project Build and Run + +1. install cuda/tensorrt/opencv + + [reference](https://github.com/l-sf/Notes/blob/main/notes/Ubuntu20.04_install_tutorials.md#%E4%BA%94cuda--cudnn--tensorrt-install) + +2. compile engine + +3. Download the onnx model from [google drive](https://drive.google.com/drive/folders/16ZqDaxlWm1aDXQsjsxLS7yFL0YqzHbxT?usp=sharing) or export it according to the tutorial, the tutorial is in README of each folder. Put your onnx file under `workspace/onnx_models` folder (create it) + + ```bash + cd Linfer/workspace + # Modify the onnx path bash compile_engine.sh + + # Uncomment particular model form compile_engine.sh or copy any of the commands from it like this + + # YOLOV8S + trtexec --onnx=./onnx_models/yolov8n.onnx \ + --saveEngine=./yolov8n.trt \ + --buildOnly \ + --minShapes=images:1x3x640x640 \ + --optShapes=images:1x3x640x640 \ + --maxShapes=images:8x3x640x640 \ + --fp16 +``` + +4. build + +```bash +# Modify CMakeLists.txt cuda/tensorrt/opencv is your own path cd Linfer +mkdir build && cd build +cmake .. && make -j4 +``` + +5. Configure: make your configuration file config.yaml please see a demo of file in config.yaml for instance if you wnt to run bytetrack tracking algorithm with YOLOV8, you'll need to build .trt from step 3 and then provide it's path in config.yaml like below + +```yaml +tasks: + - task: "track" + subtasks: + - type: "inference_bytetrack" + engine_file: "/home/e300/mahmood/code/Linfer/workspace/yolov8s.trt" + gpuid: 0 + yolo_type: "V8" + video_file: "/home/e300/mahmood/code/Linfer/workspace/videos/snow.mp4" + output_save_path: "" +``` + +6. run (to avoid any errors please provide full paths always) + +```bash +cd Linfer/workspace +./pro config.yaml +``` + +## Speed Test + +Tested on Jetson Orin Nano 8G, the test includes the entire process (image preprocessing + model inference + post-processing decoding) + +| Model | Precision | Resolution | FPS(bs=1) | +| :--------: | :-------: | :--------: | :-------: | +| yolov5_s | fp16 | 640x640 | 96.06 | +| yolox_s | fp16 | 640x640 | 79.64 | +| yolov7 | **int8** | 640x640 | 49.55 | +| yolov8_n | fp16 | 640x640 | 121.94 | +| yolov8_s | fp16 | 640x640 | 81.40 | +| yolov8_m | fp16 | 640x640 | 41.14 | +| yolov8_l | fp16 | 640x640 | 27.52 | +| yolov10_n | fp16 | 640x640 | 115.13 | +| yolov10_s | fp16 | 640x640 | 73.65 | +| yolov10_m | fp16 | 640x640 | 39.51 | +| yolov10_l | fp16 | 640x640 | 26.41 | +| rtdetr_r50 | fp16 | 640x640 | 11.25 | +| lighttrack | fp16 | 256x256 | 90.91 | +| ostrack | fp16 | 256x256 | 37.04 | +| yolop | fp16 | 640x640 | 31.4 | +| yolopv2 | fp16 | 480x640 | 21.9 | +| PP-LiteSeg | fp16 | 256x512 | 129.81 | +| MobileSeg | fp16 | 256x512 | 140.36 | + + + +## Reference + +- [tensorRT_Pro](https://github.com/shouxieai/tensorRT_Pro.git) +- [Video：详解TensorRT的C++/Python高性能部署，实战应用到项目](https://www.bilibili.com/video/BV1Xw411f7FW/?share_source=copy_web&vd_source=4bb05d1ac6ff39b7680900de14419dca) + diff --git a/apps/app_mot.cpp b/apps/app_mot.cpp index 3cbef2f..5de90fe 100644 --- a/apps/app_mot.cpp +++ b/apps/app_mot.cpp @@ -1,12 +1,13 @@ - - #include "trt_common/ilogger.hpp" #include "yolo/yolo.hpp" #include #include "bytetrack/BYTETracker.h" #include +#include using namespace std; +namespace fs = std::filesystem; + template static vector det2tracks(const Yolo::BoxArray& array, const Cond& cond){ @@ -30,7 +31,7 @@ static vector det2tracks(const Yolo::BoxArray& array, const Cond& cond){ } -void inference_bytetrack(const string& engine_file, int gpuid, Yolo::Type type, const string& video_file){ +void inference_bytetrack(const string& engine_file, int gpuid, Yolo::Type type, const string& video_file, const string& output_save_path){ auto engine = Yolo::create_infer( engine_file, // engine file @@ -59,8 +60,15 @@ void inference_bytetrack(const string& engine_file, int gpuid, Yolo::Type type, ).set_per_frame_motion({0.1, 0.1, 0.1, 0.1, 0.2, 0.2, 1, 0.2} ).set_max_time_lost(150); + + string output_path = output_save_path; + if (output_path.empty()) + { + fs::path input_path(video_file); + output_path = input_path.stem().string() + "_output" + input_path.extension().string(); + } - cv::VideoWriter writer("videos/res_mot.mp4", cv::VideoWriter::fourcc('M', 'P', 'E', 'G'), fps, cv::Size(width, height)); + cv::VideoWriter writer(output_path, cv::VideoWriter::fourcc('M', 'P', 'E', 'G'), fps, cv::Size(width, height)); auto cond = [](const Yolo::Box& b){return b.label == 0;}; shared_future> prev_fut; @@ -97,5 +105,4 @@ void inference_bytetrack(const string& engine_file, int gpuid, Yolo::Type type, writer.release(); printf("Done.\n"); -} - +} \ No newline at end of file diff --git a/apps/app_rtdetr.cpp b/apps/app_rtdetr.cpp index 0b32123..d7f76b2 100644 --- a/apps/app_rtdetr.cpp +++ b/apps/app_rtdetr.cpp @@ -1,120 +1,206 @@ - +#include #include #include #include "rtdetr/rtdetr.hpp" +#include + +namespace fs = std::filesystem; using namespace std; inline vector cocolabels = { - "person", "bicycle", "car", "motorcycle", "airplane", - "bus", "train", "truck", "boat", "traffic light", "fire hydrant", - "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", - "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", - "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", - "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", - "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", - "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", - "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", - "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", - "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", - "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", - "scissors", "teddy bear", "hair drier", "toothbrush" -}; - -inline std::tuple hsv2bgr(float h, float s, float v){ + "person", "bicycle", "car", "motorcycle", "airplane", + "bus", "train", "truck", "boat", "traffic light", "fire hydrant", + "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", + "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", + "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", + "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", + "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", + "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", + "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", + "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", + "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", + "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", + "scissors", "teddy bear", "hair drier", "toothbrush"}; + +inline std::tuple hsv2bgr(float h, float s, float v) +{ const int h_i = static_cast(h * 6); const float f = h * 6 - h_i; const float p = v * (1 - s); - const float q = v * (1 - f*s); + const float q = v * (1 - f * s); const float t = v * (1 - (1 - f) * s); float r, g, b; - switch (h_i) { - case 0:r = v; g = t; b = p;break; - case 1:r = q; g = v; b = p;break; - case 2:r = p; g = v; b = t;break; - case 3:r = p; g = q; b = v;break; - case 4:r = t; g = p; b = v;break; - case 5:r = v; g = p; b = q;break; - default:r = 1; g = 1; b = 1;break;} + switch (h_i) + { + case 0: + r = v; + g = t; + b = p; + break; + case 1: + r = q; + g = v; + b = p; + break; + case 2: + r = p; + g = v; + b = t; + break; + case 3: + r = p; + g = q; + b = v; + break; + case 4: + r = t; + g = p; + b = v; + break; + case 5: + r = v; + g = p; + b = q; + break; + default: + r = 1; + g = 1; + b = 1; + break; + } return make_tuple(static_cast(b * 255), static_cast(g * 255), static_cast(r * 255)); } -inline std::tuple random_color(int id){ - float h_plane = ((((unsigned int)id << 2) ^ 0x937151) % 100) / 100.0f;; +inline std::tuple random_color(int id) +{ + float h_plane = ((((unsigned int)id << 2) ^ 0x937151) % 100) / 100.0f; + ; float s_plane = ((((unsigned int)id << 3) ^ 0x315793) % 100) / 100.0f; return hsv2bgr(h_plane, s_plane, 1); } -inline string get_file_name(const string& path, bool include_suffix){ - if (path.empty()) return ""; +inline string get_file_name(const string &path, bool include_suffix) +{ + if (path.empty()) + return ""; int p = path.rfind('/'); int e = path.rfind('\\'); p = std::max(p, e); p += 1; - //include suffix + // include suffix if (include_suffix) return path.substr(p); int u = path.rfind('.'); if (u == -1) return path.substr(p); - if (u <= p) u = path.size(); + if (u <= p) + u = path.size(); return path.substr(p, u - p); } - -void performance(const string& engine_file, int gpuid){ +void performance(const string &engine_file, int gpuid, const string &input_dir) +{ auto infer = RTDETR::create_infer(engine_file, gpuid, 0.5); - if(infer == nullptr){ + if (infer == nullptr) + { printf("infer is nullptr.\n"); return; } int batch = 8; - std::vector images{cv::imread("imgs/bus.jpg"), cv::imread("imgs/girl.jpg"), - cv::imread("imgs/group.jpg"), cv::imread("imgs/yq.jpg")}; + vector files_; + files_.reserve(100); + cv::glob(input_dir + "/*.jpg", files_, true); + vector files(files_.begin(), files_.end()); + + if (files.empty()) + { + printf("No .jpg image files found in the input directory: %s\n", input_dir.c_str()); + return; + } + + std::vector images; + for (const auto &file : files) + { + auto image = cv::imread(file); + if (image.empty()) + { + printf("Error reading image file: %s\n", file.c_str()); + continue; + } + images.emplace_back(image); + } + + if (images.empty()) + { + printf("No valid images to process after reading files from: %s\n", input_dir.c_str()); + return; + } + for (int i = images.size(); i < batch; ++i) - images.push_back(images[i % 4]); + images.push_back(images[i % images.size()]); - // warmup + printf("Number of images to process: %zu\n", images.size()); + + // Warmup vector> boxes_array; - for(int i = 0; i < 10; ++i) + for (int i = 0; i < 10; ++i) boxes_array = infer->commits(images); boxes_array.back().get(); boxes_array.clear(); - // 测试 100 轮 + // Test 100 rounds const int ntest = 100; auto start = std::chrono::steady_clock::now(); - for(int i = 0; i < ntest; ++i) + for (int i = 0; i < ntest; ++i) boxes_array = infer->commits(images); - // 等待全部推理结束 + // Wait for all inference to finish boxes_array.back().get(); std::chrono::duration during = std::chrono::steady_clock::now() - start; double all_time = 1000.0 * during.count(); float avg_time = all_time / ntest / images.size(); - printf("Average time: %.2f ms, FPS: %.2f\n", engine_file.c_str(), avg_time, 1000 / avg_time); + printf("Average time for %s: %.2f ms, FPS: %.2f\n", engine_file.c_str(), avg_time, 1000 / avg_time); } - -void batch_inference(const string& engine_file, int gpuid){ +void batch_inference(const string &engine_file, int gpuid, const string &input_dir, const string &output_dir) +{ auto infer = RTDETR::create_infer(engine_file, gpuid, 0.5); - if(infer == nullptr){ + if (infer == nullptr) + { printf("infer is nullptr.\n"); return; } + fs::create_directories(output_dir); vector files_; files_.reserve(100); - cv::glob("imgs/*.jpg", files_, true); + cv::glob(input_dir + "/*.jpg", files_, true); vector files(files_.begin(), files_.end()); + if (files.empty()) + { + printf("No image files found in the input directory: %s\n", input_dir.c_str()); + return; + } vector images; - for(const auto& file : files){ + for (const auto &file : files) + { auto image = cv::imread(file); + if (image.empty()) + { + printf("Error reading image file: %s\n", file.c_str()); + continue; + } images.emplace_back(image); } + if (images.empty()) + { + printf("No valid images to process after reading files from: %s\n", input_dir.c_str()); + return; + } vector> boxes_array; boxes_array = infer->commits(images); @@ -122,11 +208,12 @@ void batch_inference(const string& engine_file, int gpuid){ // 等待全部推理结束 boxes_array.back().get(); - string root_res = "infer_res/rtdetr"; - for(int i = 0; i < boxes_array.size(); ++i){ + for (int i = 0; i < boxes_array.size(); ++i) + { cv::Mat image = images[i]; auto boxes = boxes_array[i].get(); - for(auto & ibox : boxes){ + for (auto &ibox : boxes) + { cv::Scalar color; std::tie(color[0], color[1], color[2]) = random_color(ibox.label); cv::rectangle(image, cv::Point(ibox.left, ibox.top), cv::Point(ibox.right, ibox.bottom), color, 2); @@ -134,28 +221,35 @@ void batch_inference(const string& engine_file, int gpuid){ auto name = cocolabels[ibox.label]; auto caption = cv::format("%s %.2f", name.c_str(), ibox.confidence); int text_width = cv::getTextSize(caption, 0, 1, 2, nullptr).width + 10; - cv::rectangle(image, cv::Point(ibox.left-2, ibox.top-32), cv::Point(ibox.left + text_width, ibox.top), color, -1); - cv::putText(image, caption, cv::Point(ibox.left, ibox.top-5), 0, 1, cv::Scalar::all(0), 2, 16); + cv::rectangle(image, cv::Point(ibox.left - 2, ibox.top - 32), cv::Point(ibox.left + text_width, ibox.top), color, -1); + cv::putText(image, caption, cv::Point(ibox.left, ibox.top - 5), 0, 1, cv::Scalar::all(0), 2, 16); } string file_name = get_file_name(files[i], false); - string save_path = cv::format("%s/%s.jpg", root_res.c_str(), file_name.c_str()); + string save_path = cv::format("%s/%s.jpg", output_dir.c_str(), file_name.c_str()); cv::imwrite(save_path, image); printf("Save to %s, %d object\n", save_path.c_str(), boxes.size()); } } - -void single_inference(const string& engine_file, int gpuid){ +void single_inference(const string &engine_file, int gpuid, const string &input_img, const string &output_img_path) +{ auto infer = RTDETR::create_infer(engine_file, gpuid, 0.6); - if(infer == nullptr){ + if (infer == nullptr) + { printf("infer is nullptr.\n"); return; } - auto image = cv::imread("imgs/bus.jpg"); + auto image = cv::imread(input_img); + if (image.empty()) + { + printf("Error reading image file: %s\n", input_img.c_str()); + return; + } auto boxes = infer->commit(image).get(); - for(auto& ibox : boxes){ + for (auto &ibox : boxes) + { cv::Scalar color; std::tie(color[0], color[1], color[2]) = random_color(ibox.label); cv::rectangle(image, cv::Point(ibox.left, ibox.top), cv::Point(ibox.right, ibox.bottom), color, 2); @@ -163,9 +257,12 @@ void single_inference(const string& engine_file, int gpuid){ auto name = cocolabels[ibox.label]; auto caption = cv::format("%s %.2f", name.c_str(), ibox.confidence); int text_width = cv::getTextSize(caption, 0, 1, 2, nullptr).width + 10; - cv::rectangle(image, cv::Point(ibox.left-2, ibox.top-32), cv::Point(ibox.left + text_width, ibox.top), color, -1); - cv::putText(image, caption, cv::Point(ibox.left, ibox.top-5), 0, 1, cv::Scalar::all(0), 2, 16); + cv::rectangle(image, cv::Point(ibox.left - 2, ibox.top - 32), cv::Point(ibox.left + text_width, ibox.top), color, -1); + cv::putText(image, caption, cv::Point(ibox.left, ibox.top - 5), 0, 1, cv::Scalar::all(0), 2, 16); } - cv::imwrite("infer_res/result.jpg", image); -} + fs::path path(output_img_path); + fs::create_directories(path.parent_path()); + cv::imwrite(output_img_path, image); + printf("Save to %s\n", output_img_path.c_str()); +} \ No newline at end of file diff --git a/apps/app_seg.cpp b/apps/app_seg.cpp index dba2e99..9817fee 100644 --- a/apps/app_seg.cpp +++ b/apps/app_seg.cpp @@ -1,61 +1,103 @@ - +#include #include #include #include "ppseg/ppseg.hpp" +#include + +namespace fs = std::filesystem; using namespace std; -void performance_seg(const string& engine_file, int gpuid){ +void performance_seg(const string &engine_file, int gpuid, const string &input_dir) +{ auto predictor = PPSeg::create_seg(engine_file, gpuid); - if(predictor == nullptr){ + if (predictor == nullptr) + { printf("predictor is nullptr.\n"); return; } - auto image = cv::imread("imgs/frame_0.jpg"); + vector files_; + files_.reserve(100); + cv::glob(input_dir + "/*.jpg", files_, true); + vector files(files_.begin(), files_.end()); + if (files.empty()) + { + printf("No image files found in the input directory: %s\n", input_dir.c_str()); + return; + } + std::vector images; + for (const auto &file : files) + { + auto image = cv::imread(file); + if (image.empty()) + { + printf("Error reading image file: %s\n", file.c_str()); + continue; + } + images.emplace_back(image); + } + + if (images.empty()) + { + printf("No valid images to process after reading files from: %s\n", input_dir.c_str()); + return; + } + int batch = 8; + for (int i = images.size(); i < batch; ++i) + images.push_back(images[i % images.size()]); - cv::Mat res; // warmup - for(int i = 0; i < 10; ++i) - res = predictor->seg(image); + cv::Mat res; + for (int i = 0; i < 10; ++i) + res = predictor->seg(images[i % images.size()]); // 测试 100 轮 const int ntest = 100; auto start = std::chrono::steady_clock::now(); - for(int i = 0; i < ntest; ++i) - res = predictor->seg(image); + for (int i = 0; i < ntest; ++i) + res = predictor->seg(images[i % images.size()]); std::chrono::duration during = std::chrono::steady_clock::now() - start; double all_time = 1000.0 * during.count(); - float avg_time = all_time / ntest; - printf("Average time: %.2f ms, FPS: %.2f\n", engine_file.c_str(), avg_time, 1000 / avg_time); + float avg_time = all_time / ntest / images.size(); + printf("Average time for %s: %.2f ms, FPS: %.2f\n", engine_file.c_str(), avg_time, 1000 / avg_time); } -void inference_seg(const string& engine_file, int gpuid){ +void inference_seg(const string &engine_file, int gpuid, const string &input_img, const string &output_img_path) +{ auto predictor = PPSeg::create_seg(engine_file, gpuid); - if(predictor == nullptr){ + if (predictor == nullptr) + { printf("predictor is nullptr.\n"); return; } - auto image = cv::imread("imgs/frame_3.jpg"); + auto image = cv::imread(input_img); + if (image.empty()) + { + printf("Error reading image file: %s\n", input_img.c_str()); + return; + } auto res = predictor->seg(image); cv::Mat color_img(image.size(), CV_8UC3, cv::Scalar(0, 0, 0)); // 遍历每个像素点，根据类别索引应用颜色映射 - for (int i = 0; i < res.rows; ++i) { - for (int j = 0; j < res.cols; ++j) { + for (int i = 0; i < res.rows; ++i) + { + for (int j = 0; j < res.cols; ++j) + { uchar pixel_value = res.at(i, j); color_img.at(i, j) = cv::Vec3b(PPSeg::color_map[pixel_value][2], - PPSeg::color_map[pixel_value][1], - PPSeg::color_map[pixel_value][0]); + PPSeg::color_map[pixel_value][1], + PPSeg::color_map[pixel_value][0]); } } cv::Mat out_color_img(image.size(), CV_8UC3, cv::Scalar(0, 0, 0)); float alpha = 0.7; out_color_img = (1 - alpha) * image + alpha * color_img; - cv::imwrite("infer_res/seg_frame_3.jpg", out_color_img); -} - - - + fs::path path(output_img_path); + fs::create_directories(path.parent_path()); + cv::imwrite(output_img_path, out_color_img); + printf("Save to %s\n", output_img_path.c_str()); +} \ No newline at end of file diff --git a/apps/app_yolo.cpp b/apps/app_yolo.cpp index c73d49a..f865003 100644 --- a/apps/app_yolo.cpp +++ b/apps/app_yolo.cpp @@ -1,7 +1,9 @@ - #include #include #include "yolo/yolo.hpp" +#include + +namespace fs = std::filesystem; using namespace std; @@ -62,8 +64,7 @@ inline string get_file_name(const string& path, bool include_suffix){ return path.substr(p, u - p); } - -void performance(const string& engine_file, int gpuid, Yolo::Type type){ +void performance(const string& engine_file, int gpuid, Yolo::Type type, const string& input_dir){ auto infer = Yolo::create_infer(engine_file, type, gpuid, 0.3, 0.45); if(infer == nullptr){ printf("infer is nullptr.\n"); @@ -71,59 +72,94 @@ void performance(const string& engine_file, int gpuid, Yolo::Type type){ } int batch = 8; - std::vector images{cv::imread("imgs/bus.jpg"), cv::imread("imgs/girl.jpg"), - cv::imread("imgs/group.jpg"), cv::imread("imgs/yq.jpg")}; + vector files_; + files_.reserve(100); + cv::glob(input_dir + "/*.jpg", files_, true); + vector files(files_.begin(), files_.end()); + + if (files.empty()) { + printf("No .jpg image files found in the input directory: %s\n", input_dir.c_str()); + return; + } + + std::vector images; + for(const auto& file : files){ + auto image = cv::imread(file); + if(image.empty()) { + printf("Error reading image file: %s\n", file.c_str()); + continue; + } + images.emplace_back(image); + } + + if (images.empty()) { + printf("No valid images to process after reading files from: %s\n", input_dir.c_str()); + return; + } + for (int i = images.size(); i < batch; ++i) - images.push_back(images[i % 4]); + images.push_back(images[i % images.size()]); + + printf("Number of images to process: %zu\n", images.size()); - // warmup + // Warmup vector> boxes_array; for(int i = 0; i < 10; ++i) boxes_array = infer->commits(images); boxes_array.back().get(); boxes_array.clear(); - // 测试 100 轮 + // Test 100 rounds const int ntest = 100; auto start = std::chrono::steady_clock::now(); - for(int i = 0; i < ntest; ++i) + for(int i = 0; i < ntest; ++i) boxes_array = infer->commits(images); - // 等待全部推理结束 + // Wait for all inference to finish boxes_array.back().get(); std::chrono::duration during = std::chrono::steady_clock::now() - start; double all_time = 1000.0 * during.count(); float avg_time = all_time / ntest / images.size(); - printf("Average time: %.2f ms, FPS: %.2f\n", engine_file.c_str(), avg_time, 1000 / avg_time); + printf("Average time for %s: %.2f ms, FPS: %.2f\n", engine_file.c_str(), avg_time, 1000 / avg_time); } -void batch_inference(const string& engine_file, int gpuid, Yolo::Type type){ +void batch_inference(const string& engine_file, int gpuid, Yolo::Type type, const string& input_dir, const string& output_dir){ auto infer = Yolo::create_infer(engine_file, type, gpuid, 0.25, 0.45); if(infer == nullptr){ printf("infer is nullptr.\n"); return; } + fs::create_directories(output_dir); vector files_; - files_.reserve(100); - cv::glob("imgs/*.jpg", files_, true); + files_.reserve(100); + cv::glob(input_dir + "/*.jpg", files_, true); vector files(files_.begin(), files_.end()); - + if (files.empty()) { + printf("No image files found in the input directory: %s\n", input_dir.c_str()); + return; + } vector images; for(const auto& file : files){ auto image = cv::imread(file); + if(image.empty()){ + printf("Error reading image file: %s\n", file.c_str()); + continue; + } images.emplace_back(image); } - + if (images.empty()) { + printf("No valid images to process after reading files from: %s\n", input_dir.c_str()); + return; + } vector> boxes_array; boxes_array = infer->commits(images); // 等待全部推理结束 boxes_array.back().get(); - string root_res = "infer_res"; - for(int i = 0; i < boxes_array.size(); ++i){ + for(int i = 0; i < boxes_array.size(); ++i){ cv::Mat image = images[i]; auto boxes = boxes_array[i].get(); for(auto & ibox : boxes){ @@ -138,20 +174,24 @@ void batch_inference(const string& engine_file, int gpuid, Yolo::Type type){ cv::putText(image, caption, cv::Point(ibox.left, ibox.top-5), 0, 1, cv::Scalar::all(0), 2, 16); } string file_name = get_file_name(files[i], false); - string save_path = cv::format("%s/%s.jpg", root_res.c_str(), file_name.c_str()); + string save_path = cv::format("%s/%s.jpg", output_dir.c_str(), file_name.c_str()); cv::imwrite(save_path, image); printf("Save to %s, %d object\n", save_path.c_str(), boxes.size()); } } -void single_inference(const string& engine_file, int gpuid, Yolo::Type type){ +void single_inference(const string& engine_file, int gpuid, Yolo::Type type, const string& input_img, const string& output_img_path){ auto infer = Yolo::create_infer(engine_file, type, gpuid, 0.25, 0.45); if(infer == nullptr){ printf("infer is nullptr.\n"); return; } - auto image = cv::imread("imgs/bus.jpg"); + auto image = cv::imread(input_img); + if(image.empty()){ + printf("Error reading image file: %s\n", input_img.c_str()); + return; + } auto boxes = infer->commit(image).get(); for(auto& ibox : boxes){ cv::Scalar color; @@ -164,6 +204,9 @@ void single_inference(const string& engine_file, int gpuid, Yolo::Type type){ cv::rectangle(image, cv::Point(ibox.left-2, ibox.top-32), cv::Point(ibox.left + text_width, ibox.top), color, -1); cv::putText(image, caption, cv::Point(ibox.left, ibox.top-5), 0, 1, cv::Scalar::all(0), 2, 16); } - cv::imwrite("infer_res/result.jpg", image); -} + fs::path path(output_img_path); + fs::create_directories(path.parent_path()); + cv::imwrite(output_img_path, image); + printf("Save to %s\n", output_img_path.c_str()); +} \ No newline at end of file diff --git a/apps/app_yolop.cpp b/apps/app_yolop.cpp index 1a6ad94..8bc9663 100644 --- a/apps/app_yolop.cpp +++ b/apps/app_yolop.cpp @@ -1,57 +1,122 @@ - #include #include #include "yolop/yolop.hpp" +#include + +namespace fs = std::filesystem; using namespace std; -void performance_yolop(const string& engine_file, YoloP::Type type, int gpuid){ +inline string get_file_name(const string &path, bool include_suffix) +{ + if (path.empty()) + return ""; + int p = path.rfind('/'); + int e = path.rfind('\\'); + p = std::max(p, e); + p += 1; + // include suffix + if (include_suffix) + return path.substr(p); + int u = path.rfind('.'); + if (u == -1) + return path.substr(p); + + if (u <= p) + u = path.size(); + return path.substr(p, u - p); +} + +void performance_yolop(const string &engine_file, YoloP::Type type, int gpuid, const string &input_dir) +{ auto detector = YoloP::create_detector(engine_file, type, gpuid, 0.4, 0.5); - if(detector == nullptr){ + if (detector == nullptr) + { printf("detector is nullptr.\n"); return; } - auto image = cv::imread("imgs/1.jpg"); + vector files_; + files_.reserve(100); + cv::glob(input_dir + "/*.jpg", files_, true); + vector files(files_.begin(), files_.end()); + + if (files.empty()) + { + printf("No image files found in the input directory: %s\n", input_dir.c_str()); + return; + } + + std::vector images; + for (const auto &file : files) + { + auto image = cv::imread(file); + if (image.empty()) + { + printf("Error reading image file: %s\n", file.c_str()); + continue; + } + images.emplace_back(image); + } + if (images.empty()) + { + printf("No valid images to process after reading files from: %s\n", input_dir.c_str()); + return; + } + int batch = 8; + for (int i = images.size(); i < batch; ++i) + images.push_back(images[i % images.size()]); + YoloP::PTMM res; // warmup - for(int i = 0; i < 10; ++i) - res = detector->detect(image); + for (int i = 0; i < 10; ++i) + res = detector->detect(images[i % images.size()]); // 测试 100 轮 const int ntest = 100; auto start = std::chrono::steady_clock::now(); - for(int i = 0; i < ntest; ++i) - res = detector->detect(image); + for (int i = 0; i < ntest; ++i) + res = detector->detect(images[i % images.size()]); std::chrono::duration during = std::chrono::steady_clock::now() - start; double all_time = 1000.0 * during.count(); - float avg_time = all_time / ntest; - printf("Average time: %.2f ms, FPS: %.2f\n", engine_file.c_str(), avg_time, 1000 / avg_time); + float avg_time = all_time / ntest / images.size(); + printf("Average time for %s: %.2f ms, FPS: %.2f\n", engine_file.c_str(), avg_time, 1000 / avg_time); } - -void inference_yolop(const string& engine_file, YoloP::Type type, int gpuid){ +void inference_yolop(const string &engine_file, YoloP::Type type, int gpuid, const string &input_img, const string &output_dir) +{ auto detector = YoloP::create_detector(engine_file, type, gpuid, 0.4, 0.5); - if(detector == nullptr){ + if (detector == nullptr) + { printf("detector is nullptr.\n"); return; } - - auto image = cv::imread("imgs/6.jpg"); + auto image = cv::imread(input_img); + if (image.empty()) + { + printf("Error reading image file: %s\n", input_img.c_str()); + return; + } auto res = detector->detect(image); - YoloP::BoxArray& boxes = get<0>(res); - cv::Mat& drive_mask = get<1>(res); - cv::Mat& lane_mask = get<2>(res); + YoloP::BoxArray &boxes = get<0>(res); + cv::Mat &drive_mask = get<1>(res); + cv::Mat &lane_mask = get<2>(res); - for(auto& ibox : boxes) + for (auto &ibox : boxes) cv::rectangle(image, cv::Point(ibox.left, ibox.top), cv::Point(ibox.right, ibox.bottom), {0, 0, 255}, 2); + fs::create_directories(output_dir); - cv::imwrite("infer_res/res.jpg", image); - cv::imwrite("infer_res/drive.jpg", drive_mask); - cv::imwrite("infer_res/lane.jpg", lane_mask); -} + string file_name = get_file_name(input_img, false); + string save_img_path = cv::format("%s/%s.jpg", output_dir.c_str(), file_name.c_str()); + string save_drive_path = cv::format("%s/drive_%s.jpg", output_dir.c_str(), file_name.c_str()); + string save_lane_path = cv::format("%s/lane_%s.jpg", output_dir.c_str(), file_name.c_str()); + cv::imwrite(save_img_path, image); + cv::imwrite(save_drive_path, drive_mask); + cv::imwrite(save_lane_path, lane_mask); + printf("Save to %s, %s, %s\n", save_img_path.c_str(), save_drive_path.c_str(), save_lane_path.c_str()); +} \ No newline at end of file diff --git a/apps/app_yolov10.cpp b/apps/app_yolov10.cpp index 6332500..ab089cf 100644 --- a/apps/app_yolov10.cpp +++ b/apps/app_yolov10.cpp @@ -1,83 +1,147 @@ - #include #include #include "yolov10/yolov10.hpp" +#include + +namespace fs = std::filesystem; using namespace std; inline vector cocolabels = { - "person", "bicycle", "car", "motorcycle", "airplane", - "bus", "train", "truck", "boat", "traffic light", "fire hydrant", - "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", - "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", - "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", - "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", - "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", - "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", - "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", - "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", - "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", - "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", - "scissors", "teddy bear", "hair drier", "toothbrush" -}; - -inline std::tuple hsv2bgr(float h, float s, float v){ + "person", "bicycle", "car", "motorcycle", "airplane", + "bus", "train", "truck", "boat", "traffic light", "fire hydrant", + "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", + "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", + "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", + "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", + "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", + "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", + "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", + "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", + "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", + "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", + "scissors", "teddy bear", "hair drier", "toothbrush"}; + +inline std::tuple hsv2bgr(float h, float s, float v) +{ const int h_i = static_cast(h * 6); const float f = h * 6 - h_i; const float p = v * (1 - s); - const float q = v * (1 - f*s); + const float q = v * (1 - f * s); const float t = v * (1 - (1 - f) * s); float r, g, b; - switch (h_i) { - case 0:r = v; g = t; b = p;break; - case 1:r = q; g = v; b = p;break; - case 2:r = p; g = v; b = t;break; - case 3:r = p; g = q; b = v;break; - case 4:r = t; g = p; b = v;break; - case 5:r = v; g = p; b = q;break; - default:r = 1; g = 1; b = 1;break;} + switch (h_i) + { + case 0: + r = v; + g = t; + b = p; + break; + case 1: + r = q; + g = v; + b = p; + break; + case 2: + r = p; + g = v; + b = t; + break; + case 3: + r = p; + g = q; + b = v; + break; + case 4: + r = t; + g = p; + b = v; + break; + case 5: + r = v; + g = p; + b = q; + break; + default: + r = 1; + g = 1; + b = 1; + break; + } return make_tuple(static_cast(b * 255), static_cast(g * 255), static_cast(r * 255)); } -inline std::tuple random_color(int id){ - float h_plane = ((((unsigned int)id << 2) ^ 0x937151) % 100) / 100.0f;; +inline std::tuple random_color(int id) +{ + float h_plane = ((((unsigned int)id << 2) ^ 0x937151) % 100) / 100.0f; + ; float s_plane = ((((unsigned int)id << 3) ^ 0x315793) % 100) / 100.0f; return hsv2bgr(h_plane, s_plane, 1); } -inline string get_file_name(const string& path, bool include_suffix){ - if (path.empty()) return ""; +inline string get_file_name(const string &path, bool include_suffix) +{ + if (path.empty()) + return ""; int p = path.rfind('/'); int e = path.rfind('\\'); p = std::max(p, e); p += 1; - //include suffix + // include suffix if (include_suffix) return path.substr(p); int u = path.rfind('.'); if (u == -1) return path.substr(p); - if (u <= p) u = path.size(); + if (u <= p) + u = path.size(); return path.substr(p, u - p); } -void performance_v10(const string& engine_file, int gpuid){ +void performance_v10(const string &engine_file, int gpuid, const string &input_dir) +{ auto infer = YOLOV10::create_infer(engine_file, gpuid, 0.5); - if(infer == nullptr){ + if (infer == nullptr) + { printf("infer is nullptr.\n"); return; } + vector files_; + files_.reserve(100); + cv::glob(input_dir + "/*.jpg", files_, true); + vector files(files_.begin(), files_.end()); + + if (files.empty()) + { + printf("No .jpg image files found in the input directory: %s\n", input_dir.c_str()); + return; + } + + std::vector images; + for (const auto &file : files) + { + auto image = cv::imread(file); + if (image.empty()) + { + printf("Error reading image file: %s\n", file.c_str()); + continue; + } + images.emplace_back(image); + } - int batch = 1; - std::vector images{cv::imread("imgs/bus.jpg"), cv::imread("imgs/girl.jpg"), - cv::imread("imgs/group.jpg"), cv::imread("imgs/yq.jpg")}; + if (images.empty()) + { + printf("No valid images to process after reading files from: %s\n", input_dir.c_str()); + return; + } + int batch = 8; for (int i = images.size(); i < batch; ++i) - images.push_back(images[i % 4]); + images.push_back(images[i % images.size()]); // warmup vector> boxes_array; - for(int i = 0; i < 10; ++i) + for (int i = 0; i < 10; ++i) boxes_array = infer->commits(images); boxes_array.back().get(); boxes_array.clear(); @@ -85,7 +149,7 @@ void performance_v10(const string& engine_file, int gpuid){ // 测试 100 轮 const int ntest = 100; auto start = std::chrono::steady_clock::now(); - for(int i = 0; i < ntest; ++i) + for (int i = 0; i < ntest; ++i) boxes_array = infer->commits(images); // 等待全部推理结束 boxes_array.back().get(); @@ -93,38 +157,59 @@ void performance_v10(const string& engine_file, int gpuid){ std::chrono::duration during = std::chrono::steady_clock::now() - start; double all_time = 1000.0 * during.count(); float avg_time = all_time / ntest / images.size(); - printf("Average time: %.2f ms, FPS: %.2f\n", engine_file.c_str(), avg_time, 1000 / avg_time); + printf("Average time for %s: %.2f ms, FPS: %.2f\n", engine_file.c_str(), avg_time, 1000 / avg_time); } -void batch_inference_v10(const string& engine_file, int gpuid){ +void batch_inference_v10(const string &engine_file, int gpuid, const string &input_dir, const string &output_dir) +{ auto infer = YOLOV10::create_infer(engine_file, gpuid, 0.5); - if(infer == nullptr){ + if (infer == nullptr) + { printf("infer is nullptr.\n"); return; } + fs::create_directories(output_dir); + vector files_; files_.reserve(100); - cv::glob("imgs/*.jpg", files_, true); + cv::glob(input_dir + "/*.jpg", files_, true); vector files(files_.begin(), files_.end()); + if (files.empty()) + { + printf("No image files found in the input directory: %s\n", input_dir.c_str()); + return; + } + vector images; - for(const auto& file : files){ + for (const auto &file : files) + { auto image = cv::imread(file); + if (image.empty()) + { + printf("Error reading image file: %s\n", file.c_str()); + continue; + } images.emplace_back(image); } - + if (images.empty()) + { + printf("No valid images to process after reading files from: %s\n", input_dir.c_str()); + return; + } vector> boxes_array; boxes_array = infer->commits(images); // 等待全部推理结束 boxes_array.back().get(); - string root_res = "infer_res/yolov10"; - for(int i = 0; i < boxes_array.size(); ++i){ + for (int i = 0; i < boxes_array.size(); ++i) + { cv::Mat image = images[i]; auto boxes = boxes_array[i].get(); - for(auto & ibox : boxes){ + for (auto &ibox : boxes) + { cv::Scalar color; std::tie(color[0], color[1], color[2]) = random_color(ibox.label); cv::rectangle(image, cv::Point(ibox.left, ibox.top), cv::Point(ibox.right, ibox.bottom), color, 2); @@ -132,28 +217,35 @@ void batch_inference_v10(const string& engine_file, int gpuid){ auto name = cocolabels[ibox.label]; auto caption = cv::format("%s %.2f", name.c_str(), ibox.confidence); int text_width = cv::getTextSize(caption, 0, 1, 2, nullptr).width + 10; - cv::rectangle(image, cv::Point(ibox.left-2, ibox.top-32), cv::Point(ibox.left + text_width, ibox.top), color, -1); - cv::putText(image, caption, cv::Point(ibox.left, ibox.top-5), 0, 1, cv::Scalar::all(0), 2, 16); + cv::rectangle(image, cv::Point(ibox.left - 2, ibox.top - 32), cv::Point(ibox.left + text_width, ibox.top), color, -1); + cv::putText(image, caption, cv::Point(ibox.left, ibox.top - 5), 0, 1, cv::Scalar::all(0), 2, 16); } string file_name = get_file_name(files[i], false); - string save_path = cv::format("%s/%s.jpg", root_res.c_str(), file_name.c_str()); + string save_path = cv::format("%s/%s.jpg", output_dir.c_str(), file_name.c_str()); cv::imwrite(save_path, image); printf("Save to %s, %d object\n", save_path.c_str(), boxes.size()); } } - -void single_inference_v10(const string& engine_file, int gpuid){ +void single_inference_v10(const string &engine_file, int gpuid, const string &input_img, const string &output_img_path) +{ auto infer = YOLOV10::create_infer(engine_file, gpuid, 0.6); - if(infer == nullptr){ + if (infer == nullptr) + { printf("infer is nullptr.\n"); return; } - auto image = cv::imread("imgs/bus.jpg"); + auto image = cv::imread(input_img); + if (image.empty()) + { + printf("Error reading image file: %s\n", input_img.c_str()); + return; + } auto boxes = infer->commit(image).get(); - for(auto& ibox : boxes){ + for (auto &ibox : boxes) + { cv::Scalar color; std::tie(color[0], color[1], color[2]) = random_color(ibox.label); cv::rectangle(image, cv::Point(ibox.left, ibox.top), cv::Point(ibox.right, ibox.bottom), color, 2); @@ -161,11 +253,11 @@ void single_inference_v10(const string& engine_file, int gpuid){ auto name = cocolabels[ibox.label]; auto caption = cv::format("%s %.2f", name.c_str(), ibox.confidence); int text_width = cv::getTextSize(caption, 0, 1, 2, nullptr).width + 10; - cv::rectangle(image, cv::Point(ibox.left-2, ibox.top-32), cv::Point(ibox.left + text_width, ibox.top), color, -1); - cv::putText(image, caption, cv::Point(ibox.left, ibox.top-5), 0, 1, cv::Scalar::all(0), 2, 16); + cv::rectangle(image, cv::Point(ibox.left - 2, ibox.top - 32), cv::Point(ibox.left + text_width, ibox.top), color, -1); + cv::putText(image, caption, cv::Point(ibox.left, ibox.top - 5), 0, 1, cv::Scalar::all(0), 2, 16); } - cv::imwrite("infer_res/result.jpg", image); -} - - - + fs::path path(output_img_path); + fs::create_directories(path.parent_path()); + cv::imwrite(output_img_path, image); + printf("Save to %s\n", output_img_path.c_str()); +} \ No newline at end of file diff --git a/config.yaml b/config.yaml new file mode 100644 index 0000000..8bf8753 --- /dev/null +++ b/config.yaml @@ -0,0 +1,82 @@ +tasks: + # - task: "rtdetr" + # subtasks: + # - type: "performance" + # engine_file: "rtdetr_r50vd_6x_coco_dynamic_fp16.trt" + # gpuid: 0 + # input_dir: "/home/e300/mahmood/code/mLinfer/workspace/imgs" + # - type: "batch_inference" + # engine_file: "rtdetr_r50vd_6x_coco_dynamic_fp16.trt" + # gpuid: 0 + # input_dir: "/home/e300/mahmood/code/mLinfer/workspace/imgs" + # output_dir: "/home/e300/mahmood/code/mLinfer/workspace/result_images_rtdetr" + # - type: "single_inference" + # engine_file: "rtdetr_r50vd_6x_coco_dynamic_fp16.trt" + # gpuid: 0 + # input_img: "/home/e300/Downloads/1684008869189.jpg" + # output_img_path: "/home/e300/mahmood/code/mLinfer/workspace/result_images_single_rtdetr/1684008869189.jpg" + - task: "yolov10" + subtasks: + - type: "performance_v10" + engine_file: "/home/e300/mahmood/code/mLinfer/workspace/yolov10s.trt" + gpuid: 0 + input_dir: "/home/e300/mahmood/code/mLinfer/workspace/imgs" + - type: "batch_inference_v10" + engine_file: "/home/e300/mahmood/code/mLinfer/workspace/yolov10s.trt" + gpuid: 0 + input_dir: "/home/e300/mahmood/code/mLinfer/workspace/imgs" + output_dir: "/home/e300/mahmood/code/mLinfer/workspace/result_images_yolov10" + - type: "single_inference_v10" + engine_file: "/home/e300/mahmood/code/mLinfer/workspace/yolov10s.trt" + gpuid: 0 + input_img: "/home/e300/Downloads/1684008869189.jpg" + output_img_path: "/home/e300/mahmood/code/mLinfer/workspace/result_images_single_yolov10/1684008869189.jpg" + # - task: "yolov10" + # subtasks: + # - type: "performance_v10" + # engine_file: "yolov10n.trt" + # gpuid: 0 + # - task: "yolo" + # subtasks: + # - type: "batch_inference" + # engine_file: "/home/e300/mahmood/code/Linfer/workspace/yolov8s.trt" + # gpuid: 0 + # yolo_type: "V8" + # input_dir: "/home/e300/mahmood/code/mLinfer/workspace/imgs" + # output_dir: "/home/e300/mahmood/code/mLinfer/workspace/result_images" + # - type: "performance" + # engine_file: "/home/e300/mahmood/code/Linfer/workspace/yolov8s.trt" + # gpuid: 0 + # yolo_type: "V8" + # input_dir: "/home/e300/mahmood/code/mLinfer/workspace/imgs" + # - type: "single_inference" + # engine_file: "/home/e300/mahmood/code/Linfer/workspace/yolov8s.trt" + # gpuid: 0 + # yolo_type: "V8" + # input_img: "/home/e300/Downloads/1684008869189.jpg" + # output_img_path: "/home/e300/mahmood/code/mLinfer/workspace/result_images_single/1684008869189.jpg" + # - task: "track" + # subtasks: + # - type: "inference_bytetrack" + # engine_file: "/home/e300/mahmood/code/Linfer/workspace/yolov8s.trt" + # gpuid: 0 + # yolo_type: "V8" + # video_file: "/home/e300/mahmood/code/Linfer/workspace/videos/snow.mp4" + # output_save_path: "" + # - task: "yolop" + # subtasks: + # - type: "inference_yolop" + # engine_file: "yolopv2-480x640.trt" + # yolo_type: "V2" + # gpuid: 0 + # - task: "seg" + # subtasks: + # - type: "performance_seg" + # engine_file: "/home/e300/mahmood/code/mLinfer/workspace/mobileseg_mbn3.trt" + # gpuid: 0 + # input_dir: "/home/e300/mahmood/code/mLinfer/workspace/imgs" + # - type: "inference_seg" + # engine_file: "/home/e300/mahmood/code/mLinfer/workspace/mobileseg_mbn3.trt" + # gpuid: 0 + # input_img: "/home/e300/mahmood/code/mLinfer/workspace/imgs/frame_3.jpg" + # output_img_path: "/home/e300/mahmood/code/mLinfer/workspace/result_images_seg/seg_frame_3.jpg" \ No newline at end of file diff --git a/main.cpp b/main.cpp index 057182f..9675550 100755 --- a/main.cpp +++ b/main.cpp @@ -1,88 +1,254 @@ - +#include +#include +#include +#include +#include #include #include "apps/yolo/yolo.hpp" #include "apps/yolop/yolop.hpp" +#include "apps/rtdetr/rtdetr.hpp" // Include the header for RTDETR using namespace std; -void performance_v10(const string& engine_file, int gpuid); -void batch_inference_v10(const string& engine_file, int gpuid); -void single_inference_v10(const string& engine_file, int gpuid); -void performance(const string& engine_file, int gpuid); -void batch_inference(const string& engine_file, int gpuid); -void single_inference(const string& engine_file, int gpuid); -void performance(const string& engine_file, int gpuid, Yolo::Type type); -void batch_inference(const string& engine_file, int gpuid, Yolo::Type type); -void single_inference(const string& engine_file, int gpuid, Yolo::Type type); -void inference_bytetrack(const string& engine_file, int gpuid, Yolo::Type type, const string& video_file); -void infer_track(int Mode, const string& path); -void inference_yolop(const string& engine_file, YoloP::Type type, int gpuid); -void performance_yolop(const string& engine_file, YoloP::Type type, int gpuid); -void inference_seg(const string& engine_file, int gpuid); -void performance_seg(const string& engine_file, int gpuid); +void performance_v10(const string &engine_file, int gpuid, const string &input_dir); +void batch_inference_v10(const string &engine_file, int gpuid, const string &input_dir, const string &output_dir); +void single_inference_v10(const string &engine_file, int gpuid, const string &input_img, const string &output_img_path); +void performance(const string &engine_file, int gpuid, const string &input_dir); +void batch_inference(const string &engine_file, int gpuid, const string &input_dir, const string &output_dir); +void single_inference(const string &engine_file, int gpuid, const string &input_img, const string &output_img_path); +void performance(const string &engine_file, int gpuid, Yolo::Type type, const string &input_dir); +void batch_inference(const string &engine_file, int gpuid, Yolo::Type type, const string &input_dir, const string &output_dir); +void single_inference(const string &engine_file, int gpuid, Yolo::Type type, const string &input_img, const string &output_img_path); +void inference_bytetrack(const string &engine_file, int gpuid, Yolo::Type type, const string &video_file, const string &output_save_path); +void infer_track(int Mode, const string &path); +void performance_yolop(const string &engine_file, YoloP::Type type, int gpuid, const string &input_dir); +void inference_yolop(const string &engine_file, YoloP::Type type, int gpuid, const string &input_img, const string &output_dir); +void performance_seg(const string &engine_file, int gpuid, const string &input_dir); +void inference_seg(const string &engine_file, int gpuid, const string &input_img, const string &output_img_path); bool test_ptq(); -void test_rtdetr(){ -// batch_inference("rtdetr_r50vd_6x_coco_dynamic_fp16.trt", 0); -// single_inference("rtdetr_r50vd_6x_coco_dynamic_fp16.trt", 0); - performance("rtdetr_r50vd_6x_coco_dynamic_fp16.trt", 0); +// Helper function to convert string to Yolo::Type +Yolo::Type stringToYoloType(const string &typeStr) +{ + if (typeStr == "V5") + return Yolo::Type::V5; + if (typeStr == "X") + return Yolo::Type::X; + if (typeStr == "V7") + return Yolo::Type::V7; + if (typeStr == "V8") + return Yolo::Type::V8; + throw std::runtime_error("Unknown Yolo type: " + typeStr); } -void test_yolov10(){ -// batch_inference_v10("yolov10l.trt", 0); -// single_inference_v10("yolov10l.trt", 0); - performance_v10("yolov10n.trt", 0); +// Helper function to convert string to YoloP::Type +YoloP::Type stringToYoloPType(const string &typeStr) +{ + if (typeStr == "V1") + return YoloP::Type::V1; + if (typeStr == "V2") + return YoloP::Type::V2; + throw std::runtime_error("Unknown YoloP type: " + typeStr); } -void test_yolo(){ -// batch_inference("yolov5s.trt", 0, Yolo::Type::V5); - performance("yolov5s.trt", 0, Yolo::Type::V5); -// batch_inference("yolov5s_ptq.trt", 0, Yolo::Type::V5); -// batch_inference("yolov5m.trt", 0, Yolo::Type::V5); -// performance("yolov5m.trt", 0, Yolo::Type::V5); -// batch_inference("yolox_s.trt", 0, Yolo::Type::X); -// performance("yolox_s.trt", 0, Yolo::Type::X); -// batch_inference("yolox_m.trt", 0, Yolo::Type::X); -// performance("yolox_m.trt", 0, Yolo::Type::X); -// batch_inference("yolov7.trt", 0, Yolo::Type::V7); -// performance("yolov7.trt", 0, Yolo::Type::V7); -// batch_inference("yolov7_qat.trt", 0, Yolo::Type::V7); -// performance("yolov7_qat.trt", 0, Yolo::Type::V7); -// batch_inference("yolov8n.trt", 0, Yolo::Type::V8); -// performance("yolov8n.trt", 0, Yolo::Type::V8); -// batch_inference("yolov8s.trt", 0, Yolo::Type::V8); -// performance("yolov8s.trt", 0, Yolo::Type::V8); -// batch_inference("yolov8l.trt", 0, Yolo::Type::V8); -// performance("yolov8l.trt", 0, Yolo::Type::V8); -// single_inference("yolov8l.trt", 0, Yolo::Type::V8); -} +int main(int argc, char *argv[]) +{ + if (argc != 2) + { + cerr << "Usage: " << argv[0] << " " << endl; + return 1; + } -void test_track(){ -// inference_bytetrack("yolov8s.trt", 0, Yolo::Type::V8, "videos/palace.mp4"); - infer_track(2, "Woman/img/%04d.jpg"); -} + const string config_file_path = argv[1]; -void test_yolop(){ - inference_yolop("yolopv2-480x640.trt", YoloP::Type::V2, 0); -// inference_yolop("yolop-640.trt", YoloP::Type::V1, 0); -// performance_yolop("yolopv2-480x640.trt", YoloP::Type::V2, 0); -// performance_yolop("yolop-640.trt", YoloP::Type::V1, 0); -} + try + { + YAML::Node config = YAML::LoadFile(config_file_path); -void test_seg(){ - inference_seg("ppliteseg_stdc2.trt", 0); -// inference_seg("mobileseg_mbn3.trt", 0); -// performance_seg("ppliteseg_stdc2.trt", 0); -// performance_seg("mobileseg_mbn3.trt", 0); -} + if (!config["tasks"] || !config["tasks"].IsSequence()) + { + cerr << "Error: Invalid or missing 'tasks' in config.yaml." << endl; + return 1; + } + + for (const auto &task_node : config["tasks"]) + { + string task_name = task_node["task"].as(); + cout << "Executing Task: " << task_name << endl; + + if (!task_node["subtasks"] || !task_node["subtasks"].IsSequence()) + { + cerr << "Error: Invalid or missing 'subtasks' in config.yaml." << endl; + continue; + } + + for (const auto &subtask_node : task_node["subtasks"]) + { + string subtask_type = subtask_node["type"].as(); + + cout << " Subtask: " << subtask_type << endl; + if (task_name == "rtdetr") + { + string engine_file = subtask_node["engine_file"].as(); + int gpuid = subtask_node["gpuid"].as(); + + if (subtask_type == "performance") + { + string input_dir = subtask_node["input_dir"].as(); + performance(engine_file, gpuid, input_dir); + } + else if (subtask_type == "batch_inference") + { + string input_dir = subtask_node["input_dir"].as(); + string output_dir = subtask_node["output_dir"].as(); + batch_inference(engine_file, gpuid, input_dir, output_dir); + } + else if (subtask_type == "single_inference") + { + string input_img = subtask_node["input_img"].as(); + string output_img_path = subtask_node["output_img_path"].as(); + single_inference(engine_file, gpuid, input_img, output_img_path); + } + else + { + cerr << " Error: Unknown subtask type for rtdetr: " << subtask_type << endl; + } + } + else if (task_name == "yolov10") + { + string engine_file = subtask_node["engine_file"].as(); + int gpuid = subtask_node["gpuid"].as(); + + if (subtask_type == "performance_v10") + { + string input_dir = subtask_node["input_dir"].as(); + performance_v10(engine_file, gpuid, input_dir); + } + else if (subtask_type == "batch_inference_v10") + { + string input_dir = subtask_node["input_dir"].as(); + string output_dir = subtask_node["output_dir"].as(); + batch_inference_v10(engine_file, gpuid, input_dir, output_dir); + } + else if (subtask_type == "single_inference_v10") + { + string input_img = subtask_node["input_img"].as(); + string output_img_path = subtask_node["output_img_path"].as(); + single_inference_v10(engine_file, gpuid, input_img, output_img_path); + } + else + { + cerr << " Error: Unknown subtask type for yolov10: " << subtask_type << endl; + } + } + else if (task_name == "yolo") + { + string engine_file = subtask_node["engine_file"].as(); + int gpuid = subtask_node["gpuid"].as(); + string yolo_type_str = subtask_node["yolo_type"].as(); + Yolo::Type yolo_type = stringToYoloType(yolo_type_str); + + if (subtask_type == "batch_inference") + { + string input_dir = subtask_node["input_dir"].as(); + string output_dir = subtask_node["output_dir"].as(); + batch_inference(engine_file, gpuid, yolo_type, input_dir, output_dir); + } + else if (subtask_type == "performance") + { + string input_dir = subtask_node["input_dir"].as(); + performance(engine_file, gpuid, yolo_type, input_dir); + } + else if (subtask_type == "single_inference") + { + string input_img = subtask_node["input_img"].as(); + string output_img_path = subtask_node["output_img_path"].as(); + single_inference(engine_file, gpuid, yolo_type, input_img, output_img_path); + } + else + { + cerr << " Error: Unknown subtask type for yolo: " << subtask_type << endl; + } + } + else if (task_name == "track") + { + string engine_file = subtask_node["engine_file"].as(); + int gpuid = subtask_node["gpuid"].as(); + string yolo_type_str = subtask_node["yolo_type"].as(); + Yolo::Type yolo_type = stringToYoloType(yolo_type_str); + string video_file = subtask_node["video_file"].as(); + string output_save_path = subtask_node["output_save_path"].as(); + + if (subtask_type == "inference_bytetrack") + { + inference_bytetrack(engine_file, gpuid, yolo_type, video_file, output_save_path); + } + else + { + cerr << " Error: Unknown subtask type for track: " << subtask_type << endl; + } + } + else if (task_name == "yolop") + { + string engine_file = subtask_node["engine_file"].as(); + int gpuid = subtask_node["gpuid"].as(); + string yolop_type_str = subtask_node["yolo_type"].as(); + YoloP::Type yolop_type = stringToYoloPType(yolop_type_str); + + if (subtask_type == "inference_yolop") + { + string input_img = subtask_node["input_img"].as(); + string output_dir = subtask_node["output_dir"].as(); + inference_yolop(engine_file, yolop_type, gpuid, input_img, output_dir); + } + else if (subtask_type == "performance_yolop") + { + string input_dir = subtask_node["input_dir"].as(); + performance_yolop(engine_file, yolop_type, gpuid, input_dir); + } + else + { + cerr << " Error: Unknown subtask type for yolop: " << subtask_type << endl; + } + } + else if (task_name == "seg") + { + string engine_file = subtask_node["engine_file"].as(); + int gpuid = subtask_node["gpuid"].as(); + if (subtask_type == "inference_seg") + { + string input_img = subtask_node["input_img"].as(); + string output_img_path = subtask_node["output_img_path"].as(); + inference_seg(engine_file, gpuid, input_img, output_img_path); + } + else if (subtask_type == "performance_seg") + { + string input_dir = subtask_node["input_dir"].as(); + performance_seg(engine_file, gpuid, input_dir); + } + else + { + cerr << " Error: Unknown subtask type for seg: " << subtask_type << endl; + } + } + else + { + cerr << " Error: Unknown task: " << task_name << endl; + } + } + cout << endl; + } + } + catch (const YAML::Exception &e) + { + cerr << "Error parsing YAML file: " << e.what() << endl; + return 1; + } + catch (const std::runtime_error &e) + { + cerr << "Error: " << e.what() << endl; + return 1; + } -int main(){ -// test_rtdetr(); - test_yolov10(); -// test_yolo(); -// test_yolop(); -// test_track(); -// test_ptq(); - test_seg(); return 0; } \ No newline at end of file diff --git a/workspace/compile_engine.sh b/workspace/compile_engine.sh index 35ee455..2a67d17 100644 --- a/workspace/compile_engine.sh +++ b/workspace/compile_engine.sh @@ -210,9 +210,9 @@ trtexec --onnx=./onnx_models/yolov10l.onnx \ # echo "************************ compile PPSeg model ***************************" # echo "" # trtexec --onnx=./onnx_models/mobileseg_mbn3.onnx \ -# --saveEngine=./mobileseg_mbn3.trt \ -# --buildOnly \ -# --fp16 + # --saveEngine=./mobileseg_mbn3.trt \ + # --buildOnly \ + # --fp16 # echo "" # trtexec --onnx=./onnx_models/ppliteseg_stdc2.onnx \