Merge pull request #79 from laugh12321/dev

Refactor tensorrt_yolo.infer module and remove deprecated typing types (Tuple, List, Set, Dict)

Author: laugh12321

CMakeLists.txt

@@ -2,7 +2,7 @@
 cmake_minimum_required(VERSION 3.15.0)                  # 设置CMake的最低版本要求
 cmake_policy(SET CMP0091 NEW)                           # 允许在CMake 3.10+中自动设置项目名作为二进制目录名
 cmake_policy(SET CMP0146 OLD)                           # 忽略对find_package的过时警告
-project(TensorRT-YOLO VERSION 5.1.0 LANGUAGES CXX CUDA) # 定义项目名称、版本和使用的编程语言（C++和CUDA）
+project(TensorRT-YOLO VERSION 5.1.1 LANGUAGES CXX CUDA) # 定义项目名称、版本和使用的编程语言（C++和CUDA）
 
 # 设置 C++ 标准
 set(CMAKE_CXX_STANDARD 17)                              # 设置C++标准为17

README.en.md

@@ -132,13 +132,13 @@ English | [简体中文](README.md)
 #### Python Inference Example
 
 > [!NOTE] 
-> `DeployDet`, `DeployOBB`, `DeploySeg`, and `DeployPose` correspond to detection (Detect), oriented bounding box (OBB), segmentation (Segment), and pose estimation (Pose) models, respectively.
+> `DeployDet`, `DeployOBB`, `DeploySeg`, `DeployPose` and `DeployCls` correspond to detection (Detect), oriented bounding box (OBB), segmentation (Segment), pose estimation (Pose) and image classification (Classify) models, respectively.
 >
 > For these models, the `CG` version utilizes CUDA Graph to further accelerate the inference process, but please note that this feature is limited to static models.
 
 ```python
 import cv2
-from tensorrt_yolo.infer import DeployDet, generate_labels_with_colors, visualize
+from tensorrt_yolo.infer import DeployDet, generate_labels, visualize
 
 # Initialize the model
 model = DeployDet("yolo11n-with-plugin.engine")
@@ -148,15 +148,15 @@ im = cv2.imread("test_image.jpg")
 result = model.predict(cv2.cvtColor(im, cv2.COLOR_BGR2RGB))
 print(f"==> detect result: {result}")
 # Visualization
-labels = generate_labels_with_colors("labels.txt")
+labels = generate_labels("labels.txt")
 vis_im = visualize(im, result, labels)
 cv2.imwrite("vis_image.jpg", vis_im)
 ```
 
 ### C++ SDK Quick Start<div id="quick-start-cpp"></div>
 
 > [!NOTE] 
-> `DeployDet`, `DeployOBB`, `DeploySeg`, and `DeployPose` correspond to detection (Detect), oriented bounding box (OBB), segmentation (Segment), and pose estimation (Pose) models, respectively.
+> `DeployDet`, `DeployOBB`, `DeploySeg`, `DeployPose` and `DeployCls` correspond to detection (Detect), oriented bounding box (OBB), segmentation (Segment), pose estimation (Pose) and image classification (Classify) models, respectively.
 >
 > For these models, the `CG` version utilizes CUDA Graph to further accelerate the inference process, but please note that this feature is limited to static models.
 

README.md

@@ -131,13 +131,13 @@
 #### Python 推理示例
 
 > [!NOTE] 
-> `DeployDet`、`DeployOBB`、`DeploySeg` 和 `DeployPose` 分别对应于检测（Detect）、方向边界框（OBB）、分割（Segment）和姿态估计（Pose）模型。
+> `DeployDet`、`DeployOBB`、`DeploySeg`、`DeployPose` 和 `DeployCls` 分别对应于检测（Detect）、方向边界框（OBB）、分割（Segment）、姿态估计（Pose）和图像分类（Classify）模型。
 >
 > 对于这些模型，`CG` 版本利用 CUDA Graph 来进一步加速推理过程，但请注意，这一功能仅限于静态模型。
 
 ```python
 import cv2
-from tensorrt_yolo.infer import DeployDet, generate_labels_with_colors, visualize
+from tensorrt_yolo.infer import DeployDet, generate_labels, visualize
 
 # 初始化模型
 model = DeployDet("yolo11n-with-plugin.engine")
@@ -147,15 +147,15 @@ im = cv2.imread("test_image.jpg")
 result = model.predict(cv2.cvtColor(im, cv2.COLOR_BGR2RGB))
 print(f"==> detect result: {result}")
 # 可视化
-labels = generate_labels_with_colors("labels.txt")
+labels = generate_labels("labels.txt")
 vis_im = visualize(im, result, labels)
 cv2.imwrite("vis_image.jpg", vis_im)
 ```
 
 ### C++ SDK快速开始<div id="quick-start-cpp"></div>
 
 > [!NOTE] 
-> `DeployDet`、`DeployOBB`、`DeploySeg` 和 `DeployPose` 分别对应于检测（Detect）、方向边界框（OBB）、分割（Segment）和姿态估计（Pose）模型。
+> `DeployDet`、`DeployOBB`、`DeploySeg`、`DeployPose` 和 `DeployCls` 分别对应于检测（Detect）、方向边界框（OBB）、分割（Segment）、姿态估计（Pose）和图像分类（Classify）模型。
 >
 > 对于这些模型，`CG` 版本利用 CUDA Graph 来进一步加速推理过程，但请注意，这一功能仅限于静态模型。
 

assets/detect.jpg

@@ -3,7 +3,6 @@
 #include <iostream>
 #include <memory>
 #include <opencv2/opencv.hpp>
-#include <random>
 
 #include "deploy/utils/utils.hpp"
 #include "deploy/vision/inference.hpp"
@@ -32,39 +31,31 @@ void createOutputDirectory(const std::string& outputPath) {
     }
 }
 
-// Generate label and color pairs
-std::vector<std::pair<std::string, cv::Scalar>> generateLabelColorPairs(const std::string& labelFile) {
-    std::ifstream                                   file(labelFile);
-    std::vector<std::pair<std::string, cv::Scalar>> labelColorPairs;
+// Generate label
+std::vector<std::string> generateLabels(const std::string& labelFile) {
+    std::ifstream            file(labelFile);
+    std::vector<std::string> labels;
     if (!file.is_open()) {
         throw std::runtime_error("Failed to open labels file: " + labelFile);
     }
 
-    auto generateRandomColor = []() {
-        std::random_device                 rd;
-        std::mt19937                       gen(rd());
-        std::uniform_int_distribution<int> dis(0, 255);
-        return cv::Scalar(dis(gen), dis(gen), dis(gen));
-    };
-
     std::string label;
     while (std::getline(file, label)) {
-        labelColorPairs.emplace_back(label, generateRandomColor());
+        labels.emplace_back(label);
     }
-    return labelColorPairs;
+    return labels;
 }
 
 // Visualize inference results
-void visualize(cv::Mat& image, deploy::ClsResult& result, std::vector<std::pair<std::string, cv::Scalar>>& labelColorPairs) {
+void visualize(cv::Mat& image, deploy::ClsResult& result, std::vector<std::string>& labels) {
     for (size_t i = 0; i < result.num; ++i) {
         int         cls       = result.classes[i];
         float       score     = result.scores[i];
-        auto&       label     = labelColorPairs[cls].first;
-        auto&       color     = labelColorPairs[cls].second;
-        std::string labelText = label + " " + cv::format("%.2f", score);
+        auto&       label     = labels[cls];
+        std::string labelText = label + " " + cv::format("%.3f", score);
 
         // Draw rectangle and label
-        cv::putText(image, labelText, cv::Point(5, 32 + i * 32), cv::FONT_HERSHEY_SIMPLEX, 0.6, color, 1);
+        cv::putText(image, labelText, cv::Point(5, 32 + i * 32), cv::FONT_HERSHEY_SIMPLEX, 0.6, cv::Scalar(251, 81, 163), 1);
     }
 }
 
@@ -119,7 +110,7 @@ int main(int argc, char** argv) {
             throw std::runtime_error("Input path does not exist or is not a regular file/directory: " + inputPath);
         }
 
-        std::vector<std::pair<std::string, cv::Scalar>> labels;
+        std::vector<std::string> labels;
         if (!outputPath.empty()) {
             if (labelPath.empty()) {
                 throw std::runtime_error("Please provide a labels file using -l or --labels.");
@@ -128,7 +119,7 @@ int main(int argc, char** argv) {
                 throw std::runtime_error("Label path does not exist: " + labelPath);
             }
 
-            labels = generateLabelColorPairs(labelPath);
+            labels = generateLabels(labelPath);
             createOutputDirectory(outputPath);
         }
 

assets/obb.png

@@ -30,7 +30,7 @@
 from loguru import logger
 from rich.progress import track
 
-from tensorrt_yolo.infer import CpuTimer, DeployCGCls, DeployCls, GpuTimer, generate_labels_with_colors, image_batches, visualize
+from tensorrt_yolo.infer import CpuTimer, DeployCGCls, DeployCls, GpuTimer, generate_labels, image_batches, visualize
 
 
 def main():
@@ -54,7 +54,7 @@ def main():
     if args.output:
         output_dir = Path(args.output)
         output_dir.mkdir(parents=True, exist_ok=True)
-        args.labels = generate_labels_with_colors(args.labels)
+        args.labels = generate_labels(args.labels)
 
     model = DeployCGCls(args.engine) if args.cudaGraph else DeployCls(args.engine)
 

assets/pose.jpg

@@ -3,7 +3,6 @@
 #include <iostream>
 #include <memory>
 #include <opencv2/opencv.hpp>
-#include <random>
 
 #include "deploy/utils/utils.hpp"
 #include "deploy/vision/inference.hpp"
@@ -32,44 +31,36 @@ void createOutputDirectory(const std::string& outputPath) {
     }
 }
 
-// Generate label and color pairs
-std::vector<std::pair<std::string, cv::Scalar>> generateLabelColorPairs(const std::string& labelFile) {
-    std::ifstream                                   file(labelFile);
-    std::vector<std::pair<std::string, cv::Scalar>> labelColorPairs;
+// Generate label
+std::vector<std::string> generateLabels(const std::string& labelFile) {
+    std::ifstream            file(labelFile);
+    std::vector<std::string> labels;
     if (!file.is_open()) {
         throw std::runtime_error("Failed to open labels file: " + labelFile);
     }
 
-    auto generateRandomColor = []() {
-        std::random_device                 rd;
-        std::mt19937                       gen(rd());
-        std::uniform_int_distribution<int> dis(0, 255);
-        return cv::Scalar(dis(gen), dis(gen), dis(gen));
-    };
-
     std::string label;
     while (std::getline(file, label)) {
-        labelColorPairs.emplace_back(label, generateRandomColor());
+        labels.emplace_back(label);
     }
-    return labelColorPairs;
+    return labels;
 }
 
 // Visualize inference results
-void visualize(cv::Mat& image, deploy::DetResult& result, std::vector<std::pair<std::string, cv::Scalar>>& labelColorPairs) {
+void visualize(cv::Mat& image, deploy::DetResult& result, std::vector<std::string>& labels) {
     for (size_t i = 0; i < result.num; ++i) {
         auto&       box       = result.boxes[i];
         int         cls       = result.classes[i];
         float       score     = result.scores[i];
-        auto&       label     = labelColorPairs[cls].first;
-        auto&       color     = labelColorPairs[cls].second;
-        std::string labelText = label + " " + cv::format("%.2f", score);
+        auto&       label     = labels[cls];
+        std::string labelText = label + " " + cv::format("%.3f", score);
 
         // Draw rectangle and label
         int      baseLine;
         cv::Size labelSize = cv::getTextSize(labelText, cv::FONT_HERSHEY_SIMPLEX, 0.6, 1, &baseLine);
-        cv::rectangle(image, cv::Point(box.left, box.top), cv::Point(box.right, box.bottom), color, 2, cv::LINE_AA);
-        cv::rectangle(image, cv::Point(box.left, box.top - labelSize.height), cv::Point(box.left + labelSize.width, box.top), color, -1);
-        cv::putText(image, labelText, cv::Point(box.left, box.top), cv::FONT_HERSHEY_SIMPLEX, 0.6, cv::Scalar(255, 255, 255), 1);
+        cv::rectangle(image, cv::Point(box.left, box.top), cv::Point(box.right, box.bottom), cv::Scalar(251, 81, 163), 2, cv::LINE_AA);
+        cv::rectangle(image, cv::Point(box.left, box.top - labelSize.height), cv::Point(box.left + labelSize.width, box.top), cv::Scalar(125, 40, 81), -1);
+        cv::putText(image, labelText, cv::Point(box.left, box.top), cv::FONT_HERSHEY_SIMPLEX, 0.6, cv::Scalar(253, 168, 208), 1);
     }
 }
 
@@ -124,7 +115,7 @@ int main(int argc, char** argv) {
             throw std::runtime_error("Input path does not exist or is not a regular file/directory: " + inputPath);
         }
 
-        std::vector<std::pair<std::string, cv::Scalar>> labels;
+        std::vector<std::string> labels;
         if (!outputPath.empty()) {
             if (labelPath.empty()) {
                 throw std::runtime_error("Please provide a labels file using -l or --labels.");
@@ -133,7 +124,7 @@ int main(int argc, char** argv) {
                 throw std::runtime_error("Label path does not exist: " + labelPath);
             }
 
-            labels = generateLabelColorPairs(labelPath);
+            labels = generateLabels(labelPath);
             createOutputDirectory(outputPath);
         }