Skip to content

Commit 6863869

Browse files
hamdiboukamchahamdiboukamcha
committed
Init Project
1 parent 198bcb5 commit 6863869

File tree

11 files changed

+1582
-0
lines changed

11 files changed

+1582
-0
lines changed

‎CMakeLists.txt‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
cmake_minimum_required(VERSION 3.18)
2+
3+
# Project declaration with C++ and CUDA support
4+
project(YOLOv11TRT LANGUAGES CXX CUDA)
5+
6+
# Set C++ standard to C++17
7+
set(CMAKE_CXX_STANDARD 17)
8+
set(CMAKE_CXX_STANDARD_REQUIRED ON)
9+
set(CMAKE_CXX_EXTENSIONS OFF)
10+
11+
# Define the path to TensorRT installation
12+
set(TENSORRT_PATH "F:/Program Files/TensorRT-8.6.1.6") # Update this to the actual path for TensorRT
13+
14+
# Define the path to OpenCV installation
15+
16+
# Allow overriding TensorRT and OpenCV paths via command line
17+
# e.g., cmake -DTENSORRT_PATH="path/to/TensorRT" -DOpenCV_DIR="path/to/OpenCV" ..
18+
option(TENSORRT_PATH_OPTION "Path to TensorRT installation" ${TENSORRT_PATH})
19+
set(TENSORRT_PATH ${TENSORRT_PATH_OPTION} CACHE PATH "Path to TensorRT installation")
20+
21+
# Find OpenCV
22+
find_package(OpenCV REQUIRED)
23+
if(NOT OpenCV_FOUND)
24+
message(FATAL_ERROR "OpenCV not found. Please install OpenCV or set OpenCV_DIR.")
25+
endif()
26+
27+
# Find CUDA
28+
find_package(CUDA REQUIRED)
29+
if(NOT CUDA_FOUND)
30+
message(FATAL_ERROR "CUDA not found. Please install the CUDA Toolkit.")
31+
endif()
32+
33+
# Include directories for TensorRT
34+
include_directories(${TENSORRT_PATH}/include)
35+
36+
# Include directory for your project
37+
include_directories(${CMAKE_SOURCE_DIR}/include)
38+
39+
# Define source files (including CUDA sources)
40+
set(SOURCES
41+
main.cpp
42+
src/yolov11.cpp
43+
src/preprocess.cu
44+
)
45+
46+
# Create executable (CMake handles CUDA sources automatically)
47+
add_executable(${PROJECT_NAME} ${SOURCES} ${HEADERS})
48+
49+
# Define API_EXPORTS macro
50+
target_compile_definitions(${PROJECT_NAME} PRIVATE API_EXPORTS)
51+
52+
# Specify include directories (modern CMake approach)
53+
target_include_directories(${PROJECT_NAME} PRIVATE
54+
src/
55+
${OpenCV_INCLUDE_DIRS}
56+
${CUDA_INCLUDE_DIRS}
57+
${TENSORRT_PATH}/include
58+
)
59+
60+
# Link TensorRT libraries
61+
# Specify full paths to TensorRT libraries to avoid relying on link_directories
62+
set(TENSORRT_LIBS
63+
"${TENSORRT_PATH}/lib/nvinfer.lib"
64+
"${TENSORRT_PATH}/lib/nvonnxparser.lib"
65+
"${TENSORRT_PATH}/lib/nvparsers.lib"
66+
"${TENSORRT_PATH}/lib/nvinfer_plugin.lib"
67+
)
68+
69+
# Link libraries to the target
70+
target_link_libraries(${PROJECT_NAME} PRIVATE
71+
${OpenCV_LIBS}
72+
${CUDA_LIBRARIES}
73+
${TENSORRT_LIBS}
74+
)
75+
76+
# Enable separable compilation for CUDA (optional but recommended)
77+
set_target_properties(${PROJECT_NAME} PROPERTIES
78+
CUDA_SEPARABLE_COMPILATION ON
79+
)
80+
81+
# (Optional) Specify CUDA architectures based on your GPU hardware
82+
# set(CMAKE_CUDA_ARCHITECTURES 75) # Example for Turing architecture
83+
84+
# (Optional) Set output directories for binaries
85+
# set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)

‎asset/output.mp4‎

6.38 MB
Binary file not shown.

‎include/YOLOv11.h‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
/**
2+
* @file YOLOv11.h
3+
* @brief Header file for the YOLOv11 object detection model using TensorRT and OpenCV.
4+
*
5+
* This class encapsulates the preprocessing, inference, and postprocessing steps required to
6+
* perform object detection using a YOLOv11 model with TensorRT.
7+
*/
8+
9+
#pragma once
10+
11+
#include "NvInfer.h"
12+
#include <opencv2/opencv.hpp>
13+
14+
using namespace nvinfer1;
15+
using namespace std;
16+
using namespace cv;
17+
18+
/**
19+
* @struct Detection
20+
* @brief A structure representing a detected object.
21+
*
22+
* Contains the confidence score, class ID, and bounding box for a detected object.
23+
*/
24+
struct Detection
25+
{
26+
float conf; //!< Confidence score of the detection.
27+
int class_id; //!< Class ID of the detected object.
28+
Rect bbox; //!< Bounding box of the detected object.
29+
};
30+
31+
/**
32+
* @class YOLOv11
33+
* @brief A class for running YOLOv11 object detection using TensorRT and OpenCV.
34+
*
35+
* This class handles model initialization, inference, and postprocessing to detect objects
36+
* in images.
37+
*/
38+
class YOLOv11
39+
{
40+
public:
41+
42+
/**
43+
* @brief Constructor to initialize the YOLOv11 object.
44+
*
45+
* Loads the model and initializes TensorRT objects.
46+
*
47+
* @param model_path Path to the model engine or ONNX file.
48+
* @param logger Reference to a TensorRT logger for error reporting.
49+
*/
50+
YOLOv11(string model_path, nvinfer1::ILogger& logger);
51+
52+
/**
53+
* @brief Destructor to clean up resources.
54+
*
55+
* Frees the allocated memory and TensorRT resources.
56+
*/
57+
~YOLOv11();
58+
59+
/**
60+
* @brief Preprocess the input image.
61+
*
62+
* Prepares the image for inference by resizing and normalizing it.
63+
*
64+
* @param image The input image to be preprocessed.
65+
*/
66+
void preprocess(Mat& image);
67+
68+
/**
69+
* @brief Run inference on the preprocessed image.
70+
*
71+
* Executes the TensorRT engine for object detection.
72+
*/
73+
void infer();
74+
75+
/**
76+
* @brief Postprocess the output from the model.
77+
*
78+
* Filters and decodes the raw output from the TensorRT engine into detection results.
79+
*
80+
* @param output A vector to store the detected objects.
81+
*/
82+
void postprocess(vector<Detection>& output);
83+
84+
/**
85+
* @brief Draw the detected objects on the image.
86+
*
87+
* Overlays bounding boxes and class labels on the image for visualization.
88+
*
89+
* @param image The input image where the detections will be drawn.
90+
* @param output A vector of detections to be visualized.
91+
*/
92+
void draw(Mat& image, const vector<Detection>& output);
93+
94+
private:
95+
/**
96+
* @brief Initialize TensorRT components from the given engine file.
97+
*
98+
* @param engine_path Path to the serialized TensorRT engine file.
99+
* @param logger Reference to a TensorRT logger for error reporting.
100+
*/
101+
void init(std::string engine_path, nvinfer1::ILogger& logger);
102+
103+
float* gpu_buffers[2]; //!< The vector of device buffers needed for engine execution.
104+
float* cpu_output_buffer; //!< Pointer to the output buffer on the host.
105+
106+
cudaStream_t stream; //!< CUDA stream for asynchronous execution.
107+
IRuntime* runtime; //!< The TensorRT runtime used to deserialize the engine.
108+
ICudaEngine* engine; //!< The TensorRT engine used to run the network.
109+
IExecutionContext* context; //!< The context for executing inference using an ICudaEngine.
110+
111+
// Model parameters
112+
int input_w; //!< Width of the input image.
113+
int input_h; //!< Height of the input image.
114+
int num_detections; //!< Number of detections output by the model.
115+
int detection_attribute_size; //!< Size of each detection attribute.
116+
int num_classes = 80; //!< Number of object classes that can be detected.
117+
const int MAX_IMAGE_SIZE = 4096 * 4096; //!< Maximum allowed input image size.
118+
float conf_threshold = 0.3f; //!< Confidence threshold for filtering detections.
119+
float nms_threshold = 0.4f; //!< Non-Maximum Suppression (NMS) threshold for filtering overlapping boxes.
120+
121+
vector<Scalar> colors; //!< A vector of colors for drawing bounding boxes.
122+
123+
/**
124+
* @brief Build the TensorRT engine from the ONNX model.
125+
*
126+
* @param onnxPath Path to the ONNX file.
127+
* @param logger Reference to a TensorRT logger for error reporting.
128+
*/
129+
void build(std::string onnxPath, nvinfer1::ILogger& logger);
130+
131+
/**
132+
* @brief Save the TensorRT engine to a file.
133+
*
134+
* @param filename Path to save the serialized engine.
135+
* @return True if the engine was saved successfully, false otherwise.
136+
*/
137+
bool saveEngine(const std::string& filename);
138+
};

‎include/common.h‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
const std::vector<std::string> CLASS_NAMES = {
2+
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train",
3+
"truck", "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench",
4+
"bird", "cat", "dog", "horse", "sheep", "cow", "elephant",
5+
"bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie",
6+
"suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat",
7+
"baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
8+
"fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich",
9+
"orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake",
10+
"chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv",
11+
"laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven",
12+
"toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors",
13+
"teddy bear", "hair drier", "toothbrush" };
14+
15+
const std::vector<std::vector<unsigned int>> COLORS = {
16+
{0, 114, 189}, {217, 83, 25}, {237, 177, 32}, {126, 47, 142}, {119, 172, 48}, {77, 190, 238},
17+
{162, 20, 47}, {76, 76, 76}, {153, 153, 153}, {255, 0, 0}, {255, 128, 0}, {191, 191, 0},
18+
{0, 255, 0}, {0, 0, 255}, {170, 0, 255}, {85, 85, 0}, {85, 170, 0}, {85, 255, 0},
19+
{170, 85, 0}, {170, 170, 0}, {170, 255, 0}, {255, 85, 0}, {255, 170, 0}, {255, 255, 0},
20+
{0, 85, 128}, {0, 170, 128}, {0, 255, 128}, {85, 0, 128}, {85, 85, 128}, {85, 170, 128},
21+
{85, 255, 128}, {170, 0, 128}, {170, 85, 128}, {170, 170, 128}, {170, 255, 128}, {255, 0, 128},
22+
{255, 85, 128}, {255, 170, 128}, {255, 255, 128}, {0, 85, 255}, {0, 170, 255}, {0, 255, 255},
23+
{85, 0, 255}, {85, 85, 255}, {85, 170, 255}, {85, 255, 255}, {170, 0, 255}, {170, 85, 255},
24+
{170, 170, 255}, {170, 255, 255}, {255, 0, 255}, {255, 85, 255}, {255, 170, 255}, {85, 0, 0},
25+
{128, 0, 0}, {170, 0, 0}, {212, 0, 0}, {255, 0, 0}, {0, 43, 0}, {0, 85, 0},
26+
{0, 128, 0}, {0, 170, 0}, {0, 212, 0}, {0, 255, 0}, {0, 0, 43}, {0, 0, 85},
27+
{0, 0, 128}, {0, 0, 170}, {0, 0, 212}, {0, 0, 255}, {0, 0, 0}, {36, 36, 36},
28+
{73, 73, 73}, {109, 109, 109}, {146, 146, 146}, {182, 182, 182}, {219, 219, 219}, {0, 114, 189},
29+
{80, 183, 189}, {128, 128, 0} };

‎include/cuda_utils.h‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
#ifndef TRTX_CUDA_UTILS_H_
2+
#define TRTX_CUDA_UTILS_H_
3+
4+
#include <cuda_runtime_api.h>
5+
6+
#ifndef CUDA_CHECK
7+
#define CUDA_CHECK(callstr)\
8+
{\
9+
cudaError_t error_code = callstr;\
10+
if (error_code != cudaSuccess) {\
11+
std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__;\
12+
assert(0);\
13+
}\
14+
}
15+
#endif // CUDA_CHECK
16+
17+
#endif // TRTX_CUDA_UTILS_H_

0 commit comments

Comments
 (0)