diff --git a/yolov9/CMakeLists.txt b/yolov9/CMakeLists.txt new file mode 100644 index 0000000..8aca589 --- /dev/null +++ b/yolov9/CMakeLists.txt @@ -0,0 +1,56 @@ +cmake_minimum_required(VERSION 3.10) + +project(TRTCreater) + +add_definitions(-w) +add_definitions(-std=c++11) +add_definitions(-DAPI_EXPORTS) +set(CMAKE_CXX_STANDARD 11) +set(CMAKE_BUILD_TYPE Debug) +set(CMAKE_CUDA_ARCHITECTURES 75 86 89) + +MESSAGE(STATUS "operation system is ${CMAKE_SYSTEM}") +IF (CMAKE_SYSTEM_NAME MATCHES "Linux") + MESSAGE(STATUS "current platform: Linux ") + set(CUDA_COMPILER_PATH "/usr/local/cuda/bin/nvcc") + set(TENSORRT_PATH "/home/benol/Package/TensorRT-8.6.1.6") + include_directories(/usr/local/cuda/include) + link_directories(/usr/local/cuda/lib64) + link_directories(/usr/local/cuda/lib) +ELSEIF (CMAKE_SYSTEM_NAME MATCHES "Windows") + MESSAGE(STATUS "current platform: Windows") + set(CUDA_COMPILER_PATH "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v12.2/bin/nvcc.exe") + set(TENSORRT_PATH "D:\\Program Files\\TensorRT-8.6.1.6") + set(OpenCV_DIR "D:\\Program Files\\opencv\\build") + include_directories(${PROJECT_SOURCE_DIR}/windows) + find_package(CUDA REQUIRED) + include_directories(${CUDA_INCLUDE_DIRS}) + link_directories(${CUDA_LIBRARIES}) +ELSE (CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64") + MESSAGE(STATUS "other platform: ${CMAKE_SYSTEM_PROCESSOR}") + include_directories(/usr/local/cuda/targets/aarch64-linux/include) + link_directories(/usr/local/cuda/targets/aarch64-linux/lib) +ENDIF (CMAKE_SYSTEM_NAME MATCHES "Linux") +set(CMAKE_CUDA_COMPILER ${CUDA_COMPILER_PATH}) +enable_language(CUDA) + +# tensorrt +include_directories(${TENSORRT_PATH}/include) +link_directories(${TENSORRT_PATH}/lib) + +find_package(OpenCV) +include_directories(${OpenCV_INCLUDE_DIRS}) + +include_directories(${PROJECT_SOURCE_DIR}/include/) +include_directories(${PROJECT_SOURCE_DIR}/plugin/) + +file(GLOB_RECURSE SRCS ${PROJECT_SOURCE_DIR}/src/*.cpp ${PROJECT_SOURCE_DIR}/src/*.cu) +file(GLOB_RECURSE PLUGIN_SRCS ${PROJECT_SOURCE_DIR}/plugin/*.cu) + +# add_library(myplugins SHARED ${PLUGIN_SRCS}) +add_library(myplugins SHARED ${PLUGIN_SRCS}) +target_link_libraries(myplugins nvinfer cudart) + +add_executable(yolov9 demo.cpp ${SRCS}) +target_link_libraries(yolov9 nvinfer cudart myplugins ${OpenCV_LIBS}) + diff --git a/yolov9/README.md b/yolov9/README.md new file mode 100644 index 0000000..2928c2a --- /dev/null +++ b/yolov9/README.md @@ -0,0 +1,106 @@ +# yolov9 + +The Pytorch implementation is [WongKinYiu/yolov9](https://github.com/WongKinYiu/yolov9). + +## Contributors + + +## Progress +- [x] YOLOv9-c: + - [x] FP32 + - [x] FP16 + - [x] INT8 +- [x] YOLOv9-e: + - [x] FP32 + - [x] FP16 + - [x] INT8 + +## Requirements + +- TensorRT 8.0+ +- OpenCV 3.4.0+ + +## Speed Test + +The speed test is done on a desktop with R7-5700G CPU and RTX 4060Ti GPU. The input size is 640x640. The FP32, FP16 and INT8 models are tested. The time only includes the inference time, not includes the pre-processing and post-processing. The time is the average of 1000 times inference. + +| frame | Model | FP32 | FP16 | INT8 | +| --- | --- | --- | --- | --- | +| pytorch | YOLOv9-c | - | 15.5ms | - | +| pytorch | YOLOv9-e | - | 19.7ms | - | +| tensorrt | YOLOv9-c | 13.5ms | 4.6ms | 3.0ms | +| tensorrt | YOLOv9-e | 8.3ms | 3.2ms | 2.15ms | + +YOLOv9-e is faster than YOLOv9-c in tensorrt, because the YOLOv9-e requires fewer layers of inference. +``` +YOLOv9-c: +[[31, 34, 37, 16, 19, 22], 1, DualDDetect, [nc]] # [A3, A4, A5, P3, P4, P5] + +YOLOv9-e: +[[35, 32, 29, 42, 45, 48], 1, DualDDetect, [nc]] + +``` + +In DualDDetect, the A3, A4, A5, P3, P4, P5 are the output of the backbone. The first 3 layers are used for the inference of the final result. + +The YOLOv9-c requires 37 layers of inference, but YOLOv9-e requires 35 layers of inference. + +## How to Run, yolov9 as example + +1. generate .wts from pytorch with .pt, or download .wts from model zoo + +``` +// download https://github.com/WongKinYiu/yolov9 +cp {tensorrtx}/yolov9/gen_wts.py {yolov9}/yolov9 +cd {yolov9}/yolov9 +python gen_wts.py +// a file 'yolov9.wts' will be generated. +``` +2. build tensorrtx/yolov9 and run + + +``` +cd {tensorrtx}/yolov9/ +// update kNumClass in config.h if your model is trained on custom dataset +mkdir build +cd build +cp {ultralytics}/ultralytics/yolov9.wts {tensorrtx}/yolov9/build +cmake .. +make +sudo ./yolov9 -s [.wts] [.engine] [c/e] // serialize model to plan file +sudo ./yolov9 -d [.engine] [image folder] // deserialize and run inference, the images in [image folder] will be processed. +// For example yolov9 +sudo ./yolov9 -s yolov9-c.wts yolov9-c.engine c +sudo ./yolov9 -d yolov9-c.engine ../images +``` + +3. check the images generated, as follows. _zidane.jpg and _bus.jpg + +4. optional, load and run the tensorrt model in python + +``` +// install python-tensorrt, pycuda, etc. +// ensure the yolov9.engine and libmyplugins.so have been built +python yolov9_trt.py +``` + + +# INT8 Quantization + +1. Prepare calibration images, you can randomly select 1000s images from your train set. For coco, you can also download my calibration images `coco_calib` from [GoogleDrive](https://drive.google.com/drive/folders/1s7jE9DtOngZMzJC1uL307J2MiaGwdRSI?usp=sharing) or [BaiduPan](https://pan.baidu.com/s/1GOm_-JobpyLMAqZWCDUhKg) pwd: a9wh + +2. unzip it in yolov8/build + +3. set the macro `USE_INT8` in config.h and change the path of calibration images in config.h, such as 'gCalibTablePath="./coco_calib/";' + +4. serialize the model and test + +

+ +

+ +## More Information + +See the readme in [home page.](https://github.com/wang-xinyu/tensorrtx) + + diff --git a/yolov9/demo.cpp b/yolov9/demo.cpp new file mode 100644 index 0000000..c67cc60 --- /dev/null +++ b/yolov9/demo.cpp @@ -0,0 +1,211 @@ +#include "config.h" +#include "model.h" +#include "cuda_utils.h" +#include "logging.h" +#include "utils.h" +#include "preprocess.h" +#include "postprocess.h" +#include +#include + +using namespace nvinfer1; + +const static int kOutputSize = kMaxNumOutputBbox * sizeof(Detection) / sizeof(float) + 1; +static Logger gLogger; + +void serialize_engine(unsigned int maxBatchSize, std::string& wts_name, std::string& sub_type, std::string& engine_name) { + // Create builder + IBuilder* builder = createInferBuilder(gLogger); + IBuilderConfig* config = builder->createBuilderConfig(); + + // Create model to populate the network, then set the outputs and create an engine + IHostMemory* serialized_engine = nullptr; + if (sub_type == "e") { + serialized_engine = build_engine_yolov9_e(maxBatchSize, builder, config, DataType::kFLOAT, wts_name); + } else if(sub_type == "c"){ + serialized_engine = build_engine_yolov9_c(maxBatchSize, builder, config, DataType::kFLOAT, wts_name); + } + else { + return; + } + assert(serialized_engine != nullptr); + + std::ofstream p(engine_name, std::ios::binary); + if (!p) { + std::cerr << "could not open plan output file" << std::endl; + assert(false); + } + p.write(reinterpret_cast(serialized_engine->data()), serialized_engine->size()); + + delete config; + delete serialized_engine; + delete builder; +} + +void deserialize_engine(std::string& engine_name, IRuntime** runtime, ICudaEngine** engine, IExecutionContext** context) { + std::ifstream file(engine_name, std::ios::binary); + if (!file.good()) { + std::cerr << "read " << engine_name << " error!" << std::endl; + assert(false); + } + size_t size = 0; + file.seekg(0, file.end); + size = file.tellg(); + file.seekg(0, file.beg); + char* serialized_engine = new char[size]; + assert(serialized_engine); + file.read(serialized_engine, size); + file.close(); + + *runtime = createInferRuntime(gLogger); + assert(*runtime); + *engine = (*runtime)->deserializeCudaEngine(serialized_engine, size); + assert(*engine); + *context = (*engine)->createExecutionContext(); + assert(*context); + delete[] serialized_engine; +} + +void prepare_buffer(ICudaEngine* engine, float** input_buffer_device, float** output_buffer_device, float** output_buffer_host) { + assert(engine->getNbBindings() == 2); + // In order to bind the buffers, we need to know the names of the input and output tensors. + // Note that indices are guaranteed to be less than IEngine::getNbBindings() + const int inputIndex = engine->getBindingIndex(kInputTensorName); + const int outputIndex = engine->getBindingIndex(kOutputTensorName); + assert(inputIndex == 0); + assert(outputIndex == 1); + // Create GPU buffers on device + CUDA_CHECK(cudaMalloc((void**)input_buffer_device, kBatchSize * 3 * kInputH * kInputW * sizeof(float))); + CUDA_CHECK(cudaMalloc((void**)output_buffer_device, kBatchSize * kOutputSize * sizeof(float))); + + *output_buffer_host = new float[kBatchSize * kOutputSize]; +} + +void infer(IExecutionContext& context, cudaStream_t& stream, void** buffers, float* output, int batchSize) { + // infer on the batch asynchronously, and DMA output back to host + context.enqueue(batchSize, buffers, stream, nullptr); + CUDA_CHECK(cudaMemcpyAsync(output, buffers[1], batchSize * kOutputSize * sizeof(float), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); +} + +bool parse_args(int argc, char** argv, std::string& wts, std::string& engine, std::string& img_dir, std::string& sub_type) { + if (argc < 4) return false; + if (std::string(argv[1]) == "-s" && argc == 5) { + wts = std::string(argv[2]); + engine = std::string(argv[3]); + sub_type = std::string(argv[4]); + } else if (std::string(argv[1]) == "-d" && argc == 4) { + engine = std::string(argv[2]); + img_dir = std::string(argv[3]); + } else { + return false; + } + return true; +} + +int main(int argc, char** argv) { + cudaSetDevice(kGpuId); + + std::string wts_name = ""; + std::string engine_name = ""; + std::string img_dir; + std::string sub_type = ""; + // speed test or inference + // const int speed_test_iter = 1000; + const int speed_test_iter = 1; + + if (!parse_args(argc, argv, wts_name, engine_name, img_dir, sub_type)) { + std::cerr << "Arguments not right!" << std::endl; + std::cerr << "./yolov9 -s [.wts] [.engine] [c/e] // serialize model to plan file" << std::endl; + std::cerr << "./yolov9 -d [.engine] ../samples // deserialize plan file and run inference" << std::endl; + return -1; + } + + + // Create a model using the API directly and serialize it to a file + if (!wts_name.empty()) { + serialize_engine(kBatchSize, wts_name, sub_type, engine_name); + return 0; + } + + // Deserialize the engine from file + IRuntime* runtime = nullptr; + ICudaEngine* engine = nullptr; + IExecutionContext* context = nullptr; + deserialize_engine(engine_name, &runtime, &engine, &context); + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreate(&stream)); + + cuda_preprocess_init(kMaxInputImageSize); + + // Prepare cpu and gpu buffers + float* device_buffers[2]; + float* output_buffer_host = nullptr; + prepare_buffer(engine, &device_buffers[0], &device_buffers[1], &output_buffer_host); + + // Read images from directory + std::vector file_names; + if (read_files_in_dir(img_dir.c_str(), file_names) < 0) { + std::cerr << "read_files_in_dir failed." << std::endl; + return -1; + } + + // batch predict + for (size_t i = 0; i < file_names.size(); i += kBatchSize) { + // Get a batch of images + std::vector img_batch; + std::vector img_name_batch; + for (size_t j = i; j < i + kBatchSize && j < file_names.size(); j++) { + cv::Mat img = cv::imread(img_dir + "/" + file_names[j]); + img_batch.push_back(img); + img_name_batch.push_back(file_names[j]); + } + + // Preprocess + cuda_batch_preprocess(img_batch, device_buffers[0], kInputW, kInputH, stream); + + // Run inference + auto start = std::chrono::system_clock::now(); + for (int j = 0; j < speed_test_iter; j++) { + infer(*context, stream, (void**)device_buffers, output_buffer_host, kBatchSize); + } + // infer(*context, stream, (void**)device_buffers, output_buffer_host, kBatchSize); + auto end = std::chrono::system_clock::now(); + std::cout << "inference time: " << std::chrono::duration_cast(end - start).count() / 1000.0 / speed_test_iter << "ms" << std::endl; + + // NMS + std::vector> res_batch; + batch_nms(res_batch, output_buffer_host, img_batch.size(), kOutputSize, kConfThresh, kNmsThresh); + + // Draw bounding boxes + draw_bbox(img_batch, res_batch); + + // Save images + for (size_t j = 0; j < img_batch.size(); j++) { + cv::imwrite("_" + img_name_batch[j], img_batch[j]); + } + } + + // Release stream and buffers + cudaStreamDestroy(stream); + CUDA_CHECK(cudaFree(device_buffers[0])); + CUDA_CHECK(cudaFree(device_buffers[1])); + delete[] output_buffer_host; + cuda_preprocess_destroy(); + // Destroy the engine + delete context; + delete engine; + delete runtime; + + // Print histogram of the output distribution + //std::cout << "\nOutput:\n\n"; + //for (unsigned int i = 0; i < kOutputSize; i++) + //{ + // std::cout << prob[i] << ", "; + // if (i % 10 == 0) std::cout << std::endl; + //} + //std::cout << std::endl; + + return 0; +} + diff --git a/yolov9/gen_wts.py b/yolov9/gen_wts.py new file mode 100644 index 0000000..a507e1f --- /dev/null +++ b/yolov9/gen_wts.py @@ -0,0 +1,70 @@ +import sys +import argparse +import os +import struct +import torch +from utils.torch_utils import select_device + +def parse_args(): + parser = argparse.ArgumentParser(description='Convert .pt file to .wts') + parser.add_argument('-w', '--weights', default='yolov9-e.pt', + help='Input weights (.pt) file path (required)') + parser.add_argument( + '-o', '--output', help='Output (.wts) file path (optional)') + parser.add_argument( + '-t', '--type', type=str, default='detect', choices=['detect', 'cls', 'seg'], + help='determines the model is detection/classification') + args = parser.parse_args() + if not os.path.isfile(args.weights): + raise SystemExit('Invalid input file') + if not args.output: + args.output = os.path.splitext(args.weights)[0] + '.wts' + elif os.path.isdir(args.output): + args.output = os.path.join( + args.output, + os.path.splitext(os.path.basename(args.weights))[0] + '.wts') + return args.weights, args.output, args.type + +pt_file, wts_file, m_type = parse_args() +print(f'Generating .wts for {m_type} model') + +# Load model +print(f'Loading {pt_file}') +device = select_device('cpu') +model = torch.load(pt_file, map_location=device) # Load FP32 weights +model.model.float() + +if m_type in ['detect', 'seg']: + # update anchor_grid info + anchor_grid = model.model[-1].anchors * model.model[-1].stride[..., None, None] + # model.model[-1].anchor_grid = anchor_grid + # delattr(model.model[-1], 'anchor_grid') # model.model[-1] is detect layer + # The parameters are saved in the OrderDict through the "register_buffer" method, and then saved to the weight. + model.model[-1].register_buffer("anchor_grid", anchor_grid) + # model.model[-1].register_buffer("strides", model.model[-1].stride) + +model.to(device).eval() + +# print(model.model) +# 将model.model保存到txt中 +with open('model.txt', 'w') as f: + f.write(str(model.model)) +f.close() +print(f'Writing into {wts_file}') +with open(wts_file, 'w') as f: + f.write('{}\n'.format(len(model.state_dict().keys()))) + for k, v in model.state_dict().items(): + vr = v.reshape(-1).cpu().numpy() + f.write('{} {} '.format(k, len(vr))) + for vv in vr: + f.write(' ') + f.write(struct.pack('>f', float(vv)).hex()) + f.write('\n') +wts_file_key = wts_file.replace('.wts', '_key.txt') +print(f'Writing into {wts_file_key}') +with open(wts_file_key, 'w') as f: + f.write('{}\n'.format(len(model.state_dict().keys()))) + for k, v in model.state_dict().items(): + vr = v.reshape(-1).cpu().numpy() + f.write('{} {} '.format(k, len(vr))) + f.write('\n') diff --git a/yolov9/images b/yolov9/images new file mode 120000 index 0000000..02cc755 --- /dev/null +++ b/yolov9/images @@ -0,0 +1 @@ +../yolov3-spp/samples \ No newline at end of file diff --git a/yolov9/include/block.h b/yolov9/include/block.h new file mode 100644 index 0000000..9369fa9 --- /dev/null +++ b/yolov9/include/block.h @@ -0,0 +1,31 @@ +#include "config.h" +#include "yololayer.h" + +#include +#include +#include +#include +#include +#include + +using namespace nvinfer1; + +// TensorRT weight files have a simple space delimited format: +// [type] [size] +void PrintDim(const ILayer* layer, std::string log = ""); +std::map loadWeights(const std::string file); +int get_width(int x, float gw, int divisor = 8); +int get_depth(int x, float gd) ; +ILayer* Proto(INetworkDefinition* network, std::map& weightMap, ITensor& input, int c_, int c2, std::string lname); +std::vector> getAnchors(std::map& weightMap, std::string lname); +// ---------------------------------------------------------------- +nvinfer1::ILayer* convBnSiLU(nvinfer1::INetworkDefinition* network, std::map& weightMap, nvinfer1::ITensor& input, int ch, int k, int s, int p, std::string lname, int g=1); +ILayer* RepNCSPELAN4(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int c3, int c4, int c5, std::string lname); +ILayer* ADown(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c2, std::string lname); +std::vector CBLinear(INetworkDefinition *network, std::map& weightMap, ITensor& input, std::vector c2s, int k, int s, int p, int g, std::string lname); +ILayer* CBFuse(INetworkDefinition *network, std::vector> input, std::vector idx, std::vector strides); +ILayer* SPPELAN(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int c3, std::string lname); +std::vector DualDDetect(INetworkDefinition *network, std::map& weightMap, std::vector dets, int cls, std::vector ch, std::string lname); +nvinfer1::IPluginV2Layer* addYoLoLayer(nvinfer1::INetworkDefinition *network, std::vector dets, bool is_segmentation); +nvinfer1::IShuffleLayer* DFL(nvinfer1::INetworkDefinition* network, std::map weightMap, nvinfer1::ITensor& input, int ch, int k, int s, int p, std::string lname); +nvinfer1::ILayer* convBnNoAct(nvinfer1::INetworkDefinition* network, std::map& weightMap, nvinfer1::ITensor& input, int ch, int k, int s, int p, std::string lname, int g); diff --git a/yolov9/include/calibrator.h b/yolov9/include/calibrator.h new file mode 100644 index 0000000..7f6acb2 --- /dev/null +++ b/yolov9/include/calibrator.h @@ -0,0 +1,36 @@ +#pragma once + +#include "macros.h" +#include +#include + +//! \class Int8EntropyCalibrator2 +//! +//! \brief Implements Entropy calibrator 2. +//! CalibrationAlgoType is kENTROPY_CALIBRATION_2. +//! +class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2 { +public: + Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache = true); + + virtual ~Int8EntropyCalibrator2(); + int getBatchSize() const TRT_NOEXCEPT override; + bool getBatch(void* bindings[], const char* names[], int nbBindings) TRT_NOEXCEPT override; + const void* readCalibrationCache(size_t& length) TRT_NOEXCEPT override; + void writeCalibrationCache(const void* cache, size_t length) TRT_NOEXCEPT override; + +private: + int batchsize_; + int input_w_; + int input_h_; + int img_idx_; + std::string img_dir_; + std::vector img_files_; + size_t input_count_; + std::string calib_table_name_; + const char* input_blob_name_; + bool read_cache_; + void* device_input_; + std::vector calib_cache_; +}; + diff --git a/yolov9/include/config.h b/yolov9/include/config.h new file mode 100644 index 0000000..5863fef --- /dev/null +++ b/yolov9/include/config.h @@ -0,0 +1,59 @@ +#pragma once + +/* -------------------------------------------------------- + * These configs are related to tensorrt model, if these are changed, + * please re-compile and re-serialize the tensorrt model. + * --------------------------------------------------------*/ + +// For INT8, you need prepare the calibration dataset, please refer to +// https://github.com/wang-xinyu/tensorrtx/tree/master/yolov5#int8-quantization +#define USE_INT8 // set USE_INT8 or USE_FP16 or USE_FP32 + +#ifdef USE_INT8 +const static char* gCalibTablePath = "./calib"; +#endif + +// These are used to define input/output tensor names, +// you can set them to whatever you want. +const static char* kInputTensorName = "images"; +const static char* kOutputTensorName = "output"; + +// Detection model and Segmentation model' number of classes +constexpr static int kNumClass = 80; + +// Classfication model's number of classes +constexpr static int kClsNumClass = 1000; + +constexpr static int kBatchSize = 1; + +// Yolo's input width and height must by divisible by 32 +constexpr static int kInputH = 640; +constexpr static int kInputW = 640; + +// Classfication model's input shape +constexpr static int kClsInputH = 224; +constexpr static int kClsInputW = 224; + +// Maximum number of output bounding boxes from yololayer plugin. +// That is maximum number of output bounding boxes before NMS. +constexpr static int kMaxNumOutputBbox = 2000; + +constexpr static int kNumAnchor = 3; + +// The bboxes whose confidence is lower than kIgnoreThresh will be ignored in yololayer plugin. +constexpr static float kIgnoreThresh = 0.05f; + +/* -------------------------------------------------------- + * These configs are NOT related to tensorrt model, if these are changed, + * please re-compile, but no need to re-serialize the tensorrt model. + * --------------------------------------------------------*/ + +// NMS overlapping thresh and final detection confidence thresh +const static float kNmsThresh = 0.45f; +const static float kConfThresh = 0.1f; + +const static int kGpuId = 0; + +// If your image size is larger than 4096 * 3112, please increase this value +const static int kMaxInputImageSize = 4096 * 3112; + diff --git a/yolov9/include/cuda_utils.h b/yolov9/include/cuda_utils.h new file mode 100644 index 0000000..8fbd319 --- /dev/null +++ b/yolov9/include/cuda_utils.h @@ -0,0 +1,18 @@ +#ifndef TRTX_CUDA_UTILS_H_ +#define TRTX_CUDA_UTILS_H_ + +#include + +#ifndef CUDA_CHECK +#define CUDA_CHECK(callstr)\ + {\ + cudaError_t error_code = callstr;\ + if (error_code != cudaSuccess) {\ + std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__;\ + assert(0);\ + }\ + } +#endif // CUDA_CHECK + +#endif // TRTX_CUDA_UTILS_H_ + diff --git a/yolov9/include/logging.h b/yolov9/include/logging.h new file mode 100644 index 0000000..6b79a8b --- /dev/null +++ b/yolov9/include/logging.h @@ -0,0 +1,504 @@ +/* + * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef TENSORRT_LOGGING_H +#define TENSORRT_LOGGING_H + +#include "NvInferRuntimeCommon.h" +#include +#include +#include +#include +#include +#include +#include +#include "macros.h" + +using Severity = nvinfer1::ILogger::Severity; + +class LogStreamConsumerBuffer : public std::stringbuf +{ +public: + LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog) + : mOutput(stream) + , mPrefix(prefix) + , mShouldLog(shouldLog) + { + } + + LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other) + : mOutput(other.mOutput) + { + } + + ~LogStreamConsumerBuffer() + { + // std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence + // std::streambuf::pptr() gives a pointer to the current position of the output sequence + // if the pointer to the beginning is not equal to the pointer to the current position, + // call putOutput() to log the output to the stream + if (pbase() != pptr()) + { + putOutput(); + } + } + + // synchronizes the stream buffer and returns 0 on success + // synchronizing the stream buffer consists of inserting the buffer contents into the stream, + // resetting the buffer and flushing the stream + virtual int sync() + { + putOutput(); + return 0; + } + + void putOutput() + { + if (mShouldLog) + { + // prepend timestamp + std::time_t timestamp = std::time(nullptr); + tm* tm_local = std::localtime(×tamp); + std::cout << "["; + std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/"; + std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/"; + std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-"; + std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":"; + std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":"; + std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] "; + // std::stringbuf::str() gets the string contents of the buffer + // insert the buffer contents pre-appended by the appropriate prefix into the stream + mOutput << mPrefix << str(); + // set the buffer to empty + str(""); + // flush the stream + mOutput.flush(); + } + } + + void setShouldLog(bool shouldLog) + { + mShouldLog = shouldLog; + } + +private: + std::ostream& mOutput; + std::string mPrefix; + bool mShouldLog; +}; + +//! +//! \class LogStreamConsumerBase +//! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer +//! +class LogStreamConsumerBase +{ +public: + LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog) + : mBuffer(stream, prefix, shouldLog) + { + } + +protected: + LogStreamConsumerBuffer mBuffer; +}; + +//! +//! \class LogStreamConsumer +//! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages. +//! Order of base classes is LogStreamConsumerBase and then std::ostream. +//! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field +//! in LogStreamConsumer and then the address of the buffer is passed to std::ostream. +//! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream. +//! Please do not change the order of the parent classes. +//! +class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream +{ +public: + //! \brief Creates a LogStreamConsumer which logs messages with level severity. + //! Reportable severity determines if the messages are severe enough to be logged. + LogStreamConsumer(Severity reportableSeverity, Severity severity) + : LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity) + , std::ostream(&mBuffer) // links the stream buffer with the stream + , mShouldLog(severity <= reportableSeverity) + , mSeverity(severity) + { + } + + LogStreamConsumer(LogStreamConsumer&& other) + : LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog) + , std::ostream(&mBuffer) // links the stream buffer with the stream + , mShouldLog(other.mShouldLog) + , mSeverity(other.mSeverity) + { + } + + void setReportableSeverity(Severity reportableSeverity) + { + mShouldLog = mSeverity <= reportableSeverity; + mBuffer.setShouldLog(mShouldLog); + } + +private: + static std::ostream& severityOstream(Severity severity) + { + return severity >= Severity::kINFO ? std::cout : std::cerr; + } + + static std::string severityPrefix(Severity severity) + { + switch (severity) + { + case Severity::kINTERNAL_ERROR: return "[F] "; + case Severity::kERROR: return "[E] "; + case Severity::kWARNING: return "[W] "; + case Severity::kINFO: return "[I] "; + case Severity::kVERBOSE: return "[V] "; + default: assert(0); return ""; + } + } + + bool mShouldLog; + Severity mSeverity; +}; + +//! \class Logger +//! +//! \brief Class which manages logging of TensorRT tools and samples +//! +//! \details This class provides a common interface for TensorRT tools and samples to log information to the console, +//! and supports logging two types of messages: +//! +//! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal) +//! - Test pass/fail messages +//! +//! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is +//! that the logic for controlling the verbosity and formatting of sample output is centralized in one location. +//! +//! In the future, this class could be extended to support dumping test results to a file in some standard format +//! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run). +//! +//! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger +//! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT +//! library and messages coming from the sample. +//! +//! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the +//! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger +//! object. + +class Logger : public nvinfer1::ILogger +{ +public: + Logger(Severity severity = Severity::kWARNING) + : mReportableSeverity(severity) + { + } + + //! + //! \enum TestResult + //! \brief Represents the state of a given test + //! + enum class TestResult + { + kRUNNING, //!< The test is running + kPASSED, //!< The test passed + kFAILED, //!< The test failed + kWAIVED //!< The test was waived + }; + + //! + //! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger + //! \return The nvinfer1::ILogger associated with this Logger + //! + //! TODO Once all samples are updated to use this method to register the logger with TensorRT, + //! we can eliminate the inheritance of Logger from ILogger + //! + nvinfer1::ILogger& getTRTLogger() + { + return *this; + } + + //! + //! \brief Implementation of the nvinfer1::ILogger::log() virtual method + //! + //! Note samples should not be calling this function directly; it will eventually go away once we eliminate the + //! inheritance from nvinfer1::ILogger + //! + void log(Severity severity, const char* msg) TRT_NOEXCEPT override + { + LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl; + } + + //! + //! \brief Method for controlling the verbosity of logging output + //! + //! \param severity The logger will only emit messages that have severity of this level or higher. + //! + void setReportableSeverity(Severity severity) + { + mReportableSeverity = severity; + } + + //! + //! \brief Opaque handle that holds logging information for a particular test + //! + //! This object is an opaque handle to information used by the Logger to print test results. + //! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used + //! with Logger::reportTest{Start,End}(). + //! + class TestAtom + { + public: + TestAtom(TestAtom&&) = default; + + private: + friend class Logger; + + TestAtom(bool started, const std::string& name, const std::string& cmdline) + : mStarted(started) + , mName(name) + , mCmdline(cmdline) + { + } + + bool mStarted; + std::string mName; + std::string mCmdline; + }; + + //! + //! \brief Define a test for logging + //! + //! \param[in] name The name of the test. This should be a string starting with + //! "TensorRT" and containing dot-separated strings containing + //! the characters [A-Za-z0-9_]. + //! For example, "TensorRT.sample_googlenet" + //! \param[in] cmdline The command line used to reproduce the test + // + //! \return a TestAtom that can be used in Logger::reportTest{Start,End}(). + //! + static TestAtom defineTest(const std::string& name, const std::string& cmdline) + { + return TestAtom(false, name, cmdline); + } + + //! + //! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments + //! as input + //! + //! \param[in] name The name of the test + //! \param[in] argc The number of command-line arguments + //! \param[in] argv The array of command-line arguments (given as C strings) + //! + //! \return a TestAtom that can be used in Logger::reportTest{Start,End}(). + static TestAtom defineTest(const std::string& name, int argc, char const* const* argv) + { + auto cmdline = genCmdlineString(argc, argv); + return defineTest(name, cmdline); + } + + //! + //! \brief Report that a test has started. + //! + //! \pre reportTestStart() has not been called yet for the given testAtom + //! + //! \param[in] testAtom The handle to the test that has started + //! + static void reportTestStart(TestAtom& testAtom) + { + reportTestResult(testAtom, TestResult::kRUNNING); + assert(!testAtom.mStarted); + testAtom.mStarted = true; + } + + //! + //! \brief Report that a test has ended. + //! + //! \pre reportTestStart() has been called for the given testAtom + //! + //! \param[in] testAtom The handle to the test that has ended + //! \param[in] result The result of the test. Should be one of TestResult::kPASSED, + //! TestResult::kFAILED, TestResult::kWAIVED + //! + static void reportTestEnd(const TestAtom& testAtom, TestResult result) + { + assert(result != TestResult::kRUNNING); + assert(testAtom.mStarted); + reportTestResult(testAtom, result); + } + + static int reportPass(const TestAtom& testAtom) + { + reportTestEnd(testAtom, TestResult::kPASSED); + return EXIT_SUCCESS; + } + + static int reportFail(const TestAtom& testAtom) + { + reportTestEnd(testAtom, TestResult::kFAILED); + return EXIT_FAILURE; + } + + static int reportWaive(const TestAtom& testAtom) + { + reportTestEnd(testAtom, TestResult::kWAIVED); + return EXIT_SUCCESS; + } + + static int reportTest(const TestAtom& testAtom, bool pass) + { + return pass ? reportPass(testAtom) : reportFail(testAtom); + } + + Severity getReportableSeverity() const + { + return mReportableSeverity; + } + +private: + //! + //! \brief returns an appropriate string for prefixing a log message with the given severity + //! + static const char* severityPrefix(Severity severity) + { + switch (severity) + { + case Severity::kINTERNAL_ERROR: return "[F] "; + case Severity::kERROR: return "[E] "; + case Severity::kWARNING: return "[W] "; + case Severity::kINFO: return "[I] "; + case Severity::kVERBOSE: return "[V] "; + default: assert(0); return ""; + } + } + + //! + //! \brief returns an appropriate string for prefixing a test result message with the given result + //! + static const char* testResultString(TestResult result) + { + switch (result) + { + case TestResult::kRUNNING: return "RUNNING"; + case TestResult::kPASSED: return "PASSED"; + case TestResult::kFAILED: return "FAILED"; + case TestResult::kWAIVED: return "WAIVED"; + default: assert(0); return ""; + } + } + + //! + //! \brief returns an appropriate output stream (cout or cerr) to use with the given severity + //! + static std::ostream& severityOstream(Severity severity) + { + return severity >= Severity::kINFO ? std::cout : std::cerr; + } + + //! + //! \brief method that implements logging test results + //! + static void reportTestResult(const TestAtom& testAtom, TestResult result) + { + severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # " + << testAtom.mCmdline << std::endl; + } + + //! + //! \brief generate a command line string from the given (argc, argv) values + //! + static std::string genCmdlineString(int argc, char const* const* argv) + { + std::stringstream ss; + for (int i = 0; i < argc; i++) + { + if (i > 0) + ss << " "; + ss << argv[i]; + } + return ss.str(); + } + + Severity mReportableSeverity; +}; + +namespace +{ + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE +//! +//! Example usage: +//! +//! LOG_VERBOSE(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) +{ + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE); +} + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO +//! +//! Example usage: +//! +//! LOG_INFO(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_INFO(const Logger& logger) +{ + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO); +} + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING +//! +//! Example usage: +//! +//! LOG_WARN(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_WARN(const Logger& logger) +{ + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING); +} + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR +//! +//! Example usage: +//! +//! LOG_ERROR(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_ERROR(const Logger& logger) +{ + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR); +} + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR +// ("fatal" severity) +//! +//! Example usage: +//! +//! LOG_FATAL(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_FATAL(const Logger& logger) +{ + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR); +} + +} // anonymous namespace + +#endif // TENSORRT_LOGGING_H diff --git a/yolov9/include/macros.h b/yolov9/include/macros.h new file mode 100644 index 0000000..17339a2 --- /dev/null +++ b/yolov9/include/macros.h @@ -0,0 +1,29 @@ +#ifndef __MACROS_H +#define __MACROS_H + +#include + +#ifdef API_EXPORTS +#if defined(_MSC_VER) +#define API __declspec(dllexport) +#else +#define API __attribute__((visibility("default"))) +#endif +#else + +#if defined(_MSC_VER) +#define API __declspec(dllimport) +#else +#define API +#endif +#endif // API_EXPORTS + +#if NV_TENSORRT_MAJOR >= 8 +#define TRT_NOEXCEPT noexcept +#define TRT_CONST_ENQUEUE const +#else +#define TRT_NOEXCEPT +#define TRT_CONST_ENQUEUE +#endif + +#endif // __MACROS_H diff --git a/yolov9/include/model.h b/yolov9/include/model.h new file mode 100644 index 0000000..eeb0e7c --- /dev/null +++ b/yolov9/include/model.h @@ -0,0 +1,6 @@ +#pragma once + +#include +#include +nvinfer1::IHostMemory* build_engine_yolov9_e(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, std::string& wts_name); +nvinfer1::IHostMemory* build_engine_yolov9_c(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, std::string& wts_name); \ No newline at end of file diff --git a/yolov9/include/postprocess.h b/yolov9/include/postprocess.h new file mode 100644 index 0000000..3a0cbe4 --- /dev/null +++ b/yolov9/include/postprocess.h @@ -0,0 +1,20 @@ +#pragma once + +#include "types.h" +#include +#include +cv::Rect get_rect(cv::Mat& img, float bbox[4]); + +void nms(std::vector& res, float *output, float conf_thresh, float nms_thresh = 0.5); + +void batch_nms(std::vector>& batch_res, float *output, int batch_size, int output_size, float conf_thresh, float nms_thresh = 0.5); + +void draw_bbox(std::vector& img_batch, std::vector>& res_batch); + +std::vector process_mask(const float* proto, int proto_size, std::vector& dets); + +void draw_mask_bbox(cv::Mat& img, std::vector& dets, std::vector& masks, std::unordered_map& labels_map); +// cuda NMS +void cuda_decode(float* predict, int num_bboxes, float confidence_threshold,float* parray,int max_objects, cudaStream_t stream); +void cuda_nms(float* parray, float nms_threshold, int max_objects, cudaStream_t stream); +void batch_process(std::vector> &res_batch, const float* decode_ptr_host, int batch_size, int bbox_element, const std::vector& img_batch); \ No newline at end of file diff --git a/yolov9/include/preprocess.h b/yolov9/include/preprocess.h new file mode 100644 index 0000000..c0dc1aa --- /dev/null +++ b/yolov9/include/preprocess.h @@ -0,0 +1,15 @@ +#pragma once + +#include +#include +#include + +void cuda_preprocess_init(int max_image_size); +void cuda_preprocess_destroy(); +void cuda_preprocess(uint8_t* src, int src_width, int src_height, + float* dst, int dst_width, int dst_height, + cudaStream_t stream); +void cuda_batch_preprocess(std::vector& img_batch, + float* dst, int dst_width, int dst_height, + cudaStream_t stream); + diff --git a/yolov9/include/types.h b/yolov9/include/types.h new file mode 100644 index 0000000..0505849 --- /dev/null +++ b/yolov9/include/types.h @@ -0,0 +1,17 @@ +#pragma once + +#include "config.h" + +struct YoloKernel { + int width; + int height; + float anchors[kNumAnchor * 2]; +}; + +struct alignas(float) Detection { + float bbox[4]; // center_x center_y w h + float conf; // bbox_conf * cls_conf + float class_id; + float mask[32]; +}; +const int bbox_element = 7; // center_x, center_y, w, h, conf, cls, obj diff --git a/yolov9/include/utils.h b/yolov9/include/utils.h new file mode 100644 index 0000000..2dea946 --- /dev/null +++ b/yolov9/include/utils.h @@ -0,0 +1,70 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +static inline int read_files_in_dir(const char* p_dir_name, std::vector& file_names) { + DIR *p_dir = opendir(p_dir_name); + if (p_dir == nullptr) { + return -1; + } + + struct dirent* p_file = nullptr; + while ((p_file = readdir(p_dir)) != nullptr) { + if (strcmp(p_file->d_name, ".") != 0 && + strcmp(p_file->d_name, "..") != 0) { + //std::string cur_file_name(p_dir_name); + //cur_file_name += "/"; + //cur_file_name += p_file->d_name; + std::string cur_file_name(p_file->d_name); + file_names.push_back(cur_file_name); + } + } + + closedir(p_dir); + return 0; +} + +// Function to trim leading and trailing whitespace from a string +static inline std::string trim_leading_whitespace(const std::string& str) { + size_t first = str.find_first_not_of(' '); + if (std::string::npos == first) { + return str; + } + size_t last = str.find_last_not_of(' '); + return str.substr(first, (last - first + 1)); +} + +// Src: https://stackoverflow.com/questions/16605967 +static inline std::string to_string_with_precision(const float a_value, const int n = 2) { + std::ostringstream out; + out.precision(n); + out << std::fixed << a_value; + return out.str(); +} + +static inline int read_labels(const std::string labels_filename, std::unordered_map& labels_map) { + + std::ifstream file(labels_filename); + // Read each line of the file + std::string line; + int index = 0; + while (std::getline(file, line)) { + // Strip the line of any leading or trailing whitespace + line = trim_leading_whitespace(line); + + // Add the stripped line to the labels_map, using the loop index as the key + labels_map[index] = line; + index++; + } + // Close the file + file.close(); + + return 0; +} + diff --git a/yolov9/plugin/yololayer.cu b/yolov9/plugin/yololayer.cu new file mode 100644 index 0000000..bdc073c --- /dev/null +++ b/yolov9/plugin/yololayer.cu @@ -0,0 +1,242 @@ +#include "yololayer.h" +#include "types.h" +#include +#include +#include "cuda_utils.h" +#include +#include + +namespace Tn { + template + void write(char*& buffer, const T& val) { + *reinterpret_cast(buffer) = val; + buffer += sizeof(T); + } + + template + void read(const char*& buffer, T& val) { + val = *reinterpret_cast(buffer); + buffer += sizeof(T); + } +} // namespace Tn + + +namespace nvinfer1 { +YoloLayerPlugin::YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, bool is_segmentation) { + mClassCount = classCount; + mYoloV8NetWidth = netWidth; + mYoloV8netHeight = netHeight; + mMaxOutObject = maxOut; + is_segmentation_ = is_segmentation; +} + +YoloLayerPlugin::~YoloLayerPlugin() {} + +YoloLayerPlugin::YoloLayerPlugin(const void* data, size_t length) { + using namespace Tn; + const char* d = reinterpret_cast(data), * a = d; + read(d, mClassCount); + read(d, mThreadCount); + read(d, mYoloV8NetWidth); + read(d, mYoloV8netHeight); + read(d, mMaxOutObject); + read(d, is_segmentation_); + + assert(d == a + length); +} + +void YoloLayerPlugin::serialize(void* buffer) const TRT_NOEXCEPT { + + using namespace Tn; + char* d = static_cast(buffer), * a = d; + write(d, mClassCount); + write(d, mThreadCount); + write(d, mYoloV8NetWidth); + write(d, mYoloV8netHeight); + write(d, mMaxOutObject); + write(d, is_segmentation_); + + assert(d == a + getSerializationSize()); +} + +size_t YoloLayerPlugin::getSerializationSize() const TRT_NOEXCEPT { + return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mYoloV8netHeight) + sizeof(mYoloV8NetWidth) + sizeof(mMaxOutObject) + sizeof(is_segmentation_); +} + +int YoloLayerPlugin::initialize() TRT_NOEXCEPT { + return 0; +} + +nvinfer1::Dims YoloLayerPlugin::getOutputDimensions(int index, const nvinfer1::Dims* inputs, int nbInputDims) TRT_NOEXCEPT { + int total_size = mMaxOutObject * sizeof(Detection) / sizeof(float); + return nvinfer1::Dims3(total_size + 1, 1, 1); +} + +void YoloLayerPlugin::setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT { + mPluginNamespace = pluginNamespace; +} + +const char* YoloLayerPlugin::getPluginNamespace() const TRT_NOEXCEPT { + return mPluginNamespace; +} + +nvinfer1::DataType YoloLayerPlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT { + return nvinfer1::DataType::kFLOAT; +} + +bool YoloLayerPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT { + + return false; +} + +bool YoloLayerPlugin::canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT { + + return false; +} + +void YoloLayerPlugin::configurePlugin(nvinfer1::PluginTensorDesc const* in, int nbInput, nvinfer1::PluginTensorDesc const* out, int nbOutput) TRT_NOEXCEPT {}; + +void YoloLayerPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT {}; + +void YoloLayerPlugin::detachFromContext() TRT_NOEXCEPT {} + +const char* YoloLayerPlugin::getPluginType() const TRT_NOEXCEPT { + + return "YoloLayer_TRT"; +} + +const char* YoloLayerPlugin::getPluginVersion() const TRT_NOEXCEPT { + return "1"; +} + +void YoloLayerPlugin::destroy() TRT_NOEXCEPT { + + delete this; +} + +nvinfer1::IPluginV2IOExt* YoloLayerPlugin::clone() const TRT_NOEXCEPT { + + YoloLayerPlugin* p = new YoloLayerPlugin(mClassCount, mYoloV8NetWidth, mYoloV8netHeight, mMaxOutObject, is_segmentation_); + p->setPluginNamespace(mPluginNamespace); + return p; +} + +int YoloLayerPlugin::enqueue(int batchSize, const void* TRT_CONST_ENQUEUE* inputs, void* const* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT { + + forwardGpu((const float* const*)inputs, (float*)outputs[0], stream, mYoloV8netHeight, mYoloV8NetWidth, batchSize); + return 0; +} + + +__device__ float Logist(float data) { return 1.0f / (1.0f + expf(-data)); }; + +__global__ void CalDetection(const float* input, float* output, int numElements, int maxoutobject, + const int grid_h, int grid_w, const int stride, int classes, int outputElem, bool is_segmentation) { + int idx = threadIdx.x + blockDim.x * blockIdx.x; + if (idx >= numElements) return; + + int total_grid = grid_h * grid_w; + int info_len = 4 + classes; + if (is_segmentation) info_len += 32; + int batchIdx = idx / total_grid; + int elemIdx = idx % total_grid; + const float* curInput = input + batchIdx * total_grid * info_len; + int outputIdx = batchIdx * outputElem; + + int class_id = 0; + float max_cls_prob = 0.0; + for (int i = 4; i < 4 + classes; i++) { + float p = Logist(curInput[elemIdx + i * total_grid]); + if (p > max_cls_prob) { + max_cls_prob = p; + class_id = i - 4; + } + } + + if (max_cls_prob < 0.1) return; + + int count = (int)atomicAdd(output + outputIdx, 1); + if (count >= maxoutobject) return; + char* data = (char*)(output + outputIdx) + sizeof(float) + count * sizeof(Detection); + Detection* det = (Detection*)(data); + + int row = elemIdx / grid_w; + int col = elemIdx % grid_w; + + det->conf = max_cls_prob; + det->class_id = class_id; + det->bbox[0] = (col + 0.5f - curInput[elemIdx + 0 * total_grid]) * stride; + det->bbox[1] = (row + 0.5f - curInput[elemIdx + 1 * total_grid]) * stride; + det->bbox[2] = (col + 0.5f + curInput[elemIdx + 2 * total_grid]) * stride; + det->bbox[3] = (row + 0.5f + curInput[elemIdx + 3 * total_grid]) * stride; + + for (int k = 0; is_segmentation && k < 32; k++) { + det->mask[k] = curInput[elemIdx + (k + 4 + classes) * total_grid]; + } +} + +void YoloLayerPlugin::forwardGpu(const float* const* inputs, float* output, cudaStream_t stream, int mYoloV8netHeight,int mYoloV8NetWidth, int batchSize) { + int outputElem = 1 + mMaxOutObject * sizeof(Detection) / sizeof(float); + cudaMemsetAsync(output, 0, sizeof(float), stream); + for (int idx = 0; idx < batchSize; ++idx) { + CUDA_CHECK(cudaMemsetAsync(output + idx * outputElem, 0, sizeof(float), stream)); + } + int numElem = 0; + int grids[3][2] = { {mYoloV8netHeight / 8, mYoloV8NetWidth / 8}, {mYoloV8netHeight / 16, mYoloV8NetWidth / 16}, {mYoloV8netHeight / 32, mYoloV8NetWidth / 32} }; + int strides[] = { 8, 16, 32 }; + for (unsigned int i = 0; i < 3; i++) { + int grid_h = grids[i][0]; + int grid_w = grids[i][1]; + int stride = strides[i]; + numElem = grid_h * grid_w * batchSize; + if (numElem < mThreadCount) mThreadCount = numElem; + + CalDetection << <(numElem + mThreadCount - 1) / mThreadCount, mThreadCount, 0, stream >> > + (inputs[i], output, numElem, mMaxOutObject, grid_h, grid_w, stride, mClassCount, outputElem, is_segmentation_); + } +} + +PluginFieldCollection YoloPluginCreator::mFC{}; +std::vector YoloPluginCreator::mPluginAttributes; + +YoloPluginCreator::YoloPluginCreator() { + mPluginAttributes.clear(); + mFC.nbFields = mPluginAttributes.size(); + mFC.fields = mPluginAttributes.data(); +} + +const char* YoloPluginCreator::getPluginName() const TRT_NOEXCEPT { + return "YoloLayer_TRT"; +} + +const char* YoloPluginCreator::getPluginVersion() const TRT_NOEXCEPT { + return "1"; +} + +const PluginFieldCollection* YoloPluginCreator::getFieldNames() TRT_NOEXCEPT { + return &mFC; +} + +IPluginV2IOExt* YoloPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) TRT_NOEXCEPT { + assert(fc->nbFields == 1); + assert(strcmp(fc->fields[0].name, "netinfo") == 0); + int* p_netinfo = (int*)(fc->fields[0].data); + int class_count = p_netinfo[0]; + int input_w = p_netinfo[1]; + int input_h = p_netinfo[2]; + int max_output_object_count = p_netinfo[3]; + bool is_segmentation = p_netinfo[4]; + YoloLayerPlugin* obj = new YoloLayerPlugin(class_count, input_w, input_h, max_output_object_count, is_segmentation); + obj->setPluginNamespace(mNamespace.c_str()); + return obj; +} + +IPluginV2IOExt* YoloPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRT_NOEXCEPT { + // This object will be deleted when the network is destroyed, which will + // call YoloLayerPlugin::destroy() + YoloLayerPlugin* obj = new YoloLayerPlugin(serialData, serialLength); + obj->setPluginNamespace(mNamespace.c_str()); + return obj; +} + +} // namespace nvinfer1 diff --git a/yolov9/plugin/yololayer.h b/yolov9/plugin/yololayer.h new file mode 100644 index 0000000..514c1f1 --- /dev/null +++ b/yolov9/plugin/yololayer.h @@ -0,0 +1,102 @@ +#pragma once +#include "macros.h" +#include "NvInfer.h" +#include +#include +#include "macros.h" +namespace nvinfer1 { +class API YoloLayerPlugin : public IPluginV2IOExt { +public: + YoloLayerPlugin(int classCount, int netWdith, int netHeight, int maxOut, bool is_segmentation); + YoloLayerPlugin(const void* data, size_t length); + ~YoloLayerPlugin(); + + int getNbOutputs() const TRT_NOEXCEPT override { + return 1; + } + + nvinfer1::Dims getOutputDimensions(int index, const nvinfer1::Dims* inputs, int nbInputDims) TRT_NOEXCEPT override; + + int initialize() TRT_NOEXCEPT override; + + virtual void terminate() TRT_NOEXCEPT override {} + + virtual size_t getWorkspaceSize(int maxBatchSize) const TRT_NOEXCEPT override { return 0; } + + virtual int enqueue(int batchSize, const void* const* inputs, void* TRT_CONST_ENQUEUE* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT override; + + virtual size_t getSerializationSize() const TRT_NOEXCEPT override; + + virtual void serialize(void* buffer) const TRT_NOEXCEPT override; + + bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const TRT_NOEXCEPT override { + return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT; + } + + + const char* getPluginType() const TRT_NOEXCEPT override; + + const char* getPluginVersion() const TRT_NOEXCEPT override; + + void destroy() TRT_NOEXCEPT override; + + IPluginV2IOExt* clone() const TRT_NOEXCEPT override; + + void setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT override; + + const char* getPluginNamespace() const TRT_NOEXCEPT override; + + nvinfer1::DataType getOutputDataType(int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const TRT_NOEXCEPT; + + bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT override; + + bool canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT override; + + void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT override; + + void configurePlugin(PluginTensorDesc const* in, int32_t nbInput, PluginTensorDesc const* out, int32_t nbOutput) TRT_NOEXCEPT override; + + void detachFromContext() TRT_NOEXCEPT override; + + private: + void forwardGpu(const float* const* inputs, float* output, cudaStream_t stream, int mYoloV8netHeight, int mYoloV8NetWidth, int batchSize); + int mThreadCount = 256; + const char* mPluginNamespace; + int mClassCount; + int mYoloV8NetWidth; + int mYoloV8netHeight; + int mMaxOutObject; + bool is_segmentation_; + }; + +class API YoloPluginCreator : public IPluginCreator { +public: + YoloPluginCreator(); + ~YoloPluginCreator() override = default; + + const char* getPluginName() const TRT_NOEXCEPT override; + + const char* getPluginVersion() const TRT_NOEXCEPT override; + + const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override; + + nvinfer1::IPluginV2IOExt* createPlugin(const char* name, const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT override; + + nvinfer1::IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRT_NOEXCEPT override; + + void setPluginNamespace(const char* libNamespace) TRT_NOEXCEPT override { + mNamespace = libNamespace; + } + + const char* getPluginNamespace() const TRT_NOEXCEPT override { + return mNamespace.c_str(); + } + + private: + std::string mNamespace; + static PluginFieldCollection mFC; + static std::vector mPluginAttributes; + }; + REGISTER_TENSORRT_PLUGIN(YoloPluginCreator); +} // namespace nvinfer1 + diff --git a/yolov9/src/block.cpp b/yolov9/src/block.cpp new file mode 100644 index 0000000..b86ae6d --- /dev/null +++ b/yolov9/src/block.cpp @@ -0,0 +1,385 @@ +#include "block.h" +#include "calibrator.h" +#include "config.h" +#include "yololayer.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace nvinfer1; + +// TensorRT weight files have a simple space delimited format: +// [type] [size] +void PrintDim(const ILayer* layer, std::string log) { + Dims dim = layer->getOutput(0)->getDimensions(); + std::cout << log <<": "<<"\t\t\t\t"; + for (int i = 0; i < dim.nbDims; i++) { + std::cout << dim.d[i] << " "; + } + std::cout << std::endl; +} +std::map loadWeights(const std::string file) { + std::cout << "Loading weights: " << file << std::endl; + std::map weightMap; + + // Open weights file + std::ifstream input(file); + assert(input.is_open() && "Unable to load weight file. please check if the .wts file path is right!!!!!!"); + + // Read number of weight blobs + int32_t count; + input >> count; + assert(count > 0 && "Invalid weight map file."); + + while (count--) { + Weights wt{ DataType::kFLOAT, nullptr, 0 }; + uint32_t size; + + // Read name and type of blob + std::string name; + input >> name >> std::dec >> size; + wt.type = DataType::kFLOAT; + + // Load blob + uint32_t* val = reinterpret_cast(malloc(sizeof(val) * size)); + for (uint32_t x = 0, y = size; x < y; ++x) { + input >> std::hex >> val[x]; + } + wt.values = val; + + wt.count = size; + weightMap[name] = wt; + } + + return weightMap; +} + +int get_width(int x, float gw, int divisor) { + return int(ceil((x * gw) / divisor)) * divisor; +} + +int get_depth(int x, float gd) { + if (x == 1) return 1; + int r = round(x * gd); + if (x * gd - int(x * gd) == 0.5 && (int(x * gd) % 2) == 0) { + --r; + } + return std::max(r, 1); +} +static nvinfer1::IScaleLayer* addBatchNorm2d(nvinfer1::INetworkDefinition* network, std::map weightMap, +nvinfer1::ITensor& input, std::string lname, float eps) { + float* gamma = (float*)weightMap[lname + ".weight"].values; + float* beta = (float*)weightMap[lname + ".bias"].values; + float* mean = (float*)weightMap[lname + ".running_mean"].values; + float* var = (float*)weightMap[lname + ".running_var"].values; + int len = weightMap[lname + ".running_var"].count; + + float* scval = reinterpret_cast(malloc(sizeof(float) * len)); + for(int i = 0; i < len; i++) { + scval[i] = gamma[i] / sqrt(var[i] + eps); + } + nvinfer1::Weights scale{nvinfer1::DataType::kFLOAT, scval, len}; + + float* shval = reinterpret_cast(malloc(sizeof(float) * len)); + for(int i = 0; i < len; i++) { + shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps); + } + nvinfer1::Weights shift{nvinfer1::DataType::kFLOAT, shval, len}; + + float* pval = reinterpret_cast(malloc(sizeof(float) * len)); + for (int i = 0; i < len; i++) { + pval[i] = 1.0; + } + nvinfer1::Weights power{ nvinfer1::DataType::kFLOAT, pval, len }; + weightMap[lname + ".scale"] = scale; + weightMap[lname + ".shift"] = shift; + weightMap[lname + ".power"] = power; + nvinfer1::IScaleLayer* output = network->addScale(input, nvinfer1::ScaleMode::kCHANNEL, shift, scale, power); + assert(output); + return output; +} +nvinfer1::ILayer* convBnSiLU(nvinfer1::INetworkDefinition* network, std::map& weightMap, nvinfer1::ITensor& input, int ch, int k, int s, int p, std::string lname, int g) { + nvinfer1::Weights bias_empty{nvinfer1::DataType::kFLOAT, nullptr, 0}; + nvinfer1::IConvolutionLayer* conv = network->addConvolutionNd(input, ch, nvinfer1::DimsHW{k, k}, weightMap[lname + ".conv.weight"], bias_empty); + assert(conv); + conv->setStrideNd(nvinfer1::DimsHW{s, s}); + conv->setPaddingNd(nvinfer1::DimsHW{p, p}); + conv->setNbGroups(g); + + nvinfer1::IScaleLayer* bn = addBatchNorm2d(network, weightMap, *conv->getOutput(0), lname + ".bn", 1e-3); + + nvinfer1::IActivationLayer* sigmoid = network->addActivation(*bn->getOutput(0), nvinfer1::ActivationType::kSIGMOID); + assert(sigmoid); + auto ew = network->addElementWise(*bn->getOutput(0), *sigmoid->getOutput(0), nvinfer1::ElementWiseOperation::kPROD); + assert(ew); + return ew; +} +nvinfer1::ILayer* convBnNoAct(nvinfer1::INetworkDefinition* network, std::map& weightMap, nvinfer1::ITensor& input, int ch, int k, int s, int p, std::string lname, int g) { + nvinfer1::Weights bias_empty{nvinfer1::DataType::kFLOAT, nullptr, 0}; + nvinfer1::IConvolutionLayer* conv = network->addConvolutionNd(input, ch, nvinfer1::DimsHW{k, k}, weightMap[lname + ".conv.weight"], bias_empty); + assert(conv); + conv->setStrideNd(nvinfer1::DimsHW{s, s}); + conv->setPaddingNd(nvinfer1::DimsHW{p, p}); + conv->setNbGroups(g); + + nvinfer1::IScaleLayer* bn = addBatchNorm2d(network, weightMap, *conv->getOutput(0), lname + ".bn", 1e-3); + return bn; +} + +std::vector> getAnchors(std::map& weightMap, std::string lname) { + std::vector> anchors; + Weights wts = weightMap[lname + ".anchor_grid"]; + int anchor_len = kNumAnchor * 2; + for (int i = 0; i < wts.count / anchor_len; i++) { + auto *p = (const float*)wts.values + i * anchor_len; + std::vector anchor(p, p + anchor_len); + anchors.push_back(anchor); + } + return anchors; +} + +ILayer* RepConvN(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int k, int s, int p, int g, int d, bool act, bool bn, bool deploy, std::string lname) { + assert(k == 3 && p == 1); + ILayer* conv1 = convBnNoAct(network, weightMap, input, c2, k, s, p, lname + ".conv1", g); + ILayer* conv2 = convBnNoAct(network, weightMap, input, c2, 1, s, p - k / 2, lname + ".conv2", g); + ILayer* ew0 = network->addElementWise(*conv1->getOutput(0), *conv2->getOutput(0), ElementWiseOperation::kSUM); + nvinfer1::IActivationLayer* sigmoid = network->addActivation(*ew0->getOutput(0), nvinfer1::ActivationType::kSIGMOID); + assert(sigmoid); + + auto ew = network->addElementWise(*ew0->getOutput(0), *sigmoid->getOutput(0), nvinfer1::ElementWiseOperation::kPROD); + assert(ew); + return ew; +} + +ILayer* RepNBottleneck(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, bool shortcut, int k, int g, float e, std::string lname) { + int c_ = int(c2 * e); + assert(k == 3 && "RepVGG only support kernel size 3"); + auto cv1 = RepConvN(network, weightMap, input, c1, c_, k, 1, 1, g, 1, true, false, false, lname + ".cv1"); + auto cv2 = convBnSiLU(network, weightMap, *cv1->getOutput(0), c2, k, 1, 1, lname + ".cv2", g); + if (shortcut && c1 == c2) { + auto ew = network->addElementWise(input, *cv2->getOutput(0), ElementWiseOperation::kSUM); + return ew; + } + return cv2; +} + +ILayer* RepNCSP(INetworkDefinition *network, std::map& weightMap, ITensor& input, + int c1, int c2, int n, bool shortcut, int g, float e, std::string lname) { + int c_ = int(c2 * e); + + auto cv1 = convBnSiLU(network, weightMap, input, c_, 1, 1, 0, lname + ".cv1", 1); + + ILayer* m = cv1; + for(int i = 0; i < n; i++) { + m = RepNBottleneck(network, weightMap, *m->getOutput(0), c_, c_, shortcut, 3, g, 1.0, lname + ".m." + std::to_string(i)); + } + + // auto m_0 = RepNBottleneck(network, weightMap, *cv1->getOutput(0), c_, c_, shortcut, 3, g, 1.0, lname + ".m.0"); + // auto m_1 = RepNBottleneck(network, weightMap, *m_0->getOutput(0), c_, c_, shortcut, 3, g, 1.0, lname + ".m.1"); + + auto cv2 = convBnSiLU(network, weightMap, input, c_, 1, 1, 0, lname + ".cv2", 1); + ITensor* inputTensors[] = { m->getOutput(0), cv2->getOutput(0) }; + auto cat = network->addConcatenation(inputTensors, 2); + + auto cv3 = convBnSiLU(network, weightMap, *cat->getOutput(0), c2, 1, 1, 0, lname + ".cv3", 1); + return cv3; +} +ILayer* RepNCSPELAN4(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int c3, int c4, int c5, std::string lname) { + + auto cv1 = convBnSiLU(network, weightMap, input, c3, 1, 1, 0, lname + ".cv1", 1); + // 将cv1的输出分成两部分 chunk(2, 1) + + nvinfer1::Dims d = cv1->getOutput(0)->getDimensions(); + nvinfer1::ISliceLayer* split1 = network->addSlice(*cv1->getOutput(0), nvinfer1::Dims3{0,0,0}, nvinfer1::Dims3{d.d[0]/2, d.d[1], d.d[2]}, nvinfer1::Dims3{1,1,1}); + nvinfer1::ISliceLayer* split2 = network->addSlice(*cv1->getOutput(0), nvinfer1::Dims3{d.d[0]/2,0,0}, nvinfer1::Dims3{d.d[0]/2, d.d[1], d.d[2]}, nvinfer1::Dims3{1,1,1}); + + auto cv2_0 = RepNCSP(network, weightMap, *split2->getOutput(0), c3/2, c4, c5, true, 1, 0.5, lname + ".cv2.0"); + auto cv2_1 = convBnSiLU(network, weightMap, *cv2_0->getOutput(0), c4, 3, 1, 1, lname + ".cv2.1", 1); + + auto cv3_0 = RepNCSP(network, weightMap, *cv2_1->getOutput(0), c4, c4, c5, true, 1, 0.5, lname + ".cv3.0"); + auto cv3_1 = convBnSiLU(network, weightMap, *cv3_0->getOutput(0), c4, 3, 1, 1, lname + ".cv3.1", 1); + + ITensor* inputTensors[] = { split1->getOutput(0), split2->getOutput(0), cv2_1->getOutput(0), cv3_1->getOutput(0) }; + auto cat = network->addConcatenation(inputTensors, 4); + auto cv4 = convBnSiLU(network, weightMap, *cat->getOutput(0), c2, 1, 1, 0, lname + ".cv4", 1); + return cv4; +} + +ILayer* ADown(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c2, std::string lname) { + int c_ = c2 / 2; + auto pool = network->addPoolingNd(input, PoolingType::kAVERAGE, DimsHW{ 2, 2 }); + pool->setStrideNd(DimsHW{ 1, 1 }); + pool->setPaddingNd(DimsHW{ 0, 0 }); + + nvinfer1::Dims d = pool->getOutput(0)->getDimensions(); + nvinfer1::ISliceLayer* split1 = network->addSlice(*pool->getOutput(0), nvinfer1::Dims3{ 0, 0, 0}, nvinfer1::Dims3{ d.d[0]/2, d.d[1], d.d[2] }, nvinfer1::Dims3{ 1, 1, 1 }); + nvinfer1::ISliceLayer* split2 = network->addSlice(*pool->getOutput(0), nvinfer1::Dims3{ d.d[0]/2, 0, 0}, nvinfer1::Dims3{ d.d[0]/2, d.d[1], d.d[2] }, nvinfer1::Dims3{ 1, 1, 1 }); + + // auto chunklayer = layer_split(1, pool->getOutput(0), network); + auto cv1 = convBnSiLU(network, weightMap, *split1->getOutput(0), c_, 3, 2, 1, lname + ".cv1", 1); + + auto pool2 = network->addPoolingNd(*split2->getOutput(0), PoolingType::kMAX, DimsHW{ 3, 3 }); + pool2->setStrideNd(DimsHW{ 2, 2 }); + pool2->setPaddingNd(DimsHW{ 1, 1 }); + auto cv2 = convBnSiLU(network, weightMap, *pool2->getOutput(0), c_, 1, 1, 0, lname + ".cv2", 1); + + ITensor* inputTensors[] = { cv1->getOutput(0), cv2->getOutput(0) }; + auto cat = network->addConcatenation(inputTensors, 2); + return cat; +} + +std::vector CBLinear(INetworkDefinition *network, std::map& weightMap, ITensor& input, std::vector c2s, int k, int s, int p, int g, std::string lname) { + + IConvolutionLayer* conv1 = network->addConvolutionNd(input, std::accumulate(c2s.begin(), c2s.end(), 0), DimsHW{ k, k }, weightMap[lname + ".conv.weight"], weightMap[lname + ".conv.bias"]); + assert(conv1); + conv1->setName((lname + ".conv").c_str()); + conv1->setStrideNd(DimsHW{ s, s }); + conv1->setPaddingNd(DimsHW{ p, p }); + + int h = input.getDimensions().d[1]; + int w = input.getDimensions().d[2]; + std::vector slices(c2s.size()); + int start = 0; + for(int i = 0; i < c2s.size(); i++) { + slices[i] = network->addSlice(*conv1->getOutput(0), Dims3{ start, 0, 0 }, Dims3{ c2s[i], h, w }, Dims3{ 1, 1, 1 }); + start += c2s[i]; + } + return slices; +} + +ILayer* CBFuse(INetworkDefinition *network, std::vector> input, std::vector idx, std::vector strides) { + ILayer** res = new ILayer*[input.size()]; + res[input.size()-1] = input[input.size()-1][0]; + + for(int i = input.size()-2; i >= 0; i--) { + auto upsample = network->addResize(*input[i][idx[i]]->getOutput(0)); + upsample->setResizeMode(ResizeMode::kNEAREST); + const float scales[] = { 1, strides[i] / strides[strides.size() - 1], strides[i] / strides[strides.size() - 1] }; + upsample->setScales(scales, 3); + res[i] = upsample; + } + + for(int i = 1; i < input.size(); i++) { + auto ew = network->addElementWise(*res[0]->getOutput(0), *res[i]->getOutput(0), ElementWiseOperation::kSUM); + res[0] = ew; + } + return res[0]; +} + +ILayer* SP(INetworkDefinition *network, std::map& weightMap, ITensor& input, int k, int s) { + int p = k / 2; + auto pool = network->addPoolingNd(input, PoolingType::kMAX, DimsHW{ k, k }); + pool->setPaddingNd(DimsHW{ p, p }); + pool->setStrideNd(DimsHW{ s, s }); + return pool; +} + +ILayer* SPPELAN(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int c3, std::string lname) { + auto cv1 = convBnSiLU(network, weightMap, input, c3, 1, 1, 0, lname + ".cv1", 1); + auto cv2 = SP(network, weightMap, *cv1->getOutput(0), 5, 1); + auto cv3 = SP(network, weightMap, *cv2->getOutput(0), 5, 1); + auto cv4 = SP(network, weightMap, *cv3->getOutput(0), 5, 1); + + ITensor* inputTensors[] = { cv1->getOutput(0), cv2->getOutput(0), cv3->getOutput(0), cv4->getOutput(0) }; + auto cat = network->addConcatenation(inputTensors, 4); + auto cv5 = convBnSiLU(network, weightMap, *cat->getOutput(0), c2, 1, 1, 0, lname + ".cv5", 1); + return cv5; +} + +ILayer* DetectBbox_Conv(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c2, int reg_max, std::string lname) { + auto cv_0 = convBnSiLU(network, weightMap, input, c2, 3, 1, 1, lname + ".0", 1); + auto cv_1 = convBnSiLU(network, weightMap, *cv_0->getOutput(0), c2, 3, 1, 1, lname + ".1", 4); + auto cv_2 = network->addConvolutionNd(*cv_1->getOutput(0), reg_max*4, DimsHW{ 1, 1 }, weightMap[lname + ".2.weight"], weightMap[lname + ".2.bias"]); + cv_2->setName((lname + ".conv").c_str()); + cv_2->setStrideNd(DimsHW{ 1, 1 }); + cv_2->setPaddingNd(DimsHW{ 0, 0 }); + cv_2->setNbGroups(4); + return cv_2; +} + +ILayer* DetectCls_Conv(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c2, int cls, std::string lname) { + auto cv_0 = convBnSiLU(network, weightMap, input, c2, 3, 1, 1, lname + ".0", 1); + auto cv_1 = convBnSiLU(network, weightMap, *cv_0->getOutput(0), c2, 3, 1, 1, lname + ".1", 1); + auto cv_2 = network->addConvolutionNd(*cv_1->getOutput(0), cls, DimsHW{ 1, 1 }, weightMap[lname + ".2.weight"], weightMap[lname + ".2.bias"]); + cv_2->setName((lname + ".conv").c_str()); + cv_2->setStrideNd(DimsHW{ 1, 1 }); + cv_2->setPaddingNd(DimsHW{ 0, 0 }); + return cv_2; +} + +nvinfer1::IShuffleLayer* DFL(nvinfer1::INetworkDefinition* network, std::map weightMap, nvinfer1::ITensor& input, int ch, int k, int s, int p, std::string lname){ + auto dim = input.getDimensions(); + int c = dim.d[0]; + int grid = dim.d[1] * dim.d[2]; + int split_num = c / ch; + + nvinfer1::IShuffleLayer* shuffle1 = network->addShuffle(input); + shuffle1->setReshapeDimensions(nvinfer1::Dims3{ split_num, ch, grid }); + shuffle1->setSecondTranspose(nvinfer1::Permutation{ 1, 0, 2 }); + nvinfer1::ISoftMaxLayer* softmax = network->addSoftMax(*shuffle1->getOutput(0)); + nvinfer1::Weights bias_empty{ nvinfer1::DataType::kFLOAT, nullptr, 0 }; + nvinfer1::IConvolutionLayer* conv = network->addConvolutionNd(*softmax->getOutput(0), 1, nvinfer1::DimsHW{1, 1}, weightMap[lname + ".conv.weight"], bias_empty); + conv->setStrideNd(nvinfer1::DimsHW{ s, s }); + conv->setPaddingNd(nvinfer1::DimsHW{ p, p }); + nvinfer1::IShuffleLayer* shuffle2 = network->addShuffle(*conv->getOutput(0)); + shuffle2->setReshapeDimensions(nvinfer1::Dims2{ 4, grid }); + return shuffle2; +} + +nvinfer1::IPluginV2Layer* addYoLoLayer(nvinfer1::INetworkDefinition *network, std::vector dets, bool is_segmentation) { + auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1"); + + nvinfer1::PluginField plugin_fields[1]; + int netinfo[5] = { kNumClass, kInputW, kInputH, kMaxNumOutputBbox, is_segmentation }; + plugin_fields[0].data = netinfo; + plugin_fields[0].length = 5; + plugin_fields[0].name = "netinfo"; + plugin_fields[0].type = nvinfer1::PluginFieldType::kFLOAT32; + + nvinfer1::PluginFieldCollection plugin_data; + plugin_data.nbFields = 1; + plugin_data.fields = plugin_fields; + nvinfer1::IPluginV2 *plugin_obj = creator->createPlugin("yololayer", &plugin_data); + std::vector input_tensors; + for (auto det: dets) { + input_tensors.push_back(det->getOutput(0)); + } + auto yolo = network->addPluginV2(&input_tensors[0], input_tensors.size(), *plugin_obj); + return yolo; +} + +std::vector DualDDetect(INetworkDefinition *network, std::map& weightMap, std::vector dets, int cls, std::vector ch, std::string lname) { + int c2 = std::max(int(ch[0] / 4), int(16 * 4)); + int c3 = std::max(ch[0], std::min(cls * 2, 128)); + int reg_max = 16; + + std::vector bboxlayers; + std::vector clslayers; + + for (int i = 0; i < dets.size(); i++) { + // Conv(x, c2, 3), Conv(c2, c2, 3, g=4), nn.Conv2d(c2, 4 * self.reg_max, 1, groups=4) + bboxlayers.push_back(DetectBbox_Conv(network, weightMap, *dets[i]->getOutput(0), c2, reg_max, lname + ".cv2." + std::to_string(i))); + // Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, self.nc, 1) + auto cls_layer = DetectCls_Conv(network, weightMap, *dets[i]->getOutput(0), c3, cls, lname + ".cv3." + std::to_string(i)); + auto dim = cls_layer->getOutput(0)->getDimensions(); + nvinfer1::IShuffleLayer* shuffle = network->addShuffle(*cls_layer->getOutput(0)); + shuffle->setReshapeDimensions(nvinfer1::Dims2{ kNumClass, dim.d[1] * dim.d[2] }); + clslayers.push_back(shuffle); + } + + std::vector ret; + for (int i = 0; i < dets.size(); i++) { + // softmax 16*4, w, h => 16, 4, w, h + auto loc = DFL(network, weightMap, *bboxlayers[i]->getOutput(0), 16, 1, 1, 0, lname + ".dfl"); + nvinfer1::ITensor* inputTensor[] = { loc->getOutput(0), clslayers[i]->getOutput(0) }; + ret.push_back(network->addConcatenation(inputTensor, 2)); + } + return ret; +} diff --git a/yolov9/src/calibrator.cpp b/yolov9/src/calibrator.cpp new file mode 100644 index 0000000..7768d63 --- /dev/null +++ b/yolov9/src/calibrator.cpp @@ -0,0 +1,97 @@ +#include "calibrator.h" +#include "cuda_utils.h" +#include "utils.h" + +#include +#include +#include +#include +#include + +static cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h) { + int w, h, x, y; + float r_w = input_w / (img.cols * 1.0); + float r_h = input_h / (img.rows * 1.0); + if (r_h > r_w) { + w = input_w; + h = r_w * img.rows; + x = 0; + y = (input_h - h) / 2; + } else { + w = r_h * img.cols; + h = input_h; + x = (input_w - w) / 2; + y = 0; + } + cv::Mat re(h, w, CV_8UC3); + cv::resize(img, re, re.size(), 0, 0, cv::INTER_LINEAR); + cv::Mat out(input_h, input_w, CV_8UC3, cv::Scalar(128, 128, 128)); + re.copyTo(out(cv::Rect(x, y, re.cols, re.rows))); + return out; +} + +Int8EntropyCalibrator2::Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache) + : batchsize_(batchsize), + input_w_(input_w), + input_h_(input_h), + img_idx_(0), + img_dir_(img_dir), + calib_table_name_(calib_table_name), + input_blob_name_(input_blob_name), + read_cache_(read_cache) { + input_count_ = 3 * input_w * input_h * batchsize; + CUDA_CHECK(cudaMalloc(&device_input_, input_count_ * sizeof(float))); + read_files_in_dir(img_dir, img_files_); +} + +Int8EntropyCalibrator2::~Int8EntropyCalibrator2() { + CUDA_CHECK(cudaFree(device_input_)); +} + +int Int8EntropyCalibrator2::getBatchSize() const TRT_NOEXCEPT { + return batchsize_; +} + +bool Int8EntropyCalibrator2::getBatch(void* bindings[], const char* names[], int nbBindings) TRT_NOEXCEPT { + if (img_idx_ + batchsize_ > (int)img_files_.size()) { + return false; + } + + std::vector input_imgs_; + for (int i = img_idx_; i < img_idx_ + batchsize_; i++) { + std::cout << img_files_[i] << " " << i << std::endl; + cv::Mat temp = cv::imread(img_dir_ + img_files_[i]); + if (temp.empty()) { + std::cerr << "Fatal error: image cannot open!" << std::endl; + return false; + } + cv::Mat pr_img = preprocess_img(temp, input_w_, input_h_); + input_imgs_.push_back(pr_img); + } + img_idx_ += batchsize_; + cv::Mat blob = cv::dnn::blobFromImages(input_imgs_, 1.0 / 255.0, cv::Size(input_w_, input_h_), cv::Scalar(0, 0, 0), true, false); + + CUDA_CHECK(cudaMemcpy(device_input_, blob.ptr(0), input_count_ * sizeof(float), cudaMemcpyHostToDevice)); + assert(!strcmp(names[0], input_blob_name_)); + bindings[0] = device_input_; + return true; +} + +const void* Int8EntropyCalibrator2::readCalibrationCache(size_t& length) TRT_NOEXCEPT { + std::cout << "reading calib cache: " << calib_table_name_ << std::endl; + calib_cache_.clear(); + std::ifstream input(calib_table_name_, std::ios::binary); + input >> std::noskipws; + if (read_cache_ && input.good()) { + std::copy(std::istream_iterator(input), std::istream_iterator(), std::back_inserter(calib_cache_)); + } + length = calib_cache_.size(); + return length ? calib_cache_.data() : nullptr; +} + +void Int8EntropyCalibrator2::writeCalibrationCache(const void* cache, size_t length) TRT_NOEXCEPT { + std::cout << "writing calib cache: " << calib_table_name_ << " size: " << length << std::endl; + std::ofstream output(calib_table_name_, std::ios::binary); + output.write(reinterpret_cast(cache), length); +} + diff --git a/yolov9/src/model.cpp b/yolov9/src/model.cpp new file mode 100644 index 0000000..fa3f9e3 --- /dev/null +++ b/yolov9/src/model.cpp @@ -0,0 +1,413 @@ +#include "model.h" +#include "calibrator.h" +#include "config.h" +#include "yololayer.h" +#include "block.h" +#include +#include +#include +#include +#include +#include + +using namespace nvinfer1; + +#ifdef USE_INT8 +void Calibrator(IBuilder* builder, IBuilderConfig* config) { + std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; + assert(builder->platformHasFastInt8()); + config->setFlag(BuilderFlag::kINT8); + Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, kInputW, kInputH, gCalibTablePath, "int8calib.table", kInputTensorName); + config->setInt8Calibrator(calibrator); +} +#endif + +IHostMemory* build_engine_yolov9_e(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, std::string& wts_name) { + /* ------ Create the builder ------ */ + INetworkDefinition* network = builder->createNetworkV2(0U); + + ITensor* data = network->addInput(kInputTensorName, dt, Dims3{ 3, kInputH, kInputW }); + assert(data); + std::map weightMap = loadWeights(wts_name); + + /* ------backbone------ */ + // [-1, 1, Conv, [64, 3, 2]], # 1-P1/2 + auto conv_1 = convBnSiLU(network, weightMap, *data, 64, 3, 2, 1, "model.1", 1); + assert(conv_1); + // [-1, 1, Conv, [128, 3, 2]], # 2-P2/4 + auto conv_2 = convBnSiLU(network, weightMap, *conv_1->getOutput(0), 128, 3, 2, 1, "model.2"); + // csp-elan block + // [-1, 1, RepNCSPELAN4, [256, 128, 64, 2]], # 3 + auto repncspelan_3 = RepNCSPELAN4(network, weightMap, *conv_2->getOutput(0), 128, 256, 128, 64, 2, "model.3"); + // avg-conv down + // [-1, 1, ADown, [256]], # 4-P3/8 + auto adown_4 = ADown(network, weightMap, *repncspelan_3->getOutput(0), 256, "model.4"); + // csp-elan block + // [-1, 1, RepNCSPELAN4, [512, 256, 128, 2]], # 5 + auto repncspelan_5 = RepNCSPELAN4(network, weightMap, *adown_4->getOutput(0), 256, 512, 256, 128, 2, "model.5"); + // avg-conv down + // [-1, 1, ADown, [512]], # 6-P4/16 + auto adown_6 = ADown(network, weightMap, *repncspelan_5->getOutput(0), 512, "model.6"); + // csp-elan block + // [-1, 1, RepNCSPELAN4, [1024, 512, 256, 2]], # 7 + auto repncspelan_7 = RepNCSPELAN4(network, weightMap, *adown_6->getOutput(0), 512, 1024, 512, 256, 2, "model.7"); + // avg-conv down + // [-1, 1, ADown, [1024]], # 8-P5/32 + auto adown_8 = ADown(network, weightMap, *repncspelan_7->getOutput(0), 1024, "model.8"); + // csp-elan block + // [-1, 1, RepNCSPELAN4, [1024, 512, 256, 2]], # 9 + auto repncspelan_9 = RepNCSPELAN4(network, weightMap, *adown_8->getOutput(0), 512, 1024, 512, 256, 2, "model.9"); + + // [1, 1, CBLinear, [[64]]], # 10 + auto cblinear_10 = CBLinear(network, weightMap, *conv_1->getOutput(0), { 64 }, 1, 1, 0, 1, "model.10"); + // [3, 1, CBLinear, [[64, 128]]], # 11 + auto cblinear_11 = CBLinear(network, weightMap, *repncspelan_3->getOutput(0), { 64, 128 }, 1, 1, 0, 1, "model.11"); + // [5, 1, CBLinear, [[64, 128, 256]]], # 12 + auto cblinear_12 = CBLinear(network, weightMap, *repncspelan_5->getOutput(0), { 64, 128, 256 }, 1, 1, 0, 1, "model.12"); + // [7, 1, CBLinear, [[64, 128, 256, 512]]], # 13 + auto cblinear_13 = CBLinear(network, weightMap, *repncspelan_7->getOutput(0), { 64, 128, 256, 512 }, 1, 1, 0, 1, "model.13"); + // [9, 1, CBLinear, [[64, 128, 256, 512, 1024]]], # 14 + auto cblinear_14 = CBLinear(network, weightMap, *repncspelan_9->getOutput(0), { 64, 128, 256, 512, 1024 }, 1, 1, 0, 1, "model.14"); + + // conv down + // [0, 1, Conv, [64, 3, 2]], # 15-P1/2 + auto conv_15 = convBnSiLU(network, weightMap, *data, 64, 3, 2, 1, "model.15", 1); + // [[10, 11, 12, 13, 14, -1], 1, CBFuse, [[0, 0, 0, 0, 0]]], # 16 + auto cbfuse_16 = CBFuse(network, { cblinear_10, cblinear_11, cblinear_12, cblinear_13, cblinear_14, std::vector{ conv_15 } }, { 0, 0, 0, 0, 0, 0 }, { 2, 4, 8, 16, 32, 2 }); + + // conv down + // [-1, 1, Conv, [128, 3, 2]], # 17-P2/4 + auto conv_17 = convBnSiLU(network, weightMap, *cbfuse_16->getOutput(0), 128, 3, 2, 1, "model.17"); + // [[11, 12, 13, 14, -1], 1, CBFuse, [[1, 1, 1, 1]]], # 18 + auto cbfuse_18 = CBFuse(network, { cblinear_11, cblinear_12, cblinear_13, cblinear_14, std::vector{ conv_17 } }, { 1, 1, 1, 1, 0 }, { 4, 8, 16, 32, 4 }); + + // csp-elan block + // [-1, 1, RepNCSPELAN4, [256, 128, 64, 2]], # 19 + auto repncspelan_19 = RepNCSPELAN4(network, weightMap, *cbfuse_18->getOutput(0), 128, 256, 128, 64, 2, "model.19"); + + // avg-conv down fuse + // [-1, 1, ADown, [256]], # 20-P3/8 + auto adown_20 = ADown(network, weightMap, *repncspelan_19->getOutput(0), 256, "model.20"); + // [[12, 13, 14, -1], 1, CBFuse, [[2, 2, 2]]], # 21 + auto cbfuse_21 = CBFuse(network, { cblinear_12, cblinear_13, cblinear_14, std::vector{ adown_20 } }, { 2, 2, 2, 0 }, { 8, 16, 32, 8 }); + + // csp-elan block + // [-1, 1, RepNCSPELAN4, [512, 256, 128, 2]], # 22 + auto repncspelan_22 = RepNCSPELAN4(network, weightMap, *cbfuse_21->getOutput(0), 256, 512, 256, 128, 2, "model.22"); + + // avg-conv down fuse + // [-1, 1, ADown, [512]], # 23-P4/16 + auto adown_23 = ADown(network, weightMap, *repncspelan_22->getOutput(0), 512, "model.23"); + // [[13, 14, -1], 1, CBFuse, [[3, 3]]], # 24 + auto cbfuse_24 = CBFuse(network, { cblinear_13, cblinear_14, std::vector{ adown_23 } }, { 3, 3, 0 }, { 16, 32, 16 }); + + // csp-elan block + // [-1, 1, RepNCSPELAN4, [1024, 512, 256, 2]], # 25 + auto repncspelan_25 = RepNCSPELAN4(network, weightMap, *cbfuse_24->getOutput(0), 512, 1024, 512, 256, 2, "model.25"); + + // avg-conv down fuse + // [-1, 1, ADown, [1024]], # 26-P5/32 + auto adown_26 = ADown(network, weightMap, *repncspelan_25->getOutput(0), 1024, "model.26"); + // [[14, -1], 1, CBFuse, [[4]]], # 27 + auto cbfuse_27 = CBFuse(network, { cblinear_14, std::vector{ adown_26 } }, { 4, 0 }, { 32, 32 }); + + // csp-elan block + // [-1, 1, RepNCSPELAN4, [1024, 512, 256, 2]], # 28 + auto repncspelan_28 = RepNCSPELAN4(network, weightMap, *cbfuse_27->getOutput(0), 512, 1024, 512, 256, 2, "model.28"); + + // elan-spp block + // [9, 1, SPPELAN, [512, 256]], # 29 + auto sppelan_29 = SPPELAN(network, weightMap, *repncspelan_9->getOutput(0), 1024, 512, 256, "model.29"); + + // # up-concat merge + // [-1, 1, nn.Upsample, [None, 2, 'nearest']], + auto upsample_30 = network->addResize(*sppelan_29->getOutput(0)); + upsample_30->setResizeMode(ResizeMode::kNEAREST); + const float scales_30[] = { 1.0, 2.0, 2.0 }; + upsample_30->setScales(scales_30, 3); + // [[-1, 7], 1, Concat, [1]], # cat backbone P4 + ITensor* input_tensor_31[] = { upsample_30->getOutput(0), repncspelan_7->getOutput(0) }; + auto cat_31 = network->addConcatenation(input_tensor_31, 2); + + // # csp-elan block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 2]], # 32 + auto repncspelan_32 = RepNCSPELAN4(network, weightMap, *cat_31->getOutput(0), 1536, 512, 512, 256, 2, "model.32"); + + // # up-concat merge + // [-1, 1, nn.Upsample, [None, 2, 'nearest']], + auto upsample_33 = network->addResize(*repncspelan_32->getOutput(0)); + upsample_33->setResizeMode(ResizeMode::kNEAREST); + const float scales_33[] = { 1.0, 2.0, 2.0 }; + upsample_33->setScales(scales_33, 3); + // [[-1, 5], 1, Concat, [1]], # cat backbone P3 + ITensor* input_tensor_34[] = { upsample_33->getOutput(0), repncspelan_5->getOutput(0) }; + auto cat_34 = network->addConcatenation(input_tensor_34, 2); + + // # csp-elan block + // [-1, 1, RepNCSPELAN4, [256, 256, 128, 2]], # 35 + auto repncspelan_35 = RepNCSPELAN4(network, weightMap, *cat_34->getOutput(0), 1024, 256, 256, 128, 2, "model.35"); + + // # elan-spp block + // [28, 1, SPPELAN, [512, 256]], # 36 + auto sppelan_36 = SPPELAN(network, weightMap, *repncspelan_28->getOutput(0), 1024, 512, 256, "model.36"); + + // # up-concat merge + // [-1, 1, nn.Upsample, [None, 2, 'nearest']], + auto upsample_37 = network->addResize(*sppelan_36->getOutput(0)); + upsample_37->setResizeMode(ResizeMode::kNEAREST); + const float scales_37[] = { 1.0, 2.0, 2.0 }; + upsample_37->setScales(scales_37, 3); + // [[-1, 25], 1, Concat, [1]], # cat backbone P4 + ITensor* input_tensor_38[] = { upsample_37->getOutput(0), repncspelan_25->getOutput(0) }; + auto cat_38 = network->addConcatenation(input_tensor_38, 2); + + // # csp-elan block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 2]], # 39 + auto repncspelan_39 = RepNCSPELAN4(network, weightMap, *cat_38->getOutput(0), 1536, 512, 512, 256, 2, "model.39"); + + // # up-concat merge + // [-1, 1, nn.Upsample, [None, 2, 'nearest']], + auto upsample_40 = network->addResize(*repncspelan_39->getOutput(0)); + upsample_40->setResizeMode(ResizeMode::kNEAREST); + const float scales_40[] = { 1.0, 2.0, 2.0 }; + upsample_40->setScales(scales_40, 3); + // [[-1, 22], 1, Concat, [1]], # cat backbone P3 + ITensor* input_tensor_41[] = { upsample_40->getOutput(0), repncspelan_22->getOutput(0) }; + auto cat_41 = network->addConcatenation(input_tensor_41, 2); + + // # csp-elan block + // [-1, 1, RepNCSPELAN4, [256, 256, 128, 2]], # 42 (P3/8-small) + auto repncspelan_42 = RepNCSPELAN4(network, weightMap, *cat_41->getOutput(0), 1024, 256, 256, 128, 2, "model.42"); + // # avg-conv-down merge + // [-1, 1, ADown, [256]], + auto adown_43 = ADown(network, weightMap, *repncspelan_42->getOutput(0), 256, "model.43"); + // [[-1, 39], 1, Concat, [1]], # cat head P4 + ITensor* input_tensor_44[] = { adown_43->getOutput(0), repncspelan_39->getOutput(0) }; + auto cat_44 = network->addConcatenation(input_tensor_44, 2); + + // # csp-elan block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 2]], # 45 (P4/16-medium) + auto repncspelan_45 = RepNCSPELAN4(network, weightMap, *cat_44->getOutput(0), 768, 512, 512, 256, 2, "model.45"); + // # avg-conv-down merge + // [-1, 1, ADown, [512]], + auto adown_46 = ADown(network, weightMap, *repncspelan_45->getOutput(0), 512, "model.46"); + // [[-1, 36], 1, Concat, [1]], # cat head P5 + ITensor* input_tensor_47[] = { adown_46->getOutput(0), sppelan_36->getOutput(0) }; + auto cat_47 = network->addConcatenation(input_tensor_47, 2); + + // # csp-elan block + // [-1, 1, RepNCSPELAN4, [512, 1024, 512, 2]], # 48 (P5/32-large) + auto repncspelan_48 = RepNCSPELAN4(network, weightMap, *cat_47->getOutput(0), 1024, 512, 1024, 512, 2, "model.48"); + + // auto DualDDetect_49 = DualDDetect(network, weightMap, std::vector{RepNCSPELAN_42, RepNCSPELAN_45, RepNCSPELAN_48}, kNumClass, {256, 512, 512}, "model.49"); + auto dualddetect_49 = DualDDetect(network, weightMap, std::vector{ repncspelan_35, repncspelan_32, sppelan_29 }, kNumClass, { 256, 512, 512 }, "model.49"); + + nvinfer1::IPluginV2Layer* yolo = addYoLoLayer(network, dualddetect_49, false); + yolo->getOutput(0)->setName(kOutputTensorName); + network->markOutput(*yolo->getOutput(0)); + + builder->setMaxBatchSize(kBatchSize); + config->setMaxWorkspaceSize(16 * (1 << 20)); + +#if defined(USE_FP16) + config->setFlag(nvinfer1::BuilderFlag::kFP16); +#elif defined(USE_INT8) + std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; + assert(builder->platformHasFastInt8()); + config->setFlag(nvinfer1::BuilderFlag::kINT8); + auto* calibrator = new Int8EntropyCalibrator2(1, kInputW, kInputH, gCalibTablePath, "int8calib.table", kInputTensorName); + config->setInt8Calibrator(calibrator); +#endif + + std::cout << "Building engine, please wait for a while..." << std::endl; + IHostMemory* serialized_model = builder->buildSerializedNetwork(*network, *config); + std::cout << "Build engine successfully!" << std::endl; + + delete network; + + // Release host memory + for (auto& mem : weightMap) { + free((void*)(mem.second.values)); + } + + return serialized_model; +} +IHostMemory* build_engine_yolov9_c(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, std::string& wts_name) { + /* ------ Create the builder ------ */ + INetworkDefinition* network = builder->createNetworkV2(0U); + + ITensor* data = network->addInput(kInputTensorName, dt, Dims3{ 3, kInputH, kInputW }); + assert(data); + std::map weightMap = loadWeights(wts_name); + + // # conv down + // [-1, 1, Conv, [64, 3, 2]], # 1-P1/2 + auto conv_1 = convBnSiLU(network, weightMap, *data, 64, 3, 2, 1, "model.1", 1); + // # conv down + // [-1, 1, Conv, [128, 3, 2]], # 2-P2/4 + auto conv_2 = convBnSiLU(network, weightMap, *conv_1->getOutput(0), 128, 3, 2, 1, "model.2"); + // # elan-1 block + // [-1, 1, RepNCSPELAN4, [256, 128, 64, 1]], # 3 + auto repncspelan_3 = RepNCSPELAN4(network, weightMap, *conv_2->getOutput(0), 128, 256, 128, 64, 1, "model.3"); + // # avg-conv down + // [-1, 1, ADown, [256]], # 4-P3/8 + auto adown_4 = ADown(network, weightMap, *repncspelan_3->getOutput(0), 256, "model.4"); + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 256, 128, 1]], # 5 + auto repncspelan_5 = RepNCSPELAN4(network, weightMap, *adown_4->getOutput(0), 256, 512, 256, 128, 1, "model.5"); + // # avg-conv down + // [-1, 1, ADown, [512]], # 6-P4/16 + auto adown_6 = ADown(network, weightMap, *repncspelan_5->getOutput(0), 512, "model.6"); + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 7 + auto repncspelan_7 = RepNCSPELAN4(network, weightMap, *adown_6->getOutput(0), 512, 512, 512, 256, 1, "model.7"); + // # avg-conv down + // [-1, 1, ADown, [512]], # 8-P5/32 + auto adown_8 = ADown(network, weightMap, *repncspelan_7->getOutput(0), 512, "model.8"); + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 9 + auto repncspelan_9 = RepNCSPELAN4(network, weightMap, *adown_8->getOutput(0), 512, 512, 512, 256, 1, "model.9"); + // # elan-spp block + // [-1, 1, SPPELAN, [512, 256]], # 10 + auto sppelan_10 = SPPELAN(network, weightMap, *repncspelan_9->getOutput(0), 512, 512, 256, "model.10"); + + // # up-concat merge + // [-1, 1, nn.Upsample, [None, 2, 'nearest']], + auto upsample_11 = network->addResize(*sppelan_10->getOutput(0)); + upsample_11->setResizeMode(ResizeMode::kNEAREST); + const float scales_11[] = { 1.0, 2.0, 2.0 }; + upsample_11->setScales(scales_11, 3); + // [[-1, 7], 1, Concat, [1]], # cat backbone P4 + ITensor* input_tensor_12[] = { upsample_11->getOutput(0), repncspelan_7->getOutput(0) }; + auto cat_12 = network->addConcatenation(input_tensor_12, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 13 + auto repncspelan_13 = RepNCSPELAN4(network, weightMap, *cat_12->getOutput(0), 1536, 512, 512, 256, 1, "model.13"); + + // # up-concat merge + // [-1, 1, nn.Upsample, [None, 2, 'nearest']], + auto upsample_14 = network->addResize(*repncspelan_13->getOutput(0)); + upsample_14->setResizeMode(ResizeMode::kNEAREST); + const float scales_14[] = { 1.0, 2.0, 2.0 }; + upsample_14->setScales(scales_14, 3); + // [[-1, 5], 1, Concat, [1]], # cat backbone P3 + ITensor* input_tensor_15[] = { upsample_14->getOutput(0), repncspelan_5->getOutput(0) }; + auto cat_15 = network->addConcatenation(input_tensor_15, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [256, 256, 128, 1]], # 16 (P3/8-small) + auto repncspelan_16 = RepNCSPELAN4(network, weightMap, *cat_15->getOutput(0), 1024, 256, 256, 128, 1, "model.16"); + + // # avg-conv-down merge + // [-1, 1, ADown, [256]], + auto adown_17 = ADown(network, weightMap, *repncspelan_16->getOutput(0), 256, "model.17"); + // [[-1, 13], 1, Concat, [1]], # cat head P4 + ITensor* input_tensor_18[] = { adown_17->getOutput(0), repncspelan_13->getOutput(0) }; + auto cat_18 = network->addConcatenation(input_tensor_18, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 19 (P4/16-medium) + auto repncspelan_19 = RepNCSPELAN4(network, weightMap, *cat_18->getOutput(0), 768, 512, 512, 256, 1, "model.19"); + + // # avg-conv-down merge + // [-1, 1, ADown, [512]], + auto adown_20 = ADown(network, weightMap, *repncspelan_19->getOutput(0), 512, "model.20"); + // [[-1, 10], 1, Concat, [1]], # cat head P5 + ITensor* input_tensor_21[] = { adown_20->getOutput(0), sppelan_10->getOutput(0) }; + auto cat_21 = network->addConcatenation(input_tensor_21, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 22 (P5/32-large) + auto repncspelan_22 = RepNCSPELAN4(network, weightMap, *cat_21->getOutput(0), 1024, 512, 512, 256, 1, "model.22"); + + // # multi-level reversible auxiliary branch + + // # routing + // [5, 1, CBLinear, [[256]]], # 23 + auto cblinear_23 = CBLinear(network, weightMap, *repncspelan_5->getOutput(0), { 256 }, 1, 1, 0, 1, "model.23"); + // [7, 1, CBLinear, [[256, 512]]], # 24 + auto cblinear_24 = CBLinear(network, weightMap, *repncspelan_7->getOutput(0), { 256, 512 }, 1, 1, 0, 1, "model.24"); + // [9, 1, CBLinear, [[256, 512, 512]]], # 25 + auto cblinear_25 = CBLinear(network, weightMap, *repncspelan_9->getOutput(0), { 256, 512, 512 }, 1, 1, 0, 1, "model.25"); + + // # conv down + // [0, 1, Conv, [64, 3, 2]], # 26-P1/2 + auto conv_26 = convBnSiLU(network, weightMap, *data, 64, 3, 2, 1, "model.26", 1); + + // # conv down + // [-1, 1, Conv, [128, 3, 2]], # 27-P2/4 + auto conv_27 = convBnSiLU(network, weightMap, *conv_26->getOutput(0), 128, 3, 2, 1, "model.27"); + + // # elan-1 block + // [-1, 1, RepNCSPELAN4, [256, 128, 64, 1]], # 28 + auto repncspelan_28 = RepNCSPELAN4(network, weightMap, *conv_27->getOutput(0), 128, 256, 128, 64, 1, "model.28"); + + // # avg-conv down fuse + // [-1, 1, ADown, [256]], # 29-P3/8 + auto adown_29 = ADown(network, weightMap, *repncspelan_28->getOutput(0), 256, "model.29"); + // [[23, 24, 25, -1], 1, CBFuse, [[0, 0, 0]]], # 30 + auto cbfuse = CBFuse(network, { cblinear_23, cblinear_24, cblinear_25, std::vector{ adown_29 } }, { 0, 0, 0, 0 }, { 8, 16, 32, 8 }); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 256, 128, 1]], # 31 + auto repncspelan_31 = RepNCSPELAN4(network, weightMap, *cbfuse->getOutput(0), 256, 512, 256, 128, 1, "model.31"); + + // # avg-conv down fuse + // [-1, 1, ADown, [512]], # 32-P4/16 + auto adown_32 = ADown(network, weightMap, *repncspelan_31->getOutput(0), 512, "model.32"); + // [[24, 25, -1], 1, CBFuse, [[1, 1]]], # 33 + auto cbfuse_33 = CBFuse(network, { cblinear_24, cblinear_25, std::vector{ adown_32 } }, { 1, 1, 0 }, { 16, 32, 16 }); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 34 + auto repncspelan_34 = RepNCSPELAN4(network, weightMap, *cbfuse_33->getOutput(0), 512, 512, 512, 256, 1, "model.34"); + + // # avg-conv down fuse + // [-1, 1, ADown, [512]], # 35-P5/32 + auto adown_35 = ADown(network, weightMap, *repncspelan_34->getOutput(0), 512, "model.35"); + + + // [[25, -1], 1, CBFuse, [[2]]], # 36 + auto cbfuse_36 = CBFuse(network, { cblinear_25, std::vector{ adown_35 } }, { 2, 0 }, { 32, 32 }); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 37 + auto repncspelan_37 = RepNCSPELAN4(network, weightMap, *cbfuse_36->getOutput(0), 512, 512, 512, 256, 1, "model.37"); + + // # detection head + // # detect + // [[31, 34, 37, 16, 19, 22], 1, DualDDetect, [nc]], # DualDDetect(A3, A4, A5, P3, P4, P5) + auto dualddetect_38 = DualDDetect(network, weightMap, std::vector{ repncspelan_31, repncspelan_34, repncspelan_37 }, kNumClass, { 512, 512, 512 }, "model.38"); + + nvinfer1::IPluginV2Layer* yolo = addYoLoLayer(network, dualddetect_38, false); + yolo->getOutput(0)->setName(kOutputTensorName); + network->markOutput(*yolo->getOutput(0)); + + builder->setMaxBatchSize(kBatchSize); + config->setMaxWorkspaceSize(16 * (1 << 20)); + +#if defined(USE_FP16) + config->setFlag(nvinfer1::BuilderFlag::kFP16); +#elif defined(USE_INT8) + std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; + assert(builder->platformHasFastInt8()); + config->setFlag(nvinfer1::BuilderFlag::kINT8); + auto* calibrator = new Int8EntropyCalibrator2(1, kInputW, kInputH, gCalibTablePath, "int8calib.table", kInputTensorName); + config->setInt8Calibrator(calibrator); +#endif + + std::cout << "Building engine, please wait for a while..." << std::endl; + IHostMemory* serialized_model = builder->buildSerializedNetwork(*network, *config); + std::cout << "Build engine successfully!" << std::endl; + + delete network; + + // Release host memory + for (auto& mem : weightMap) { + free((void*)(mem.second.values)); + } + + return serialized_model; + +} \ No newline at end of file diff --git a/yolov9/src/postprocess.cpp b/yolov9/src/postprocess.cpp new file mode 100644 index 0000000..be25b62 --- /dev/null +++ b/yolov9/src/postprocess.cpp @@ -0,0 +1,228 @@ +#include "postprocess.h" +#include "utils.h" + +cv::Rect get_rect(cv::Mat& img, float bbox[4]) { + float l, r, t, b; + float r_w = kInputW / (img.cols * 1.0); + float r_h = kInputH / (img.rows * 1.0); + if (r_h > r_w) { + l = bbox[0] - bbox[2] / 2.f; + r = bbox[0] + bbox[2] / 2.f; + t = bbox[1] - bbox[3] / 2.f - (kInputH - r_w * img.rows) / 2; + b = bbox[1] + bbox[3] / 2.f - (kInputH - r_w * img.rows) / 2; + l = l / r_w; + r = r / r_w; + t = t / r_w; + b = b / r_w; + } else { + l = bbox[0] - bbox[2] / 2.f - (kInputW - r_h * img.cols) / 2; + r = bbox[0] + bbox[2] / 2.f - (kInputW - r_h * img.cols) / 2; + t = bbox[1] - bbox[3] / 2.f; + b = bbox[1] + bbox[3] / 2.f; + l = l / r_h; + r = r / r_h; + t = t / r_h; + b = b / r_h; + } + return cv::Rect(round(l), round(t), round(r - l), round(b - t)); +} + +static float iou(float lbox[4], float rbox[4]) { + float interBox[] = { + (std::max)(lbox[0] - lbox[2] / 2.f , rbox[0] - rbox[2] / 2.f), //left + (std::min)(lbox[0] + lbox[2] / 2.f , rbox[0] + rbox[2] / 2.f), //right + (std::max)(lbox[1] - lbox[3] / 2.f , rbox[1] - rbox[3] / 2.f), //top + (std::min)(lbox[1] + lbox[3] / 2.f , rbox[1] + rbox[3] / 2.f), //bottom + }; + + if (interBox[2] > interBox[3] || interBox[0] > interBox[1]) + return 0.0f; + + float interBoxS = (interBox[1] - interBox[0])*(interBox[3] - interBox[2]); + return interBoxS / (lbox[2] * lbox[3] + rbox[2] * rbox[3] - interBoxS); +} + +static bool cmp(const Detection& a, const Detection& b) { + return a.conf > b.conf; +} + +void nms(std::vector& res, float* output, float conf_thresh, float nms_thresh) { + int det_size = sizeof(Detection) / sizeof(float); + std::map> m; + for (int i = 0; i < output[0] && i < kMaxNumOutputBbox; i++) { + if (output[1 + det_size * i + 4] <= conf_thresh) continue; + Detection det; + memcpy(&det, &output[1 + det_size * i], det_size * sizeof(float)); + if (m.count(det.class_id) == 0) m.emplace(det.class_id, std::vector()); + // x1x2y1y2 -> xywh + float c_x = (det.bbox[0]+det.bbox[2])/2; + float c_y = (det.bbox[1]+det.bbox[3])/2; + float w = det.bbox[2]-det.bbox[0]; + float h = det.bbox[3]-det.bbox[1]; + det.bbox[0] = c_x; + det.bbox[1] = c_y; + det.bbox[2] = w; + det.bbox[3] = h; + m[det.class_id].push_back(det); + } + for (auto it = m.begin(); it != m.end(); it++) { + auto& dets = it->second; + std::sort(dets.begin(), dets.end(), cmp); + for (size_t m = 0; m < dets.size(); ++m) { + auto& item = dets[m]; + res.push_back(item); + for (size_t n = m + 1; n < dets.size(); ++n) { + if (iou(item.bbox, dets[n].bbox) > nms_thresh) { + dets.erase(dets.begin() + n); + --n; + } + } + } + } +} + +void batch_nms(std::vector>& res_batch, float *output, int batch_size, int output_size, float conf_thresh, float nms_thresh) { + res_batch.resize(batch_size); + for (int i = 0; i < batch_size; i++) { + nms(res_batch[i], &output[i * output_size], conf_thresh, nms_thresh); + } +} + +void draw_bbox(std::vector& img_batch, std::vector>& res_batch) { + for (size_t i = 0; i < img_batch.size(); i++) { + auto& res = res_batch[i]; + cv::Mat img = img_batch[i]; + for (size_t j = 0; j < res.size(); j++) { + cv::Rect r = get_rect(img, res[j].bbox); + cv::rectangle(img, r, cv::Scalar(0x27, 0xC1, 0x36), 2); + cv::putText(img, std::to_string((int)res[j].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2); + } + } + // draw num of objets to img + for (size_t i = 0; i < img_batch.size(); i++) { + cv::putText(img_batch[i], std::to_string(res_batch[i].size()), cv::Point(0, 20), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2); + } +} + +static cv::Rect get_downscale_rect(float bbox[4], float scale) { + float left = bbox[0] - bbox[2] / 2; + float top = bbox[1] - bbox[3] / 2; + float right = bbox[0] + bbox[2] / 2; + float bottom = bbox[1] + bbox[3] / 2; + left /= scale; + top /= scale; + right /= scale; + bottom /= scale; + return cv::Rect(round(left), round(top), round(right - left), round(bottom - top)); +} + +// std::vector process_mask(const float* proto, int proto_size, std::vector& dets) { +// std::vector masks; +// for (size_t i = 0; i < dets.size(); i++) { +// cv::Mat mask_mat = cv::Mat::zeros(kInputH / 4, kInputW / 4, CV_32FC1); +// auto r = get_downscale_rect(dets[i].bbox, 4); +// for (int x = r.x; x < r.x + r.width; x++) { +// for (int y = r.y; y < r.y + r.height; y++) { +// float e = 0.0f; +// for (int j = 0; j < 32; j++) { +// e += dets[i].mask[j] * proto[j * proto_size / 32 + y * mask_mat.cols + x]; +// } +// e = 1.0f / (1.0f + expf(-e)); +// mask_mat.at(y, x) = e; +// } +// } +// cv::resize(mask_mat, mask_mat, cv::Size(kInputW, kInputH)); +// masks.push_back(mask_mat); +// } +// return masks; +// } + +cv::Mat scale_mask(cv::Mat mask, cv::Mat img) { + int x, y, w, h; + float r_w = kInputW / (img.cols * 1.0); + float r_h = kInputH / (img.rows * 1.0); + if (r_h > r_w) { + w = kInputW; + h = r_w * img.rows; + x = 0; + y = (kInputH - h) / 2; + } else { + w = r_h * img.cols; + h = kInputH; + x = (kInputW - w) / 2; + y = 0; + } + cv::Rect r(x, y, w, h); + cv::Mat res; + cv::resize(mask(r), res, img.size()); + return res; +} + +void draw_mask_bbox(cv::Mat& img, std::vector& dets, std::vector& masks, std::unordered_map& labels_map) { + static std::vector colors = {0xFF3838, 0xFF9D97, 0xFF701F, 0xFFB21D, 0xCFD231, 0x48F90A, + 0x92CC17, 0x3DDB86, 0x1A9334, 0x00D4BB, 0x2C99A8, 0x00C2FF, + 0x344593, 0x6473FF, 0x0018EC, 0x8438FF, 0x520085, 0xCB38FF, + 0xFF95C8, 0xFF37C7}; + for (size_t i = 0; i < dets.size(); i++) { + cv::Mat img_mask = scale_mask(masks[i], img); + auto color = colors[(int)dets[i].class_id % colors.size()]; + auto bgr = cv::Scalar(color & 0xFF, color >> 8 & 0xFF, color >> 16 & 0xFF); + + cv::Rect r = get_rect(img, dets[i].bbox); + for (int x = r.x; x < r.x + r.width; x++) { + for (int y = r.y; y < r.y + r.height; y++) { + float val = img_mask.at(y, x); + if (val <= 0.5) continue; + img.at(y, x)[0] = img.at(y, x)[0] / 2 + bgr[0] / 2; + img.at(y, x)[1] = img.at(y, x)[1] / 2 + bgr[1] / 2; + img.at(y, x)[2] = img.at(y, x)[2] / 2 + bgr[2] / 2; + } + } + + cv::rectangle(img, r, bgr, 2); + + // Get the size of the text + cv::Size textSize = cv::getTextSize(labels_map[(int)dets[i].class_id] + " " + to_string_with_precision(dets[i].conf), cv::FONT_HERSHEY_PLAIN, 1.2, 2, NULL); + // Set the top left corner of the rectangle + cv::Point topLeft(r.x, r.y - textSize.height); + + // Set the bottom right corner of the rectangle + cv::Point bottomRight(r.x + textSize.width, r.y + textSize.height); + + // Set the thickness of the rectangle lines + int lineThickness = 2; + + // Draw the rectangle on the image + cv::rectangle(img, topLeft, bottomRight, bgr, -1); + + cv::putText(img, labels_map[(int)dets[i].class_id] + " " + to_string_with_precision(dets[i].conf), cv::Point(r.x, r.y + 4), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar::all(0xFF), 2); + + } +} +void process_decode_ptr_host(std::vector &res, const float* decode_ptr_host, int bbox_element, cv::Mat& img, int count) { + Detection det; + for (int i = 0; i < count; i++) { + int basic_pos = 1 + i * bbox_element; + int keep_flag = decode_ptr_host[basic_pos + 6]; + if (keep_flag == 1) { + det.bbox[0] = decode_ptr_host[basic_pos + 0]; + det.bbox[1] = decode_ptr_host[basic_pos + 1]; + det.bbox[2] = decode_ptr_host[basic_pos + 2]; + det.bbox[3] = decode_ptr_host[basic_pos + 3]; + det.conf = decode_ptr_host[basic_pos + 4]; + det.class_id = decode_ptr_host[basic_pos + 5]; + res.push_back(det); + } + } +} +void batch_process(std::vector> &res_batch, const float* decode_ptr_host, int batch_size, int bbox_element, const std::vector& img_batch) { + res_batch.resize(batch_size); + int count = static_cast(*decode_ptr_host); + count = count > kMaxNumOutputBbox ? kMaxNumOutputBbox : count; + // std::min(count, kMaxNumOutputBbox); + for (int i = 0; i < batch_size; i++) { + auto& img = const_cast(img_batch[i]); + process_decode_ptr_host(res_batch[i], &decode_ptr_host[i * count], bbox_element, img, count); + } +} + diff --git a/yolov9/src/postprocess.cu b/yolov9/src/postprocess.cu new file mode 100644 index 0000000..a322654 --- /dev/null +++ b/yolov9/src/postprocess.cu @@ -0,0 +1,91 @@ +// +// Created by lindsay on 23-7-17. +// +#include "postprocess.h" + +static __global__ void decode_kernel(float *predict, int num_bboxes, float confidence_threshold, float *parray, int max_objects) { + + float count = predict[0]; + int position = (blockDim.x * blockIdx.x + threadIdx.x); + if (position >= count) + return; + float *pitem = predict + 1 + position * 6; + int index = atomicAdd(parray, 1); + if (index >= max_objects) + return; + float confidence = pitem[4]; + if (confidence < confidence_threshold) + return; + float *pout_item = parray + 1 + index * bbox_element; + float left = pitem[0]; + float top = pitem[1]; + float right = pitem[2]; + float bottom = pitem[3]; + float label = pitem[5]; + *pout_item++ = left; + *pout_item++ = top; + *pout_item++ = right; + *pout_item++ = bottom; + *pout_item++ = confidence; + *pout_item++ = label; + *pout_item++ = 1; // 1 = keep, 0 = ignore +} + +static __device__ float box_iou(float aleft, float atop, float aright, float abottom, float bleft, float btop, float bright, float bbottom) { + + float cleft = max(aleft, bleft); + float ctop = max(atop, btop); + float cright = min(aright, bright); + float cbottom = min(abottom, bbottom); + + float c_area = max(cright - cleft, 0.0f) * max(cbottom - ctop, 0.0f); + if (c_area == 0.0f) + return 0.0f; + + float a_area = max(0.0f, aright - aleft) * max(0.0f, abottom - atop); + float b_area = max(0.0f, bright - bleft) * max(0.0f, bbottom - btop); + return c_area / (a_area + b_area - c_area); +} + +static __global__ void nms_kernel(float *bboxes, int max_objects, float threshold) { + + int position = (blockDim.x * blockIdx.x + threadIdx.x); + int count = bboxes[0]; + + // float count = 0.0f; + if (position >= count) + return; + + float *pcurrent = bboxes + 1 + position * bbox_element; + for (int i = 1; i < count; ++i) { + float *pitem = bboxes + 1 + i * bbox_element; + if (i == position || pcurrent[5] != pitem[5]) continue; + + if (pitem[4] >= pcurrent[4]) { + if (pitem[4] == pcurrent[4] && i < position) + continue; + + float iou = box_iou(pcurrent[0], pcurrent[1], pcurrent[2], pcurrent[3],pitem[0], pitem[1], pitem[2], pitem[3]); + + if (iou > threshold) { + pcurrent[6] = 0; + return; + } + } + } +} +// 置信度过滤 +void cuda_decode(float *predict, int num_bboxes, float confidence_threshold, float *parray, int max_objects, + cudaStream_t stream) { + int block = 256; + int grid = ceil(num_bboxes / (float) block); + decode_kernel << < grid, block, 0, stream >> > ((float *) predict, num_bboxes, confidence_threshold, parray, max_objects); + +} + +void cuda_nms(float *parray, float nms_threshold, int max_objects, cudaStream_t stream) { + int block = max_objects < 256 ? max_objects : 256; + int grid = ceil(max_objects / (float) block); + nms_kernel << < grid, block, 0, stream >> > (parray, max_objects, nms_threshold); + +} diff --git a/yolov9/src/preprocess.cu b/yolov9/src/preprocess.cu new file mode 100644 index 0000000..2e2aca0 --- /dev/null +++ b/yolov9/src/preprocess.cu @@ -0,0 +1,153 @@ +#include "preprocess.h" +#include "cuda_utils.h" + +static uint8_t* img_buffer_host = nullptr; +static uint8_t* img_buffer_device = nullptr; + +struct AffineMatrix { + float value[6]; +}; + +__global__ void warpaffine_kernel( + uint8_t* src, int src_line_size, int src_width, + int src_height, float* dst, int dst_width, + int dst_height, uint8_t const_value_st, + AffineMatrix d2s, int edge) { + int position = blockDim.x * blockIdx.x + threadIdx.x; + if (position >= edge) return; + + float m_x1 = d2s.value[0]; + float m_y1 = d2s.value[1]; + float m_z1 = d2s.value[2]; + float m_x2 = d2s.value[3]; + float m_y2 = d2s.value[4]; + float m_z2 = d2s.value[5]; + + int dx = position % dst_width; + int dy = position / dst_width; + float src_x = m_x1 * dx + m_y1 * dy + m_z1 + 0.5f; + float src_y = m_x2 * dx + m_y2 * dy + m_z2 + 0.5f; + float c0, c1, c2; + + if (src_x <= -1 || src_x >= src_width || src_y <= -1 || src_y >= src_height) { + // out of range + c0 = const_value_st; + c1 = const_value_st; + c2 = const_value_st; + } else { + int y_low = floorf(src_y); + int x_low = floorf(src_x); + int y_high = y_low + 1; + int x_high = x_low + 1; + + uint8_t const_value[] = {const_value_st, const_value_st, const_value_st}; + float ly = src_y - y_low; + float lx = src_x - x_low; + float hy = 1 - ly; + float hx = 1 - lx; + float w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx; + uint8_t* v1 = const_value; + uint8_t* v2 = const_value; + uint8_t* v3 = const_value; + uint8_t* v4 = const_value; + + if (y_low >= 0) { + if (x_low >= 0) + v1 = src + y_low * src_line_size + x_low * 3; + + if (x_high < src_width) + v2 = src + y_low * src_line_size + x_high * 3; + } + + if (y_high < src_height) { + if (x_low >= 0) + v3 = src + y_high * src_line_size + x_low * 3; + + if (x_high < src_width) + v4 = src + y_high * src_line_size + x_high * 3; + } + + c0 = w1 * v1[0] + w2 * v2[0] + w3 * v3[0] + w4 * v4[0]; + c1 = w1 * v1[1] + w2 * v2[1] + w3 * v3[1] + w4 * v4[1]; + c2 = w1 * v1[2] + w2 * v2[2] + w3 * v3[2] + w4 * v4[2]; + } + + // bgr to rgb + float t = c2; + c2 = c0; + c0 = t; + + // normalization + c0 = c0 / 255.0f; + c1 = c1 / 255.0f; + c2 = c2 / 255.0f; + + // rgbrgbrgb to rrrgggbbb + int area = dst_width * dst_height; + float* pdst_c0 = dst + dy * dst_width + dx; + float* pdst_c1 = pdst_c0 + area; + float* pdst_c2 = pdst_c1 + area; + *pdst_c0 = c0; + *pdst_c1 = c1; + *pdst_c2 = c2; +} + +void cuda_preprocess( + uint8_t* src, int src_width, int src_height, + float* dst, int dst_width, int dst_height, + cudaStream_t stream) { + + int img_size = src_width * src_height * 3; + // copy data to pinned memory + memcpy(img_buffer_host, src, img_size); + // copy data to device memory + CUDA_CHECK(cudaMemcpyAsync(img_buffer_device, img_buffer_host, img_size, cudaMemcpyHostToDevice, stream)); + + AffineMatrix s2d, d2s; + float scale = std::min(dst_height / (float)src_height, dst_width / (float)src_width); + + s2d.value[0] = scale; + s2d.value[1] = 0; + s2d.value[2] = -scale * src_width * 0.5 + dst_width * 0.5; + s2d.value[3] = 0; + s2d.value[4] = scale; + s2d.value[5] = -scale * src_height * 0.5 + dst_height * 0.5; + + cv::Mat m2x3_s2d(2, 3, CV_32F, s2d.value); + cv::Mat m2x3_d2s(2, 3, CV_32F, d2s.value); + cv::invertAffineTransform(m2x3_s2d, m2x3_d2s); + + memcpy(d2s.value, m2x3_d2s.ptr(0), sizeof(d2s.value)); + + int jobs = dst_height * dst_width; + int threads = 256; + int blocks = ceil(jobs / (float)threads); + + warpaffine_kernel<<>>( + img_buffer_device, src_width * 3, src_width, + src_height, dst, dst_width, + dst_height, 128, d2s, jobs); +} + +void cuda_batch_preprocess(std::vector& img_batch, + float* dst, int dst_width, int dst_height, + cudaStream_t stream) { + int dst_size = dst_width * dst_height * 3; + for (size_t i = 0; i < img_batch.size(); i++) { + cuda_preprocess(img_batch[i].ptr(), img_batch[i].cols, img_batch[i].rows, &dst[dst_size * i], dst_width, dst_height, stream); + CUDA_CHECK(cudaStreamSynchronize(stream)); + } +} + +void cuda_preprocess_init(int max_image_size) { + // prepare input data in pinned memory + CUDA_CHECK(cudaMallocHost((void**)&img_buffer_host, max_image_size * 3)); + // prepare input data in device memory + CUDA_CHECK(cudaMalloc((void**)&img_buffer_device, max_image_size * 3)); +} + +void cuda_preprocess_destroy() { + CUDA_CHECK(cudaFree(img_buffer_device)); + CUDA_CHECK(cudaFreeHost(img_buffer_host)); +} + diff --git a/yolov9/windows/dirent.h b/yolov9/windows/dirent.h new file mode 100644 index 0000000..b44059a --- /dev/null +++ b/yolov9/windows/dirent.h @@ -0,0 +1,1167 @@ +/* + * Dirent interface for Microsoft Visual Studio + * + * Copyright (C) 1998-2019 Toni Ronkko + * This file is part of dirent. Dirent may be freely distributed + * under the MIT license. For all details and documentation, see + * https://github.com/tronkko/dirent + */ +#ifndef DIRENT_H +#define DIRENT_H + +/* Hide warnings about unreferenced local functions */ +#if defined(__clang__) +# pragma clang diagnostic ignored "-Wunused-function" +#elif defined(_MSC_VER) +# pragma warning(disable:4505) +#elif defined(__GNUC__) +# pragma GCC diagnostic ignored "-Wunused-function" +#endif + +/* + * Include windows.h without Windows Sockets 1.1 to prevent conflicts with + * Windows Sockets 2.0. + */ +#ifndef WIN32_LEAN_AND_MEAN +# define WIN32_LEAN_AND_MEAN +#endif +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +/* Indicates that d_type field is available in dirent structure */ +#define _DIRENT_HAVE_D_TYPE + +/* Indicates that d_namlen field is available in dirent structure */ +#define _DIRENT_HAVE_D_NAMLEN + +/* Entries missing from MSVC 6.0 */ +#if !defined(FILE_ATTRIBUTE_DEVICE) +# define FILE_ATTRIBUTE_DEVICE 0x40 +#endif + +/* File type and permission flags for stat(), general mask */ +#if !defined(S_IFMT) +# define S_IFMT _S_IFMT +#endif + +/* Directory bit */ +#if !defined(S_IFDIR) +# define S_IFDIR _S_IFDIR +#endif + +/* Character device bit */ +#if !defined(S_IFCHR) +# define S_IFCHR _S_IFCHR +#endif + +/* Pipe bit */ +#if !defined(S_IFFIFO) +# define S_IFFIFO _S_IFFIFO +#endif + +/* Regular file bit */ +#if !defined(S_IFREG) +# define S_IFREG _S_IFREG +#endif + +/* Read permission */ +#if !defined(S_IREAD) +# define S_IREAD _S_IREAD +#endif + +/* Write permission */ +#if !defined(S_IWRITE) +# define S_IWRITE _S_IWRITE +#endif + +/* Execute permission */ +#if !defined(S_IEXEC) +# define S_IEXEC _S_IEXEC +#endif + +/* Pipe */ +#if !defined(S_IFIFO) +# define S_IFIFO _S_IFIFO +#endif + +/* Block device */ +#if !defined(S_IFBLK) +# define S_IFBLK 0 +#endif + +/* Link */ +#if !defined(S_IFLNK) +# define S_IFLNK 0 +#endif + +/* Socket */ +#if !defined(S_IFSOCK) +# define S_IFSOCK 0 +#endif + +/* Read user permission */ +#if !defined(S_IRUSR) +# define S_IRUSR S_IREAD +#endif + +/* Write user permission */ +#if !defined(S_IWUSR) +# define S_IWUSR S_IWRITE +#endif + +/* Execute user permission */ +#if !defined(S_IXUSR) +# define S_IXUSR 0 +#endif + +/* Read group permission */ +#if !defined(S_IRGRP) +# define S_IRGRP 0 +#endif + +/* Write group permission */ +#if !defined(S_IWGRP) +# define S_IWGRP 0 +#endif + +/* Execute group permission */ +#if !defined(S_IXGRP) +# define S_IXGRP 0 +#endif + +/* Read others permission */ +#if !defined(S_IROTH) +# define S_IROTH 0 +#endif + +/* Write others permission */ +#if !defined(S_IWOTH) +# define S_IWOTH 0 +#endif + +/* Execute others permission */ +#if !defined(S_IXOTH) +# define S_IXOTH 0 +#endif + +/* Maximum length of file name */ +#if !defined(PATH_MAX) +# define PATH_MAX MAX_PATH +#endif +#if !defined(FILENAME_MAX) +# define FILENAME_MAX MAX_PATH +#endif +#if !defined(NAME_MAX) +# define NAME_MAX FILENAME_MAX +#endif + +/* File type flags for d_type */ +#define DT_UNKNOWN 0 +#define DT_REG S_IFREG +#define DT_DIR S_IFDIR +#define DT_FIFO S_IFIFO +#define DT_SOCK S_IFSOCK +#define DT_CHR S_IFCHR +#define DT_BLK S_IFBLK +#define DT_LNK S_IFLNK + +/* Macros for converting between st_mode and d_type */ +#define IFTODT(mode) ((mode) & S_IFMT) +#define DTTOIF(type) (type) + +/* + * File type macros. Note that block devices, sockets and links cannot be + * distinguished on Windows and the macros S_ISBLK, S_ISSOCK and S_ISLNK are + * only defined for compatibility. These macros should always return false + * on Windows. + */ +#if !defined(S_ISFIFO) +# define S_ISFIFO(mode) (((mode) & S_IFMT) == S_IFIFO) +#endif +#if !defined(S_ISDIR) +# define S_ISDIR(mode) (((mode) & S_IFMT) == S_IFDIR) +#endif +#if !defined(S_ISREG) +# define S_ISREG(mode) (((mode) & S_IFMT) == S_IFREG) +#endif +#if !defined(S_ISLNK) +# define S_ISLNK(mode) (((mode) & S_IFMT) == S_IFLNK) +#endif +#if !defined(S_ISSOCK) +# define S_ISSOCK(mode) (((mode) & S_IFMT) == S_IFSOCK) +#endif +#if !defined(S_ISCHR) +# define S_ISCHR(mode) (((mode) & S_IFMT) == S_IFCHR) +#endif +#if !defined(S_ISBLK) +# define S_ISBLK(mode) (((mode) & S_IFMT) == S_IFBLK) +#endif + +/* Return the exact length of the file name without zero terminator */ +#define _D_EXACT_NAMLEN(p) ((p)->d_namlen) + +/* Return the maximum size of a file name */ +#define _D_ALLOC_NAMLEN(p) ((PATH_MAX)+1) + + +#ifdef __cplusplus +extern "C" { +#endif + + +/* Wide-character version */ +struct _wdirent { + /* Always zero */ + long d_ino; + + /* File position within stream */ + long d_off; + + /* Structure size */ + unsigned short d_reclen; + + /* Length of name without \0 */ + size_t d_namlen; + + /* File type */ + int d_type; + + /* File name */ + wchar_t d_name[PATH_MAX+1]; +}; +typedef struct _wdirent _wdirent; + +struct _WDIR { + /* Current directory entry */ + struct _wdirent ent; + + /* Private file data */ + WIN32_FIND_DATAW data; + + /* True if data is valid */ + int cached; + + /* Win32 search handle */ + HANDLE handle; + + /* Initial directory name */ + wchar_t *patt; +}; +typedef struct _WDIR _WDIR; + +/* Multi-byte character version */ +struct dirent { + /* Always zero */ + long d_ino; + + /* File position within stream */ + long d_off; + + /* Structure size */ + unsigned short d_reclen; + + /* Length of name without \0 */ + size_t d_namlen; + + /* File type */ + int d_type; + + /* File name */ + char d_name[PATH_MAX+1]; +}; +typedef struct dirent dirent; + +struct DIR { + struct dirent ent; + struct _WDIR *wdirp; +}; +typedef struct DIR DIR; + + +/* Dirent functions */ +static DIR *opendir (const char *dirname); +static _WDIR *_wopendir (const wchar_t *dirname); + +static struct dirent *readdir (DIR *dirp); +static struct _wdirent *_wreaddir (_WDIR *dirp); + +static int readdir_r( + DIR *dirp, struct dirent *entry, struct dirent **result); +static int _wreaddir_r( + _WDIR *dirp, struct _wdirent *entry, struct _wdirent **result); + +static int closedir (DIR *dirp); +static int _wclosedir (_WDIR *dirp); + +static void rewinddir (DIR* dirp); +static void _wrewinddir (_WDIR* dirp); + +static int scandir (const char *dirname, struct dirent ***namelist, + int (*filter)(const struct dirent*), + int (*compare)(const struct dirent**, const struct dirent**)); + +static int alphasort (const struct dirent **a, const struct dirent **b); + +static int versionsort (const struct dirent **a, const struct dirent **b); + + +/* For compatibility with Symbian */ +#define wdirent _wdirent +#define WDIR _WDIR +#define wopendir _wopendir +#define wreaddir _wreaddir +#define wclosedir _wclosedir +#define wrewinddir _wrewinddir + + +/* Internal utility functions */ +static WIN32_FIND_DATAW *dirent_first (_WDIR *dirp); +static WIN32_FIND_DATAW *dirent_next (_WDIR *dirp); + +static int dirent_mbstowcs_s( + size_t *pReturnValue, + wchar_t *wcstr, + size_t sizeInWords, + const char *mbstr, + size_t count); + +static int dirent_wcstombs_s( + size_t *pReturnValue, + char *mbstr, + size_t sizeInBytes, + const wchar_t *wcstr, + size_t count); + +static void dirent_set_errno (int error); + + +/* + * Open directory stream DIRNAME for read and return a pointer to the + * internal working area that is used to retrieve individual directory + * entries. + */ +static _WDIR* +_wopendir( + const wchar_t *dirname) +{ + _WDIR *dirp; +#if WINAPI_FAMILY_PARTITION(WINAPI_PARTITION_DESKTOP) + /* Desktop */ + DWORD n; +#else + /* WinRT */ + size_t n; +#endif + wchar_t *p; + + /* Must have directory name */ + if (dirname == NULL || dirname[0] == '\0') { + dirent_set_errno (ENOENT); + return NULL; + } + + /* Allocate new _WDIR structure */ + dirp = (_WDIR*) malloc (sizeof (struct _WDIR)); + if (!dirp) { + return NULL; + } + + /* Reset _WDIR structure */ + dirp->handle = INVALID_HANDLE_VALUE; + dirp->patt = NULL; + dirp->cached = 0; + + /* + * Compute the length of full path plus zero terminator + * + * Note that on WinRT there's no way to convert relative paths + * into absolute paths, so just assume it is an absolute path. + */ +#if WINAPI_FAMILY_PARTITION(WINAPI_PARTITION_DESKTOP) + /* Desktop */ + n = GetFullPathNameW (dirname, 0, NULL, NULL); +#else + /* WinRT */ + n = wcslen (dirname); +#endif + + /* Allocate room for absolute directory name and search pattern */ + dirp->patt = (wchar_t*) malloc (sizeof (wchar_t) * n + 16); + if (dirp->patt == NULL) { + goto exit_closedir; + } + + /* + * Convert relative directory name to an absolute one. This + * allows rewinddir() to function correctly even when current + * working directory is changed between opendir() and rewinddir(). + * + * Note that on WinRT there's no way to convert relative paths + * into absolute paths, so just assume it is an absolute path. + */ +#if WINAPI_FAMILY_PARTITION(WINAPI_PARTITION_DESKTOP) + /* Desktop */ + n = GetFullPathNameW (dirname, n, dirp->patt, NULL); + if (n <= 0) { + goto exit_closedir; + } +#else + /* WinRT */ + wcsncpy_s (dirp->patt, n+1, dirname, n); +#endif + + /* Append search pattern \* to the directory name */ + p = dirp->patt + n; + switch (p[-1]) { + case '\\': + case '/': + case ':': + /* Directory ends in path separator, e.g. c:\temp\ */ + /*NOP*/; + break; + + default: + /* Directory name doesn't end in path separator */ + *p++ = '\\'; + } + *p++ = '*'; + *p = '\0'; + + /* Open directory stream and retrieve the first entry */ + if (!dirent_first (dirp)) { + goto exit_closedir; + } + + /* Success */ + return dirp; + + /* Failure */ + exit_closedir: + _wclosedir (dirp); + return NULL; +} + +/* + * Read next directory entry. + * + * Returns pointer to static directory entry which may be overwritten by + * subsequent calls to _wreaddir(). + */ +static struct _wdirent* +_wreaddir( + _WDIR *dirp) +{ + struct _wdirent *entry; + + /* + * Read directory entry to buffer. We can safely ignore the return value + * as entry will be set to NULL in case of error. + */ + (void) _wreaddir_r (dirp, &dirp->ent, &entry); + + /* Return pointer to statically allocated directory entry */ + return entry; +} + +/* + * Read next directory entry. + * + * Returns zero on success. If end of directory stream is reached, then sets + * result to NULL and returns zero. + */ +static int +_wreaddir_r( + _WDIR *dirp, + struct _wdirent *entry, + struct _wdirent **result) +{ + WIN32_FIND_DATAW *datap; + + /* Read next directory entry */ + datap = dirent_next (dirp); + if (datap) { + size_t n; + DWORD attr; + + /* + * Copy file name as wide-character string. If the file name is too + * long to fit in to the destination buffer, then truncate file name + * to PATH_MAX characters and zero-terminate the buffer. + */ + n = 0; + while (n < PATH_MAX && datap->cFileName[n] != 0) { + entry->d_name[n] = datap->cFileName[n]; + n++; + } + entry->d_name[n] = 0; + + /* Length of file name excluding zero terminator */ + entry->d_namlen = n; + + /* File type */ + attr = datap->dwFileAttributes; + if ((attr & FILE_ATTRIBUTE_DEVICE) != 0) { + entry->d_type = DT_CHR; + } else if ((attr & FILE_ATTRIBUTE_DIRECTORY) != 0) { + entry->d_type = DT_DIR; + } else { + entry->d_type = DT_REG; + } + + /* Reset dummy fields */ + entry->d_ino = 0; + entry->d_off = 0; + entry->d_reclen = sizeof (struct _wdirent); + + /* Set result address */ + *result = entry; + + } else { + + /* Return NULL to indicate end of directory */ + *result = NULL; + + } + + return /*OK*/0; +} + +/* + * Close directory stream opened by opendir() function. This invalidates the + * DIR structure as well as any directory entry read previously by + * _wreaddir(). + */ +static int +_wclosedir( + _WDIR *dirp) +{ + int ok; + if (dirp) { + + /* Release search handle */ + if (dirp->handle != INVALID_HANDLE_VALUE) { + FindClose (dirp->handle); + } + + /* Release search pattern */ + free (dirp->patt); + + /* Release directory structure */ + free (dirp); + ok = /*success*/0; + + } else { + + /* Invalid directory stream */ + dirent_set_errno (EBADF); + ok = /*failure*/-1; + + } + return ok; +} + +/* + * Rewind directory stream such that _wreaddir() returns the very first + * file name again. + */ +static void +_wrewinddir( + _WDIR* dirp) +{ + if (dirp) { + /* Release existing search handle */ + if (dirp->handle != INVALID_HANDLE_VALUE) { + FindClose (dirp->handle); + } + + /* Open new search handle */ + dirent_first (dirp); + } +} + +/* Get first directory entry (internal) */ +static WIN32_FIND_DATAW* +dirent_first( + _WDIR *dirp) +{ + WIN32_FIND_DATAW *datap; + DWORD error; + + /* Open directory and retrieve the first entry */ + dirp->handle = FindFirstFileExW( + dirp->patt, FindExInfoStandard, &dirp->data, + FindExSearchNameMatch, NULL, 0); + if (dirp->handle != INVALID_HANDLE_VALUE) { + + /* a directory entry is now waiting in memory */ + datap = &dirp->data; + dirp->cached = 1; + + } else { + + /* Failed to open directory: no directory entry in memory */ + dirp->cached = 0; + datap = NULL; + + /* Set error code */ + error = GetLastError (); + switch (error) { + case ERROR_ACCESS_DENIED: + /* No read access to directory */ + dirent_set_errno (EACCES); + break; + + case ERROR_DIRECTORY: + /* Directory name is invalid */ + dirent_set_errno (ENOTDIR); + break; + + case ERROR_PATH_NOT_FOUND: + default: + /* Cannot find the file */ + dirent_set_errno (ENOENT); + } + + } + return datap; +} + +/* + * Get next directory entry (internal). + * + * Returns + */ +static WIN32_FIND_DATAW* +dirent_next( + _WDIR *dirp) +{ + WIN32_FIND_DATAW *p; + + /* Get next directory entry */ + if (dirp->cached != 0) { + + /* A valid directory entry already in memory */ + p = &dirp->data; + dirp->cached = 0; + + } else if (dirp->handle != INVALID_HANDLE_VALUE) { + + /* Get the next directory entry from stream */ + if (FindNextFileW (dirp->handle, &dirp->data) != FALSE) { + /* Got a file */ + p = &dirp->data; + } else { + /* The very last entry has been processed or an error occurred */ + FindClose (dirp->handle); + dirp->handle = INVALID_HANDLE_VALUE; + p = NULL; + } + + } else { + + /* End of directory stream reached */ + p = NULL; + + } + + return p; +} + +/* + * Open directory stream using plain old C-string. + */ +static DIR* +opendir( + const char *dirname) +{ + struct DIR *dirp; + + /* Must have directory name */ + if (dirname == NULL || dirname[0] == '\0') { + dirent_set_errno (ENOENT); + return NULL; + } + + /* Allocate memory for DIR structure */ + dirp = (DIR*) malloc (sizeof (struct DIR)); + if (!dirp) { + return NULL; + } + { + int error; + wchar_t wname[PATH_MAX + 1]; + size_t n; + + /* Convert directory name to wide-character string */ + error = dirent_mbstowcs_s( + &n, wname, PATH_MAX + 1, dirname, PATH_MAX + 1); + if (error) { + /* + * Cannot convert file name to wide-character string. This + * occurs if the string contains invalid multi-byte sequences or + * the output buffer is too small to contain the resulting + * string. + */ + goto exit_free; + } + + + /* Open directory stream using wide-character name */ + dirp->wdirp = _wopendir (wname); + if (!dirp->wdirp) { + goto exit_free; + } + + } + + /* Success */ + return dirp; + + /* Failure */ + exit_free: + free (dirp); + return NULL; +} + +/* + * Read next directory entry. + */ +static struct dirent* +readdir( + DIR *dirp) +{ + struct dirent *entry; + + /* + * Read directory entry to buffer. We can safely ignore the return value + * as entry will be set to NULL in case of error. + */ + (void) readdir_r (dirp, &dirp->ent, &entry); + + /* Return pointer to statically allocated directory entry */ + return entry; +} + +/* + * Read next directory entry into called-allocated buffer. + * + * Returns zero on success. If the end of directory stream is reached, then + * sets result to NULL and returns zero. + */ +static int +readdir_r( + DIR *dirp, + struct dirent *entry, + struct dirent **result) +{ + WIN32_FIND_DATAW *datap; + + /* Read next directory entry */ + datap = dirent_next (dirp->wdirp); + if (datap) { + size_t n; + int error; + + /* Attempt to convert file name to multi-byte string */ + error = dirent_wcstombs_s( + &n, entry->d_name, PATH_MAX + 1, datap->cFileName, PATH_MAX + 1); + + /* + * If the file name cannot be represented by a multi-byte string, + * then attempt to use old 8+3 file name. This allows traditional + * Unix-code to access some file names despite of unicode + * characters, although file names may seem unfamiliar to the user. + * + * Be ware that the code below cannot come up with a short file + * name unless the file system provides one. At least + * VirtualBox shared folders fail to do this. + */ + if (error && datap->cAlternateFileName[0] != '\0') { + error = dirent_wcstombs_s( + &n, entry->d_name, PATH_MAX + 1, + datap->cAlternateFileName, PATH_MAX + 1); + } + + if (!error) { + DWORD attr; + + /* Length of file name excluding zero terminator */ + entry->d_namlen = n - 1; + + /* File attributes */ + attr = datap->dwFileAttributes; + if ((attr & FILE_ATTRIBUTE_DEVICE) != 0) { + entry->d_type = DT_CHR; + } else if ((attr & FILE_ATTRIBUTE_DIRECTORY) != 0) { + entry->d_type = DT_DIR; + } else { + entry->d_type = DT_REG; + } + + /* Reset dummy fields */ + entry->d_ino = 0; + entry->d_off = 0; + entry->d_reclen = sizeof (struct dirent); + + } else { + + /* + * Cannot convert file name to multi-byte string so construct + * an erroneous directory entry and return that. Note that + * we cannot return NULL as that would stop the processing + * of directory entries completely. + */ + entry->d_name[0] = '?'; + entry->d_name[1] = '\0'; + entry->d_namlen = 1; + entry->d_type = DT_UNKNOWN; + entry->d_ino = 0; + entry->d_off = -1; + entry->d_reclen = 0; + + } + + /* Return pointer to directory entry */ + *result = entry; + + } else { + + /* No more directory entries */ + *result = NULL; + + } + + return /*OK*/0; +} + +/* + * Close directory stream. + */ +static int +closedir( + DIR *dirp) +{ + int ok; + if (dirp) { + + /* Close wide-character directory stream */ + ok = _wclosedir (dirp->wdirp); + dirp->wdirp = NULL; + + /* Release multi-byte character version */ + free (dirp); + + } else { + + /* Invalid directory stream */ + dirent_set_errno (EBADF); + ok = /*failure*/-1; + + } + return ok; +} + +/* + * Rewind directory stream to beginning. + */ +static void +rewinddir( + DIR* dirp) +{ + /* Rewind wide-character string directory stream */ + _wrewinddir (dirp->wdirp); +} + +/* + * Scan directory for entries. + */ +static int +scandir( + const char *dirname, + struct dirent ***namelist, + int (*filter)(const struct dirent*), + int (*compare)(const struct dirent**, const struct dirent**)) +{ + struct dirent **files = NULL; + size_t size = 0; + size_t allocated = 0; + const size_t init_size = 1; + DIR *dir = NULL; + struct dirent *entry; + struct dirent *tmp = NULL; + size_t i; + int result = 0; + + /* Open directory stream */ + dir = opendir (dirname); + if (dir) { + + /* Read directory entries to memory */ + while (1) { + + /* Enlarge pointer table to make room for another pointer */ + if (size >= allocated) { + void *p; + size_t num_entries; + + /* Compute number of entries in the enlarged pointer table */ + if (size < init_size) { + /* Allocate initial pointer table */ + num_entries = init_size; + } else { + /* Double the size */ + num_entries = size * 2; + } + + /* Allocate first pointer table or enlarge existing table */ + p = realloc (files, sizeof (void*) * num_entries); + if (p != NULL) { + /* Got the memory */ + files = (dirent**) p; + allocated = num_entries; + } else { + /* Out of memory */ + result = -1; + break; + } + + } + + /* Allocate room for temporary directory entry */ + if (tmp == NULL) { + tmp = (struct dirent*) malloc (sizeof (struct dirent)); + if (tmp == NULL) { + /* Cannot allocate temporary directory entry */ + result = -1; + break; + } + } + + /* Read directory entry to temporary area */ + if (readdir_r (dir, tmp, &entry) == /*OK*/0) { + + /* Did we get an entry? */ + if (entry != NULL) { + int pass; + + /* Determine whether to include the entry in result */ + if (filter) { + /* Let the filter function decide */ + pass = filter (tmp); + } else { + /* No filter function, include everything */ + pass = 1; + } + + if (pass) { + /* Store the temporary entry to pointer table */ + files[size++] = tmp; + tmp = NULL; + + /* Keep up with the number of files */ + result++; + } + + } else { + + /* + * End of directory stream reached => sort entries and + * exit. + */ + qsort (files, size, sizeof (void*), + (int (*) (const void*, const void*)) compare); + break; + + } + + } else { + /* Error reading directory entry */ + result = /*Error*/ -1; + break; + } + + } + + } else { + /* Cannot open directory */ + result = /*Error*/ -1; + } + + /* Release temporary directory entry */ + free (tmp); + + /* Release allocated memory on error */ + if (result < 0) { + for (i = 0; i < size; i++) { + free (files[i]); + } + free (files); + files = NULL; + } + + /* Close directory stream */ + if (dir) { + closedir (dir); + } + + /* Pass pointer table to caller */ + if (namelist) { + *namelist = files; + } + return result; +} + +/* Alphabetical sorting */ +static int +alphasort( + const struct dirent **a, const struct dirent **b) +{ + return strcoll ((*a)->d_name, (*b)->d_name); +} + +/* Sort versions */ +static int +versionsort( + const struct dirent **a, const struct dirent **b) +{ + /* FIXME: implement strverscmp and use that */ + return alphasort (a, b); +} + +/* Convert multi-byte string to wide character string */ +static int +dirent_mbstowcs_s( + size_t *pReturnValue, + wchar_t *wcstr, + size_t sizeInWords, + const char *mbstr, + size_t count) +{ + int error; + +#if defined(_MSC_VER) && _MSC_VER >= 1400 + + /* Microsoft Visual Studio 2005 or later */ + error = mbstowcs_s (pReturnValue, wcstr, sizeInWords, mbstr, count); + +#else + + /* Older Visual Studio or non-Microsoft compiler */ + size_t n; + + /* Convert to wide-character string (or count characters) */ + n = mbstowcs (wcstr, mbstr, sizeInWords); + if (!wcstr || n < count) { + + /* Zero-terminate output buffer */ + if (wcstr && sizeInWords) { + if (n >= sizeInWords) { + n = sizeInWords - 1; + } + wcstr[n] = 0; + } + + /* Length of resulting multi-byte string WITH zero terminator */ + if (pReturnValue) { + *pReturnValue = n + 1; + } + + /* Success */ + error = 0; + + } else { + + /* Could not convert string */ + error = 1; + + } + +#endif + return error; +} + +/* Convert wide-character string to multi-byte string */ +static int +dirent_wcstombs_s( + size_t *pReturnValue, + char *mbstr, + size_t sizeInBytes, /* max size of mbstr */ + const wchar_t *wcstr, + size_t count) +{ + int error; + +#if defined(_MSC_VER) && _MSC_VER >= 1400 + + /* Microsoft Visual Studio 2005 or later */ + error = wcstombs_s (pReturnValue, mbstr, sizeInBytes, wcstr, count); + +#else + + /* Older Visual Studio or non-Microsoft compiler */ + size_t n; + + /* Convert to multi-byte string (or count the number of bytes needed) */ + n = wcstombs (mbstr, wcstr, sizeInBytes); + if (!mbstr || n < count) { + + /* Zero-terminate output buffer */ + if (mbstr && sizeInBytes) { + if (n >= sizeInBytes) { + n = sizeInBytes - 1; + } + mbstr[n] = '\0'; + } + + /* Length of resulting multi-bytes string WITH zero-terminator */ + if (pReturnValue) { + *pReturnValue = n + 1; + } + + /* Success */ + error = 0; + + } else { + + /* Cannot convert string */ + error = 1; + + } + +#endif + return error; +} + +/* Set errno variable */ +static void +dirent_set_errno( + int error) +{ +#if defined(_MSC_VER) && _MSC_VER >= 1400 + + /* Microsoft Visual Studio 2005 and later */ + _set_errno (error); + +#else + + /* Non-Microsoft compiler or older Microsoft compiler */ + errno = error; + +#endif +} + + +#ifdef __cplusplus +} +#endif +#endif /*DIRENT_H*/ + diff --git a/yolov9/yolov9_trt.py b/yolov9/yolov9_trt.py new file mode 100644 index 0000000..15a7737 --- /dev/null +++ b/yolov9/yolov9_trt.py @@ -0,0 +1,454 @@ +""" +An example that uses TensorRT's Python api to make inferences. +""" +import ctypes +import os +import shutil +import random +import sys +import threading +import time +import cv2 +import numpy as np +import pycuda.autoinit +import pycuda.driver as cuda +import tensorrt as trt + +CONF_THRESH = 0.5 +IOU_THRESHOLD = 0.4 + + +def get_img_path_batches(batch_size, img_dir): + ret = [] + batch = [] + for root, dirs, files in os.walk(img_dir): + for name in files: + if len(batch) == batch_size: + ret.append(batch) + batch = [] + batch.append(os.path.join(root, name)) + if len(batch) > 0: + ret.append(batch) + return ret + + +def plot_one_box(x, img, color=None, label=None, line_thickness=None): + """ + description: Plots one bounding box on image img, + this function comes from yolov9 project. + param: + x: a box likes [x1,y1,x2,y2] + img: a opencv image object + color: color to draw rectangle, such as (0,255,0) + label: str + line_thickness: int + return: + no return + + """ + tl = ( + line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 + ) # line/font thickness + color = color or [random.randint(0, 255) for _ in range(3)] + c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3])) + cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA) + if label: + tf = max(tl - 1, 1) # font thickness + t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] + c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 + cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled + cv2.putText( + img, + label, + (c1[0], c1[1] - 2), + 0, + tl / 3, + [225, 255, 255], + thickness=tf, + lineType=cv2.LINE_AA, + ) + + +class yolov9TRT(object): + """ + description: A yolov9 class that warps TensorRT ops, preprocess and postprocess ops. + """ + + def __init__(self, engine_file_path): + # Create a Context on this device, + self.ctx = cuda.Device(0).make_context() + stream = cuda.Stream() + TRT_LOGGER = trt.Logger(trt.Logger.INFO) + runtime = trt.Runtime(TRT_LOGGER) + + # Deserialize the engine from file + with open(engine_file_path, "rb") as f: + engine = runtime.deserialize_cuda_engine(f.read()) + context = engine.create_execution_context() + + host_inputs = [] + cuda_inputs = [] + host_outputs = [] + cuda_outputs = [] + bindings = [] + + for binding in engine: + print('bingding:', binding, engine.get_binding_shape(binding)) + size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size + dtype = trt.nptype(engine.get_binding_dtype(binding)) + # Allocate host and device buffers + host_mem = cuda.pagelocked_empty(size, dtype) + cuda_mem = cuda.mem_alloc(host_mem.nbytes) + # Append the device buffer to device bindings. + bindings.append(int(cuda_mem)) + # Append to the appropriate list. + if engine.binding_is_input(binding): + self.input_w = engine.get_binding_shape(binding)[-1] + self.input_h = engine.get_binding_shape(binding)[-2] + host_inputs.append(host_mem) + cuda_inputs.append(cuda_mem) + else: + host_outputs.append(host_mem) + cuda_outputs.append(cuda_mem) + + # Store + self.stream = stream + self.context = context + self.engine = engine + self.host_inputs = host_inputs + self.cuda_inputs = cuda_inputs + self.host_outputs = host_outputs + self.cuda_outputs = cuda_outputs + self.bindings = bindings + self.batch_size = engine.max_batch_size + + def infer(self, raw_image_generator): + threading.Thread.__init__(self) + # Make self the active context, pushing it on top of the context stack. + self.ctx.push() + # Restore + stream = self.stream + context = self.context + engine = self.engine + host_inputs = self.host_inputs + cuda_inputs = self.cuda_inputs + host_outputs = self.host_outputs + cuda_outputs = self.cuda_outputs + bindings = self.bindings + # Do image preprocess + batch_image_raw = [] + batch_origin_h = [] + batch_origin_w = [] + batch_input_image = np.empty(shape=[self.batch_size, 3, self.input_h, self.input_w]) + for i, image_raw in enumerate(raw_image_generator): + input_image, image_raw, origin_h, origin_w = self.preprocess_image(image_raw) + batch_image_raw.append(image_raw) + batch_origin_h.append(origin_h) + batch_origin_w.append(origin_w) + np.copyto(batch_input_image[i], input_image) + batch_input_image = np.ascontiguousarray(batch_input_image) + + # Copy input image to host buffer + np.copyto(host_inputs[0], batch_input_image.ravel()) + start = time.time() + # Transfer input data to the GPU. + cuda.memcpy_htod_async(cuda_inputs[0], host_inputs[0], stream) + # Run inference. + context.execute_async(batch_size=self.batch_size, bindings=bindings, stream_handle=stream.handle) + # Transfer predictions back from the GPU. + cuda.memcpy_dtoh_async(host_outputs[0], cuda_outputs[0], stream) + # Synchronize the stream + stream.synchronize() + end = time.time() + # Remove any context from the top of the context stack, deactivating it. + self.ctx.pop() + # Here we use the first row of output in that batch_size = 1 + output = host_outputs[0] + # Do postprocess + for i in range(self.batch_size): + result_boxes, result_scores, result_classid = self.post_process( + output[i * 38001: (i + 1) * 38001], batch_origin_h[i], batch_origin_w[i] + ) + # Draw rectangles and labels on the original image + for j in range(len(result_boxes)): + box = result_boxes[j] + plot_one_box( + box, + batch_image_raw[i], + label="{}:{:.2f}".format( + categories[int(result_classid[j])], result_scores[j] + ), + ) + return batch_image_raw, end - start + + def destroy(self): + # Remove any context from the top of the context stack, deactivating it. + self.ctx.pop() + + def get_raw_image(self, image_path_batch): + """ + description: Read an image from image path + """ + for img_path in image_path_batch: + yield cv2.imread(img_path) + + def get_raw_image_zeros(self, image_path_batch=None): + """ + description: Ready data for warmup + """ + for _ in range(self.batch_size): + yield np.zeros([self.input_h, self.input_w, 3], dtype=np.uint8) + + def preprocess_image(self, raw_bgr_image): + """ + description: Convert BGR image to RGB, + resize and pad it to target size, normalize to [0,1], + transform to NCHW format. + param: + input_image_path: str, image path + return: + image: the processed image + image_raw: the original image + h: original height + w: original width + """ + image_raw = raw_bgr_image + h, w, c = image_raw.shape + image = cv2.cvtColor(image_raw, cv2.COLOR_BGR2RGB) + # Calculate widht and height and paddings + r_w = self.input_w / w + r_h = self.input_h / h + if r_h > r_w: + tw = self.input_w + th = int(r_w * h) + tx1 = tx2 = 0 + ty1 = int((self.input_h - th) / 2) + ty2 = self.input_h - th - ty1 + else: + tw = int(r_h * w) + th = self.input_h + tx1 = int((self.input_w - tw) / 2) + tx2 = self.input_w - tw - tx1 + ty1 = ty2 = 0 + # Resize the image with long side while maintaining ratio + image = cv2.resize(image, (tw, th)) + # Pad the short side with (128,128,128) + image = cv2.copyMakeBorder( + image, ty1, ty2, tx1, tx2, cv2.BORDER_CONSTANT, None, (128, 128, 128) + ) + image = image.astype(np.float32) + # Normalize to [0,1] + image /= 255.0 + # HWC to CHW format: + image = np.transpose(image, [2, 0, 1]) + # CHW to NCHW format + image = np.expand_dims(image, axis=0) + # Convert the image to row-major order, also known as "C order": + image = np.ascontiguousarray(image) + return image, image_raw, h, w + + def xywh2xyxy(self, origin_h, origin_w, x): + """ + description: Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right + param: + origin_h: height of original image + origin_w: width of original image + x: A boxes numpy, each row is a box [center_x, center_y, w, h] + return: + y: A boxes numpy, each row is a box [x1, y1, x2, y2] + """ + y = np.zeros_like(x) + r_w = self.input_w / origin_w + r_h = self.input_h / origin_h + if r_h > r_w: + y[:, 0] = x[:, 0] + y[:, 2] = x[:, 2] + y[:, 1] = x[:, 1] - (self.input_h - r_w * origin_h) / 2 + y[:, 3] = x[:, 3] - (self.input_h - r_w * origin_h) / 2 + y /= r_w + else: + y[:, 0] = x[:, 0] - (self.input_w - r_h * origin_w) / 2 + y[:, 2] = x[:, 2] - (self.input_w - r_h * origin_w) / 2 + y[:, 1] = x[:, 1] + y[:, 3] = x[:, 3] + y /= r_h + + return y + def post_process(self, output, origin_h, origin_w): + """ + description: postprocess the prediction + param: + output: A numpy likes [num_boxes,cx,cy,w,h,conf,cls_id, cx,cy,w,h,conf,cls_id, ...] + origin_h: height of original image + origin_w: width of original image + return: + result_boxes: finally boxes, a boxes numpy, each row is a box [x1, y1, x2, y2] + result_scores: finally scores, a numpy, each element is the score correspoing to box + result_classid: finally classid, a numpy, each element is the classid correspoing to box + """ + # Get the num of boxes detected + num = int(output[0]) + # Reshape to a two dimentional ndarray + pred = np.reshape(output[1:], (-1, 38))[:num, :] + # Do nms + boxes = self.non_max_suppression(pred, origin_h, origin_w, conf_thres=CONF_THRESH, nms_thres=IOU_THRESHOLD) + result_boxes = boxes[:, :4] if len(boxes) else np.array([]) + result_scores = boxes[:, 4] if len(boxes) else np.array([]) + result_classid = boxes[:, 5] if len(boxes) else np.array([]) + return result_boxes, result_scores, result_classid + + def bbox_iou(self, box1, box2, x1y1x2y2=True): + """ + description: compute the IoU of two bounding boxes + param: + box1: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h)) + box2: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h)) + x1y1x2y2: select the coordinate format + return: + iou: computed iou + """ + if not x1y1x2y2: + # Transform from center and width to exact coordinates + b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2 + b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2 + b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2 + b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2 + else: + # Get the coordinates of bounding boxes + b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3] + b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3] + + # Get the coordinates of the intersection rectangle + inter_rect_x1 = np.maximum(b1_x1, b2_x1) + inter_rect_y1 = np.maximum(b1_y1, b2_y1) + inter_rect_x2 = np.minimum(b1_x2, b2_x2) + inter_rect_y2 = np.minimum(b1_y2, b2_y2) + # Intersection area + inter_area = np.clip(inter_rect_x2 - inter_rect_x1 + 1, 0, None) * \ + np.clip(inter_rect_y2 - inter_rect_y1 + 1, 0, None) + # Union Area + b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1) + b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1) + + iou = inter_area / (b1_area + b2_area - inter_area + 1e-16) + + return iou + + def non_max_suppression(self, prediction, origin_h, origin_w, conf_thres=0.5, nms_thres=0.4): + """ + description: Removes detections with lower object confidence score than 'conf_thres' and performs + Non-Maximum Suppression to further filter detections. + param: + prediction: detections, (x1, y1, x2, y2, conf, cls_id) + origin_h: original image height + origin_w: original image width + conf_thres: a confidence threshold to filter detections + nms_thres: a iou threshold to filter detections + return: + boxes: output after nms with the shape (x1, y1, x2, y2, conf, cls_id) + """ + # Get the boxes that score > CONF_THRESH + boxes = prediction[prediction[:, 4] >= conf_thres] + # Trandform bbox from [center_x, center_y, w, h] to [x1, y1, x2, y2] + boxes[:, :4] = self.xywh2xyxy(origin_h, origin_w, boxes[:, :4]) + # clip the coordinates + boxes[:, 0] = np.clip(boxes[:, 0], 0, origin_w - 1) + boxes[:, 2] = np.clip(boxes[:, 2], 0, origin_w - 1) + boxes[:, 1] = np.clip(boxes[:, 1], 0, origin_h - 1) + boxes[:, 3] = np.clip(boxes[:, 3], 0, origin_h - 1) + # Object confidence + confs = boxes[:, 4] + # Sort by the confs + boxes = boxes[np.argsort(-confs)] + # Perform non-maximum suppression + keep_boxes = [] + while boxes.shape[0]: + large_overlap = self.bbox_iou(np.expand_dims(boxes[0, :4], 0), boxes[:, :4]) > nms_thres + label_match = boxes[0, -1] == boxes[:, -1] + # Indices of boxes with lower confidence scores, large IOUs and matching labels + invalid = large_overlap & label_match + keep_boxes += [boxes[0]] + boxes = boxes[~invalid] + boxes = np.stack(keep_boxes, 0) if len(keep_boxes) else np.array([]) + return boxes + + +class inferThread(threading.Thread): + def __init__(self, yolov9_wrapper, image_path_batch): + threading.Thread.__init__(self) + self.yolov9_wrapper = yolov9_wrapper + self.image_path_batch = image_path_batch + + def run(self): + batch_image_raw, use_time = self.yolov9_wrapper.infer(self.yolov9_wrapper.get_raw_image(self.image_path_batch)) + for i, img_path in enumerate(self.image_path_batch): + parent, filename = os.path.split(img_path) + save_name = os.path.join('output', filename) + # Save image + cv2.imwrite(save_name, batch_image_raw[i]) + print('input->{}, time->{:.2f}ms, saving into output/'.format(self.image_path_batch, use_time * 1000)) + + +class warmUpThread(threading.Thread): + def __init__(self, yolov9_wrapper): + threading.Thread.__init__(self) + self.yolov9_wrapper = yolov9_wrapper + + def run(self): + batch_image_raw, use_time = self.yolov9_wrapper.infer(self.yolov9_wrapper.get_raw_image_zeros()) + print('warm_up->{}, time->{:.2f}ms'.format(batch_image_raw[0].shape, use_time * 1000)) + + +if __name__ == "__main__": + # load custom plugin and engine + PLUGIN_LIBRARY = "build/libmyplugins.so" + engine_file_path = "yolov9-c.engine" + + if len(sys.argv) > 1: + engine_file_path = sys.argv[1] + if len(sys.argv) > 2: + PLUGIN_LIBRARY = sys.argv[2] + + ctypes.CDLL(PLUGIN_LIBRARY) + + # load coco labels + + categories = ["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"] + + if os.path.exists('output/'): + shutil.rmtree('output/') + os.makedirs('output/') + # a yolov9TRT instance + yolov9_wrapper = yolov9TRT(engine_file_path) + try: + print('batch size is', yolov9_wrapper.batch_size) + + image_dir = "images/" + image_path_batches = get_img_path_batches(yolov9_wrapper.batch_size, image_dir) + + for i in range(10): + # create a new thread to do warm_up + thread1 = warmUpThread(yolov9_wrapper) + thread1.start() + thread1.join() + for batch in image_path_batches: + # create a new thread to do inference + thread1 = inferThread(yolov9_wrapper, batch) + thread1.start() + thread1.join() + finally: + # destroy the instance + yolov9_wrapper.destroy()