tensorrt-yolov9 (#1449)
* tensorrt-yolov9 * format change * change format * Update block.cpp * format: add space * update block.cpp: add space * update config.h --------- Co-authored-by: Wang Xinyu <wangxinyu_es@163.com>
This commit is contained in:
WuxinrongY
GitHub
parent
e585b15c36
commit
e73bffcd25
56
yolov9/CMakeLists.txt
Normal file
56
yolov9/CMakeLists.txt
Normal file
@ -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})
|
||||
|
||||
106
yolov9/README.md
Normal file
106
yolov9/README.md
Normal file
@ -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
|
||||
|
||||
<p align="center">
|
||||
<img src="https://user-images.githubusercontent.com/15235574/78247927-4d9fac00-751e-11ea-8b1b-704a0aeb3fcf.jpg" height="360px;">
|
||||
</p>
|
||||
|
||||
## More Information
|
||||
|
||||
See the readme in [home page.](https://github.com/wang-xinyu/tensorrtx)
|
||||
|
||||
|
||||
211
yolov9/demo.cpp
Normal file
211
yolov9/demo.cpp
Normal file
@ -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 <chrono>
|
||||
#include <fstream>
|
||||
|
||||
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<const char*>(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<std::string> 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<cv::Mat> img_batch;
|
||||
std::vector<std::string> 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<std::chrono::microseconds>(end - start).count() / 1000.0 / speed_test_iter << "ms" << std::endl;
|
||||
|
||||
// NMS
|
||||
std::vector<std::vector<Detection>> 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;
|
||||
}
|
||||
|
||||
70
yolov9/gen_wts.py
Normal file
70
yolov9/gen_wts.py
Normal file
@ -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')
|
||||
1
yolov9/images
Symbolic link
1
yolov9/images
Symbolic link
@ -0,0 +1 @@
|
||||
../yolov3-spp/samples
|
||||
31
yolov9/include/block.h
Normal file
31
yolov9/include/block.h
Normal file
@ -0,0 +1,31 @@
|
||||
#include "config.h"
|
||||
#include "yololayer.h"
|
||||
|
||||
#include <iostream>
|
||||
#include <fstream>
|
||||
#include <map>
|
||||
#include <cassert>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
|
||||
using namespace nvinfer1;
|
||||
|
||||
// TensorRT weight files have a simple space delimited format:
|
||||
// [type] [size] <data x size in hex>
|
||||
void PrintDim(const ILayer* layer, std::string log = "");
|
||||
std::map<std::string, Weights> 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<std::string, Weights>& weightMap, ITensor& input, int c_, int c2, std::string lname);
|
||||
std::vector<std::vector<float>> getAnchors(std::map<std::string, Weights>& weightMap, std::string lname);
|
||||
// ----------------------------------------------------------------
|
||||
nvinfer1::ILayer* convBnSiLU(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights>& weightMap, nvinfer1::ITensor& input, int ch, int k, int s, int p, std::string lname, int g=1);
|
||||
ILayer* RepNCSPELAN4(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, int c3, int c4, int c5, std::string lname);
|
||||
ILayer* ADown(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c2, std::string lname);
|
||||
std::vector<ILayer*> CBLinear(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::vector<int> c2s, int k, int s, int p, int g, std::string lname);
|
||||
ILayer* CBFuse(INetworkDefinition *network, std::vector<std::vector<ILayer*>> input, std::vector<int> idx, std::vector<int> strides);
|
||||
ILayer* SPPELAN(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, int c3, std::string lname);
|
||||
std::vector<IConcatenationLayer*> DualDDetect(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, std::vector<ILayer*> dets, int cls, std::vector<int> ch, std::string lname);
|
||||
nvinfer1::IPluginV2Layer* addYoLoLayer(nvinfer1::INetworkDefinition *network, std::vector<nvinfer1::IConcatenationLayer*> dets, bool is_segmentation);
|
||||
nvinfer1::IShuffleLayer* DFL(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights> weightMap, nvinfer1::ITensor& input, int ch, int k, int s, int p, std::string lname);
|
||||
nvinfer1::ILayer* convBnNoAct(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights>& weightMap, nvinfer1::ITensor& input, int ch, int k, int s, int p, std::string lname, int g);
|
||||
36
yolov9/include/calibrator.h
Normal file
36
yolov9/include/calibrator.h
Normal file
@ -0,0 +1,36 @@
|
||||
#pragma once
|
||||
|
||||
#include "macros.h"
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
//! \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<std::string> img_files_;
|
||||
size_t input_count_;
|
||||
std::string calib_table_name_;
|
||||
const char* input_blob_name_;
|
||||
bool read_cache_;
|
||||
void* device_input_;
|
||||
std::vector<char> calib_cache_;
|
||||
};
|
||||
|
||||
59
yolov9/include/config.h
Normal file
59
yolov9/include/config.h
Normal file
@ -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;
|
||||
|
||||
18
yolov9/include/cuda_utils.h
Normal file
18
yolov9/include/cuda_utils.h
Normal file
@ -0,0 +1,18 @@
|
||||
#ifndef TRTX_CUDA_UTILS_H_
|
||||
#define TRTX_CUDA_UTILS_H_
|
||||
|
||||
#include <cuda_runtime_api.h>
|
||||
|
||||
#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_
|
||||
|
||||
504
yolov9/include/logging.h
Normal file
504
yolov9/include/logging.h
Normal file
@ -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 <cassert>
|
||||
#include <ctime>
|
||||
#include <iomanip>
|
||||
#include <iostream>
|
||||
#include <ostream>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#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
|
||||
29
yolov9/include/macros.h
Normal file
29
yolov9/include/macros.h
Normal file
@ -0,0 +1,29 @@
|
||||
#ifndef __MACROS_H
|
||||
#define __MACROS_H
|
||||
|
||||
#include <NvInfer.h>
|
||||
|
||||
#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
|
||||
6
yolov9/include/model.h
Normal file
6
yolov9/include/model.h
Normal file
@ -0,0 +1,6 @@
|
||||
#pragma once
|
||||
|
||||
#include <NvInfer.h>
|
||||
#include <string>
|
||||
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);
|
||||
20
yolov9/include/postprocess.h
Normal file
20
yolov9/include/postprocess.h
Normal file
@ -0,0 +1,20 @@
|
||||
#pragma once
|
||||
|
||||
#include "types.h"
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <cuda_runtime.h>
|
||||
cv::Rect get_rect(cv::Mat& img, float bbox[4]);
|
||||
|
||||
void nms(std::vector<Detection>& res, float *output, float conf_thresh, float nms_thresh = 0.5);
|
||||
|
||||
void batch_nms(std::vector<std::vector<Detection>>& batch_res, float *output, int batch_size, int output_size, float conf_thresh, float nms_thresh = 0.5);
|
||||
|
||||
void draw_bbox(std::vector<cv::Mat>& img_batch, std::vector<std::vector<Detection>>& res_batch);
|
||||
|
||||
std::vector<cv::Mat> process_mask(const float* proto, int proto_size, std::vector<Detection>& dets);
|
||||
|
||||
void draw_mask_bbox(cv::Mat& img, std::vector<Detection>& dets, std::vector<cv::Mat>& masks, std::unordered_map<int, std::string>& 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<std::vector<Detection>> &res_batch, const float* decode_ptr_host, int batch_size, int bbox_element, const std::vector<cv::Mat>& img_batch);
|
||||
15
yolov9/include/preprocess.h
Normal file
15
yolov9/include/preprocess.h
Normal file
@ -0,0 +1,15 @@
|
||||
#pragma once
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cstdint>
|
||||
#include <opencv2/opencv.hpp>
|
||||
|
||||
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<cv::Mat>& img_batch,
|
||||
float* dst, int dst_width, int dst_height,
|
||||
cudaStream_t stream);
|
||||
|
||||
17
yolov9/include/types.h
Normal file
17
yolov9/include/types.h
Normal file
@ -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
|
||||
70
yolov9/include/utils.h
Normal file
70
yolov9/include/utils.h
Normal file
@ -0,0 +1,70 @@
|
||||
#pragma once
|
||||
|
||||
#include <dirent.h>
|
||||
#include <fstream>
|
||||
#include <unordered_map>
|
||||
#include <string>
|
||||
#include <sstream>
|
||||
#include <vector>
|
||||
#include <cstring>
|
||||
|
||||
static inline int read_files_in_dir(const char* p_dir_name, std::vector<std::string>& 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<int, std::string>& 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;
|
||||
}
|
||||
|
||||
242
yolov9/plugin/yololayer.cu
Normal file
242
yolov9/plugin/yololayer.cu
Normal file
@ -0,0 +1,242 @@
|
||||
#include "yololayer.h"
|
||||
#include "types.h"
|
||||
#include <assert.h>
|
||||
#include <math.h>
|
||||
#include "cuda_utils.h"
|
||||
#include <vector>
|
||||
#include <iostream>
|
||||
|
||||
namespace Tn {
|
||||
template<typename T>
|
||||
void write(char*& buffer, const T& val) {
|
||||
*reinterpret_cast<T*>(buffer) = val;
|
||||
buffer += sizeof(T);
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
void read(const char*& buffer, T& val) {
|
||||
val = *reinterpret_cast<const T*>(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<const char*>(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<char*>(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<PluginField> 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
|
||||
102
yolov9/plugin/yololayer.h
Normal file
102
yolov9/plugin/yololayer.h
Normal file
@ -0,0 +1,102 @@
|
||||
#pragma once
|
||||
#include "macros.h"
|
||||
#include "NvInfer.h"
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#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<PluginField> mPluginAttributes;
|
||||
};
|
||||
REGISTER_TENSORRT_PLUGIN(YoloPluginCreator);
|
||||
} // namespace nvinfer1
|
||||
|
||||
385
yolov9/src/block.cpp
Normal file
385
yolov9/src/block.cpp
Normal file
@ -0,0 +1,385 @@
|
||||
#include "block.h"
|
||||
#include "calibrator.h"
|
||||
#include "config.h"
|
||||
#include "yololayer.h"
|
||||
|
||||
#include <iostream>
|
||||
#include <fstream>
|
||||
#include <map>
|
||||
#include <cassert>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
#include <algorithm>
|
||||
|
||||
using namespace nvinfer1;
|
||||
|
||||
// TensorRT weight files have a simple space delimited format:
|
||||
// [type] [size] <data x size in hex>
|
||||
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<std::string, Weights> loadWeights(const std::string file) {
|
||||
std::cout << "Loading weights: " << file << std::endl;
|
||||
std::map<std::string, Weights> 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<uint32_t*>(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<int>(r, 1);
|
||||
}
|
||||
static nvinfer1::IScaleLayer* addBatchNorm2d(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights> 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<float*>(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<float*>(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<float*>(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<std::string, nvinfer1::Weights>& 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<std::string, nvinfer1::Weights>& 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<std::vector<float>> getAnchors(std::map<std::string, Weights>& weightMap, std::string lname) {
|
||||
std::vector<std::vector<float>> 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<float> anchor(p, p + anchor_len);
|
||||
anchors.push_back(anchor);
|
||||
}
|
||||
return anchors;
|
||||
}
|
||||
|
||||
ILayer* RepConvN(INetworkDefinition *network, std::map<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<ILayer*> CBLinear(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::vector<int> 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<ILayer*> 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<std::vector<ILayer*>> input, std::vector<int> idx, std::vector<int> 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, nvinfer1::Weights> 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<nvinfer1::IConcatenationLayer*> 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<nvinfer1::ITensor*> 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<IConcatenationLayer*> DualDDetect(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, std::vector<ILayer*> dets, int cls, std::vector<int> 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<ILayer*> bboxlayers;
|
||||
std::vector<ILayer*> 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<IConcatenationLayer*> 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;
|
||||
}
|
||||
97
yolov9/src/calibrator.cpp
Normal file
97
yolov9/src/calibrator.cpp
Normal file
@ -0,0 +1,97 @@
|
||||
#include "calibrator.h"
|
||||
#include "cuda_utils.h"
|
||||
#include "utils.h"
|
||||
|
||||
#include <iostream>
|
||||
#include <iterator>
|
||||
#include <fstream>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <opencv2/dnn/dnn.hpp>
|
||||
|
||||
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<cv::Mat> 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<float>(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<char>(input), std::istream_iterator<char>(), 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<const char*>(cache), length);
|
||||
}
|
||||
|
||||
413
yolov9/src/model.cpp
Normal file
413
yolov9/src/model.cpp
Normal file
@ -0,0 +1,413 @@
|
||||
#include "model.h"
|
||||
#include "calibrator.h"
|
||||
#include "config.h"
|
||||
#include "yololayer.h"
|
||||
#include "block.h"
|
||||
#include <iostream>
|
||||
#include <fstream>
|
||||
#include <map>
|
||||
#include <cassert>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
|
||||
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<std::string, Weights> 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<ILayer*>{ 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<ILayer*>{ 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<ILayer*>{ 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<ILayer*>{ 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<ILayer*>{ 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<ILayer*>{RepNCSPELAN_42, RepNCSPELAN_45, RepNCSPELAN_48}, kNumClass, {256, 512, 512}, "model.49");
|
||||
auto dualddetect_49 = DualDDetect(network, weightMap, std::vector<ILayer*>{ 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<std::string, Weights> 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<ILayer*>{ 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<ILayer*>{ 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<ILayer*>{ 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<ILayer*>{ 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;
|
||||
|
||||
}
|
||||
228
yolov9/src/postprocess.cpp
Normal file
228
yolov9/src/postprocess.cpp
Normal file
@ -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<Detection>& res, float* output, float conf_thresh, float nms_thresh) {
|
||||
int det_size = sizeof(Detection) / sizeof(float);
|
||||
std::map<float, std::vector<Detection>> 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<Detection>());
|
||||
// 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<std::vector<Detection>>& 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<cv::Mat>& img_batch, std::vector<std::vector<Detection>>& 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<cv::Mat> process_mask(const float* proto, int proto_size, std::vector<Detection>& dets) {
|
||||
// std::vector<cv::Mat> 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<float>(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<Detection>& dets, std::vector<cv::Mat>& masks, std::unordered_map<int, std::string>& labels_map) {
|
||||
static std::vector<uint32_t> 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<float>(y, x);
|
||||
if (val <= 0.5) continue;
|
||||
img.at<cv::Vec3b>(y, x)[0] = img.at<cv::Vec3b>(y, x)[0] / 2 + bgr[0] / 2;
|
||||
img.at<cv::Vec3b>(y, x)[1] = img.at<cv::Vec3b>(y, x)[1] / 2 + bgr[1] / 2;
|
||||
img.at<cv::Vec3b>(y, x)[2] = img.at<cv::Vec3b>(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<Detection> &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<std::vector<Detection>> &res_batch, const float* decode_ptr_host, int batch_size, int bbox_element, const std::vector<cv::Mat>& img_batch) {
|
||||
res_batch.resize(batch_size);
|
||||
int count = static_cast<int>(*decode_ptr_host);
|
||||
count = count > kMaxNumOutputBbox ? kMaxNumOutputBbox : count;
|
||||
// std::min(count, kMaxNumOutputBbox);
|
||||
for (int i = 0; i < batch_size; i++) {
|
||||
auto& img = const_cast<cv::Mat&>(img_batch[i]);
|
||||
process_decode_ptr_host(res_batch[i], &decode_ptr_host[i * count], bbox_element, img, count);
|
||||
}
|
||||
}
|
||||
|
||||
91
yolov9/src/postprocess.cu
Normal file
91
yolov9/src/postprocess.cu
Normal file
@ -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);
|
||||
|
||||
}
|
||||
153
yolov9/src/preprocess.cu
Normal file
153
yolov9/src/preprocess.cu
Normal file
@ -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<float>(0), sizeof(d2s.value));
|
||||
|
||||
int jobs = dst_height * dst_width;
|
||||
int threads = 256;
|
||||
int blocks = ceil(jobs / (float)threads);
|
||||
|
||||
warpaffine_kernel<<<blocks, threads, 0, stream>>>(
|
||||
img_buffer_device, src_width * 3, src_width,
|
||||
src_height, dst, dst_width,
|
||||
dst_height, 128, d2s, jobs);
|
||||
}
|
||||
|
||||
void cuda_batch_preprocess(std::vector<cv::Mat>& 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));
|
||||
}
|
||||
|
||||
1167
yolov9/windows/dirent.h
Normal file
1167
yolov9/windows/dirent.h
Normal file
File diff suppressed because it is too large
Load Diff
454
yolov9/yolov9_trt.py
Normal file
454
yolov9/yolov9_trt.py
Normal file
@ -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()
|
||||
Loading…
Reference in New Issue
Block a user