yolov7 v0.1 (#1105)
* yolov7 v0.1 * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Update yololayer.h * Update yololayer.h * Delete yolov7.engine * Delete yolov7-tiny.engine * Delete yolov7/output directory * Delete dirent.h * Update README.md * Update README.md * Update README.md * Update README.md
This commit is contained in:
lindsayshuo
GitHub
parent
d988519e73
commit
64da26ef5d
55
yolov7/CMakeLists.txt
Normal file
55
yolov7/CMakeLists.txt
Normal file
@ -0,0 +1,55 @@
|
||||
cmake_minimum_required(VERSION 2.6)
|
||||
|
||||
project(yolov7)
|
||||
|
||||
add_definitions(-std=c++11)
|
||||
add_definitions(-DAPI_EXPORTS)
|
||||
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
|
||||
set(CMAKE_CXX_STANDARD 11)
|
||||
set(CMAKE_BUILD_TYPE Debug)
|
||||
|
||||
find_package(CUDA REQUIRED)
|
||||
|
||||
if(WIN32)
|
||||
enable_language(CUDA)
|
||||
endif(WIN32)
|
||||
|
||||
include_directories(${PROJECT_SOURCE_DIR}/include)
|
||||
# include and link dirs of cuda and tensorrt, you need adapt them if yours are different
|
||||
|
||||
|
||||
if (CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64")
|
||||
message("embed_platform on")
|
||||
include_directories(/usr/local/cuda/targets/aarch64-linux/include)
|
||||
link_directories(/usr/local/cuda/targets/aarch64-linux/lib)
|
||||
else()
|
||||
message("embed_platform off")
|
||||
# cuda
|
||||
include_directories(/usr/local/cuda/include)
|
||||
link_directories(/usr/local/cuda/lib64)
|
||||
|
||||
# tensorrt
|
||||
include_directories(/home/na/TensorRT-7.2.2.3/include)
|
||||
link_directories(/home/na/TensorRT-7.2.2.3/lib)
|
||||
endif()
|
||||
|
||||
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -g -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")
|
||||
cuda_add_library(myplugins SHARED yololayer.cu)
|
||||
target_link_libraries(myplugins nvinfer cudart)
|
||||
|
||||
find_package(OpenCV)
|
||||
include_directories(${OpenCV_INCLUDE_DIRS})
|
||||
|
||||
cuda_add_executable(yolov7 calibrator.cpp yolov7.cpp preprocess.cu)
|
||||
|
||||
target_link_libraries(yolov7 nvinfer)
|
||||
target_link_libraries(yolov7 cudart)
|
||||
target_link_libraries(yolov7 myplugins)
|
||||
target_link_libraries(yolov7 ${OpenCV_LIBS})
|
||||
|
||||
if(UNIX)
|
||||
add_definitions(-O2 -pthread)
|
||||
endif(UNIX)
|
||||
|
||||
|
||||
86
yolov7/README.md
Normal file
86
yolov7/README.md
Normal file
@ -0,0 +1,86 @@
|
||||
# yolov7
|
||||
|
||||
The Pytorch implementation is [WongKinYiu/yolov7](https://github.com/WongKinYiu/yolov7).
|
||||
|
||||
## environment
|
||||
|
||||
### ubuntu18.04 cuda11.1 tensorrt7.2.2.3
|
||||
|
||||
## Different versions of yolov7
|
||||
|
||||
Currently, we support yolov7 v0.1
|
||||
|
||||
- For yolov7 v0.1, download .pt from [yolov7 release v0.1](https://github.com/WongKinYiu/yolov7/releases/tag/v0.10), then follow how-to-run in current page.
|
||||
|
||||
|
||||
## Config
|
||||
|
||||
- Choose the model tiny/v7/x/d6/w6/e6/e6e from command line arguments.
|
||||
- Input shape defined in yololayer.h
|
||||
- Number of classes defined in yololayer.h, **DO NOT FORGET TO ADAPT THIS, If using your own model**
|
||||
- INT8/FP16/FP32 can be selected by the macro in yolov7.cpp, **INT8 need more steps, pls follow `How to Run` first and then go the `INT8 Quantization` below**
|
||||
- GPU id can be selected by the macro in yolov7.cpp
|
||||
- NMS thresh in yolov7.cpp
|
||||
- BBox confidence thresh in yolov7.cpp
|
||||
- Batch size in yolov7.cpp
|
||||
|
||||
## How to Run, yolov7-tiny as example
|
||||
|
||||
1. generate .wts from pytorch with .pt, or download .wts from model zoo
|
||||
|
||||
```
|
||||
// download https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-tiny.pt
|
||||
cp {tensorrtx}/yolov7/gen_wts.py {WongKinYiu}/yolov7
|
||||
cd {WongKinYiu}/yolov7
|
||||
python gen_wts.py -w yolov7.pt -o yolov7.wts
|
||||
// a file 'yolov7.wts' will be generated.
|
||||
```
|
||||
|
||||
|
||||
|
||||
2. build tensorrtx/yolov7 and run
|
||||
|
||||
```
|
||||
cd {tensorrtx}/yolov7/
|
||||
// update CLASS_NUM in yololayer.h if your model is trained on custom dataset
|
||||
mkdir build
|
||||
cd build
|
||||
cp {ultralytics}/yolov7/yolov7s.wts {tensorrtx}/yolov7/build
|
||||
cmake ..
|
||||
make
|
||||
sudo ./yolov7 -s [.wts] [.engine] [t/v7/x/w6/e6/d6/e6e gd gw] // serialize model to plan file
|
||||
sudo ./yolov7 -d [.engine] [image folder] // deserialize and run inference, the images in [image folder] will be processed.
|
||||
// For example yolov7s
|
||||
sudo ./yolov7 -s yolov7.wts yolov7.engine v7
|
||||
sudo ./yolov7 -d yolov7s.engine ../samples
|
||||
```
|
||||
|
||||
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 yolov7.engine and libmyplugins.so have been built
|
||||
python yolov7_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 yolov7/build
|
||||
|
||||
3. set the macro `USE_INT8` in yolov7.cpp and make
|
||||
|
||||
4. serialize the model and test
|
||||
|
||||
# Acknowledgment
|
||||
The code is modified on this basis: https://github.com/QIANXUNZDL123/tensorrtx-yolov7 ,
|
||||
Thanks for the donation for this tensorrt implementation from https://github.com/QIANXUNZDL123/tensorrtx-yolov7
|
||||
|
||||
## More Information
|
||||
|
||||
See the readme in [home page.](https://github.com/wang-xinyu/tensorrtx)
|
||||
|
||||
80
yolov7/calibrator.cpp
Normal file
80
yolov7/calibrator.cpp
Normal file
@ -0,0 +1,80 @@
|
||||
#include <iostream>
|
||||
#include <iterator>
|
||||
#include <fstream>
|
||||
#include <opencv2/dnn/dnn.hpp>
|
||||
#include "calibrator.h"
|
||||
#include "cuda_utils.h"
|
||||
#include "utils.h"
|
||||
|
||||
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);
|
||||
}
|
||||
|
||||
40
yolov7/calibrator.h
Normal file
40
yolov7/calibrator.h
Normal file
@ -0,0 +1,40 @@
|
||||
#ifndef ENTROPY_CALIBRATOR_H
|
||||
#define ENTROPY_CALIBRATOR_H
|
||||
|
||||
#include <NvInfer.h>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include "macros.h"
|
||||
|
||||
//! \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_;
|
||||
};
|
||||
|
||||
#endif // ENTROPY_CALIBRATOR_H
|
||||
678
yolov7/common.hpp
Normal file
678
yolov7/common.hpp
Normal file
@ -0,0 +1,678 @@
|
||||
#ifndef YOLOV7_COMMON_H_
|
||||
#define YOLOV7_COMMON_H_
|
||||
|
||||
#include <fstream>
|
||||
#include <map>
|
||||
#include <sstream>
|
||||
#include <vector>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include "NvInfer.h"
|
||||
#include "yololayer.h"
|
||||
using namespace nvinfer1;
|
||||
|
||||
|
||||
|
||||
cv::Rect get_rect(cv::Mat& img, float bbox[4]) {
|
||||
float l, r, t, b;
|
||||
float r_w = Yolo::INPUT_W / (img.cols * 1.0);
|
||||
float r_h = Yolo::INPUT_H / (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 - (Yolo::INPUT_H - r_w * img.rows) / 2;
|
||||
b = bbox[1] + bbox[3] / 2.f - (Yolo::INPUT_H - 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 - (Yolo::INPUT_W - r_h * img.cols) / 2;
|
||||
r = bbox[0] + bbox[2] / 2.f - (Yolo::INPUT_W - 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));
|
||||
}
|
||||
//=================IOU===============
|
||||
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);
|
||||
}
|
||||
|
||||
float giou(float lbox[4], float rbox[4]) {
|
||||
|
||||
float interBoxMax[] = {
|
||||
(std::min)(lbox[0] - lbox[2] / 2.f , rbox[0] - rbox[2] / 2.f), //left
|
||||
(std::max)(lbox[0] + lbox[2] / 2.f , rbox[0] + rbox[2] / 2.f), //right
|
||||
(std::min)(lbox[1] - lbox[3] / 2.f , rbox[1] - rbox[3] / 2.f), //top
|
||||
(std::max)(lbox[1] + lbox[3] / 2.f , rbox[1] + rbox[3] / 2.f), //bottom
|
||||
};
|
||||
|
||||
float interBoxMaxS = (interBoxMax[1] - interBoxMax[0]) * (interBoxMax[3] - interBoxMax[2]);
|
||||
|
||||
|
||||
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]);//
|
||||
float iou1 = interBoxS / (lbox[2] * lbox[3] + rbox[2] * rbox[3] - interBoxS);//正常计算IOU
|
||||
|
||||
float iou2 = (interBoxMaxS - (lbox[2] * lbox[3] + rbox[2] * rbox[3] - interBoxS)) / interBoxMaxS;
|
||||
|
||||
return iou1 - iou2;
|
||||
}
|
||||
|
||||
float diou(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]);//
|
||||
float iou = interBoxS / (lbox[2] * lbox[3] + rbox[2] * rbox[3] - interBoxS);
|
||||
|
||||
float interBoxMax[] = {
|
||||
(std::min)(lbox[0] - lbox[2] / 2.f , rbox[0] - rbox[2] / 2.f), //left
|
||||
(std::max)(lbox[0] + lbox[2] / 2.f , rbox[0] + rbox[2] / 2.f), //right
|
||||
(std::min)(lbox[1] - lbox[3] / 2.f , rbox[1] - rbox[3] / 2.f), //top
|
||||
(std::max)(lbox[1] + lbox[3] / 2.f , rbox[1] + rbox[3] / 2.f), //bottom
|
||||
};
|
||||
|
||||
|
||||
float d_center1 = (lbox[0] - rbox[0]) * (lbox[0] - rbox[0]) + (lbox[1] - rbox[1]) * (lbox[1] - rbox[1]);
|
||||
|
||||
float d_center2 = (interBoxMax[0] - interBoxMax[2]) * (interBoxMax[0] - interBoxMax[2])
|
||||
+ (interBoxMax[1] - interBoxMax[3]) * (interBoxMax[1] - interBoxMax[3]);
|
||||
|
||||
return iou - d_center1 / d_center2;
|
||||
}
|
||||
|
||||
float ciou(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]);//
|
||||
float iou = interBoxS / (lbox[2] * lbox[3] + rbox[2] * rbox[3] - interBoxS);
|
||||
|
||||
float interBoxMax[] = {
|
||||
(std::min)(lbox[0] - lbox[2] / 2.f , rbox[0] - rbox[2] / 2.f), //left
|
||||
(std::max)(lbox[0] + lbox[2] / 2.f , rbox[0] + rbox[2] / 2.f), //right
|
||||
(std::min)(lbox[1] - lbox[3] / 2.f , rbox[1] - rbox[3] / 2.f), //top
|
||||
(std::max)(lbox[1] + lbox[3] / 2.f , rbox[1] + rbox[3] / 2.f), //bottom
|
||||
};
|
||||
|
||||
|
||||
float d_center1 = (lbox[0] - rbox[0]) * (lbox[0] - rbox[0]) + (lbox[1] - rbox[1]) * (lbox[1] - rbox[1]);
|
||||
|
||||
float d_center2 = (interBoxMax[0] - interBoxMax[2]) * (interBoxMax[0] - interBoxMax[2])
|
||||
+ (interBoxMax[1] - interBoxMax[3]) * (interBoxMax[1] - interBoxMax[3]);
|
||||
|
||||
float v = 4 * (atan(lbox[2] / lbox[3]) * atan(rbox[2] / rbox[3])) * (atan(lbox[2] / lbox[3]) * atan(rbox[2] / rbox[3])) / (acos(-1) * acos(-1));
|
||||
float alpha = v / (1 - iou + v);
|
||||
float ciou = iou - d_center1 / d_center2 - alpha * v;
|
||||
return ciou;
|
||||
}
|
||||
|
||||
|
||||
bool cmp(const Yolo::Detection& a, const Yolo::Detection& b) {
|
||||
return a.conf > b.conf;
|
||||
}
|
||||
|
||||
void nms(std::vector<Yolo::Detection>& res, float *output, float conf_thresh, float nms_thresh = 0.5) {
|
||||
int det_size = sizeof(Yolo::Detection) / sizeof(float);
|
||||
std::map<float, std::vector<Yolo::Detection>> m;
|
||||
for (int i = 0; i < output[0] && i < Yolo::MAX_OUTPUT_BBOX_COUNT; i++) {
|
||||
if (output[1 + det_size * i + 4] <= conf_thresh) continue;
|
||||
Yolo::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<Yolo::Detection>());
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TensorRT weight files have a simple space delimited format:
|
||||
// [type] [size] <data x size in hex>
|
||||
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;
|
||||
}
|
||||
|
||||
|
||||
|
||||
IScaleLayer* addBatchNorm2d(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, 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);
|
||||
}
|
||||
Weights scale{ 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);
|
||||
}
|
||||
Weights shift{ DataType::kFLOAT, shval, len };
|
||||
|
||||
float* pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
|
||||
for (int i = 0; i < len; i++) {
|
||||
pval[i] = 1.0;
|
||||
}
|
||||
Weights power{ DataType::kFLOAT, pval, len };
|
||||
|
||||
weightMap[lname + ".scale"] = scale;
|
||||
weightMap[lname + ".shift"] = shift;
|
||||
weightMap[lname + ".power"] = power;
|
||||
IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
|
||||
assert(scale_1);
|
||||
return scale_1;
|
||||
}
|
||||
|
||||
|
||||
IElementWiseLayer * convBnSilu(INetworkDefinition * network, std::map<std::string, Weights>&weightMap, ITensor & input, int c2, int k, int s, int p, std::string lname) {
|
||||
Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
|
||||
|
||||
IConvolutionLayer* conv1 = network->addConvolutionNd(input, c2, DimsHW{ k, k }, weightMap[lname + ".conv.weight"], emptywts);
|
||||
assert(conv1);
|
||||
conv1->setName((lname + ".conv").c_str());
|
||||
conv1->setStrideNd(DimsHW{ s, s });
|
||||
conv1->setPaddingNd(DimsHW{ p, p });
|
||||
|
||||
|
||||
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn", 1e-5);
|
||||
|
||||
|
||||
// silu = x * sigmoid(x)
|
||||
IActivationLayer* sig1 = network->addActivation(*bn1->getOutput(0), ActivationType::kSIGMOID);
|
||||
assert(sig1);
|
||||
IElementWiseLayer* ew1 = network->addElementWise(*bn1->getOutput(0), *sig1->getOutput(0), ElementWiseOperation::kPROD);
|
||||
assert(ew1);
|
||||
return ew1;
|
||||
}
|
||||
ILayer* ReOrg(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int inch) {
|
||||
ISliceLayer* s1 = network->addSlice(input, Dims3{ 0, 0, 0 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 });
|
||||
ISliceLayer* s2 = network->addSlice(input, Dims3{ 0, 1, 0 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 });
|
||||
ISliceLayer* s3 = network->addSlice(input, Dims3{ 0, 0, 1 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 });
|
||||
ISliceLayer* s4 = network->addSlice(input, Dims3{ 0, 1, 1 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 });
|
||||
ITensor* inputTensors[] = { s1->getOutput(0), s2->getOutput(0), s3->getOutput(0), s4->getOutput(0) };
|
||||
auto cat = network->addConcatenation(inputTensors, 4);
|
||||
return cat;
|
||||
}
|
||||
ILayer* DownC(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, const std::string& lname)
|
||||
{
|
||||
int c_ = int(c2 * 0.5);
|
||||
IElementWiseLayer* cv1 = convBnSilu(network, weightMap, input, c1, 1, 1, 0, lname + ".cv1");
|
||||
IElementWiseLayer* cv2 = convBnSilu(network, weightMap, *cv1->getOutput(0), c_, 3, 2, 1, lname + ".cv2");
|
||||
|
||||
IPoolingLayer* m1 = network->addPoolingNd(input, PoolingType::kMAX, DimsHW{ 2, 2 });
|
||||
m1->setStrideNd(DimsHW{ 2, 2 });
|
||||
IElementWiseLayer* cv3 = convBnSilu(network, weightMap, *m1->getOutput(0), c_, 1, 1, 0, lname + ".cv3");
|
||||
|
||||
ITensor* input_tensors[] = { cv2->getOutput(0), cv3->getOutput(0) };
|
||||
IConcatenationLayer* concat = network->addConcatenation(input_tensors, 2);
|
||||
|
||||
return concat;
|
||||
|
||||
}
|
||||
|
||||
// SPPCSPC
|
||||
IElementWiseLayer* SPPCSPC(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int c2, const std::string& lname) {
|
||||
int c_ = int(2 * c2 * 0.5);
|
||||
IElementWiseLayer* cv1 = convBnSilu(network, weightMap, input, c_, 1, 1, 0, lname + ".cv1");
|
||||
IElementWiseLayer* cv2 = convBnSilu(network, weightMap, input, c_, 1, 1, 0, lname + ".cv2");
|
||||
|
||||
IElementWiseLayer* cv3 = convBnSilu(network, weightMap, *cv1->getOutput(0), c_, 3, 1, 1, lname + ".cv3");
|
||||
IElementWiseLayer* cv4 = convBnSilu(network, weightMap, *cv3->getOutput(0), c_, 1, 1, 0, lname + ".cv4");
|
||||
|
||||
IPoolingLayer* m1 = network->addPoolingNd(*cv4->getOutput(0), PoolingType::kMAX, DimsHW{ 5, 5 });
|
||||
m1->setStrideNd(DimsHW{ 1, 1 });
|
||||
m1->setPaddingNd(DimsHW{ 2, 2 });
|
||||
IPoolingLayer* m2 = network->addPoolingNd(*cv4->getOutput(0), PoolingType::kMAX, DimsHW{ 9, 9 });
|
||||
m2->setStrideNd(DimsHW{ 1, 1 });
|
||||
m2->setPaddingNd(DimsHW{ 4, 4 });
|
||||
IPoolingLayer* m3 = network->addPoolingNd(*cv4->getOutput(0), PoolingType::kMAX, DimsHW{ 13, 13 });
|
||||
m3->setStrideNd(DimsHW{ 1, 1 });
|
||||
m3->setPaddingNd(DimsHW{ 6, 6 });
|
||||
|
||||
ITensor* input_tensors[] = { cv4->getOutput(0), m1->getOutput(0), m2->getOutput(0), m3->getOutput(0) };
|
||||
IConcatenationLayer* concat = network->addConcatenation(input_tensors, 4);
|
||||
// 0U
|
||||
concat->setAxis(0);
|
||||
|
||||
IElementWiseLayer* cv5 = convBnSilu(network, weightMap, *concat->getOutput(0), c_, 1, 1, 0, lname + ".cv5");
|
||||
IElementWiseLayer* cv6 = convBnSilu(network, weightMap, *cv5->getOutput(0), c_, 3, 1, 1, lname + ".cv6");
|
||||
|
||||
ITensor* input_tensors2[] = { cv6->getOutput(0), cv2->getOutput(0) };
|
||||
IConcatenationLayer* concat1 = network->addConcatenation(input_tensors2, 2);
|
||||
// 0U
|
||||
concat1->setAxis(0);
|
||||
|
||||
|
||||
IElementWiseLayer* cv7 = convBnSilu(network, weightMap, *concat1->getOutput(0), c2, 1, 1, 0, lname + ".cv7");
|
||||
return cv7;
|
||||
}
|
||||
|
||||
// RepConv
|
||||
IElementWiseLayer* RepConv(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int c2, int k, int s, const std::string& lname) {
|
||||
Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
|
||||
// 256 * 128 * 3 *3
|
||||
IConvolutionLayer* rbr_dense_conv = network->addConvolutionNd(input, c2, DimsHW{ k, k }, weightMap[lname + ".rbr_dense.0.weight"], emptywts);
|
||||
assert(rbr_dense_conv);
|
||||
rbr_dense_conv->setPaddingNd(DimsHW{ k / 2, k / 2 });
|
||||
rbr_dense_conv->setStrideNd(DimsHW{ s, s });
|
||||
rbr_dense_conv->setName((lname + ".rbr_dense.0").c_str());
|
||||
IScaleLayer* rbr_dense_bn = addBatchNorm2d(network, weightMap, *rbr_dense_conv->getOutput(0), lname + ".rbr_dense.1", 1e-3);
|
||||
|
||||
IConvolutionLayer* rbr_1x1_conv = network->addConvolutionNd(input, c2, DimsHW{ 1, 1 }, weightMap[lname + ".rbr_1x1.0.weight"], emptywts);
|
||||
assert(rbr_1x1_conv);
|
||||
rbr_1x1_conv->setStrideNd(DimsHW{ s, s });
|
||||
rbr_1x1_conv->setName((lname + ".rbr_1x1.0").c_str());
|
||||
IScaleLayer* rbr_1x1_bn = addBatchNorm2d(network, weightMap, *rbr_1x1_conv->getOutput(0), lname + ".rbr_1x1.1", 1e-3);
|
||||
|
||||
IElementWiseLayer* ew1 = network->addElementWise(*rbr_dense_bn->getOutput(0), *rbr_1x1_bn->getOutput(0), ElementWiseOperation::kSUM);
|
||||
assert(ew1);
|
||||
// silu
|
||||
IActivationLayer* sigmoid = network->addActivation(*ew1->getOutput(0), ActivationType::kSIGMOID);
|
||||
IElementWiseLayer* ew2 = network->addElementWise(*ew1->getOutput(0), *sigmoid->getOutput(0), ElementWiseOperation::kPROD);
|
||||
return ew2;
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
IConcatenationLayer* MPC3(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int outch, std::string lname1, std::string lname2, std::string lname3) {
|
||||
Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
|
||||
|
||||
IPoolingLayer* mp1 = network->addPoolingNd(input, PoolingType::kMAX, DimsHW{ 2, 2 });
|
||||
mp1->setStrideNd(DimsHW{ 2, 2 });
|
||||
IElementWiseLayer* conv2 = convBnSilu(network, weightMap, *mp1->getOutput(0), outch, 1, 1, 0, lname1); // 左侧分支
|
||||
|
||||
IElementWiseLayer* conv3 = convBnSilu(network, weightMap, input, outch, 1, 1, 0, lname2);
|
||||
IElementWiseLayer* conv4 = convBnSilu(network, weightMap, *conv3->getOutput(0), outch, 3, 2, 1, lname3);
|
||||
ITensor* input_tensor_1[] = { conv4->getOutput(0), conv2->getOutput(0)};
|
||||
IConcatenationLayer* concat1 = network->addConcatenation(input_tensor_1, 2);
|
||||
concat1->setAxis(0);
|
||||
|
||||
return concat1;
|
||||
}
|
||||
|
||||
IElementWiseLayer* C9_1(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int outch,int outch2, std::string lname1, std::string lname2, std::string lname3,
|
||||
std::string lname4, std::string lname5, std::string lname6, std::string lname7, std::string lname8, std::string lname9)
|
||||
{
|
||||
Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
|
||||
|
||||
IElementWiseLayer* conv1 = convBnSilu(network, weightMap, input, outch, 1, 1, 0, lname1);//单独得分支
|
||||
|
||||
|
||||
IElementWiseLayer* conv2 = convBnSilu(network, weightMap, input, outch, 1, 1, 0, lname2);
|
||||
IElementWiseLayer* conv3 = convBnSilu(network, weightMap, *conv2->getOutput(0), outch, 3, 1, 1, lname3);
|
||||
IElementWiseLayer* conv4 = convBnSilu(network, weightMap, *conv3->getOutput(0), outch, 3, 1, 1, lname4);
|
||||
IElementWiseLayer* conv5 = convBnSilu(network, weightMap, *conv4->getOutput(0), outch, 3, 1, 1, lname5);
|
||||
IElementWiseLayer* conv6 = convBnSilu(network, weightMap, *conv5->getOutput(0), outch, 3, 1, 1, lname6);
|
||||
IElementWiseLayer* conv7 = convBnSilu(network, weightMap, *conv6->getOutput(0), outch, 3, 1, 1, lname7);
|
||||
IElementWiseLayer* conv8 = convBnSilu(network, weightMap, *conv7->getOutput(0), outch, 3, 1, 1, lname8);
|
||||
ITensor* input_tensor_9[] = { conv8->getOutput(0), conv6->getOutput(0), conv4->getOutput(0), conv2->getOutput(0),conv1->getOutput(0) };
|
||||
IConcatenationLayer* concat9 = network->addConcatenation(input_tensor_9, 5);
|
||||
//concat9->setAxis(0);
|
||||
IElementWiseLayer* conv10 = convBnSilu(network, weightMap, *concat9->getOutput(0), outch2, 1, 1, 0, lname9);
|
||||
|
||||
|
||||
return conv10;
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
ILayer* convBlock(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int outch, int ksize, int s, int g, std::string lname) {
|
||||
Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
|
||||
int p = ksize / 3;
|
||||
IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[lname + ".conv.weight"], emptywts);
|
||||
assert(conv1);
|
||||
conv1->setStrideNd(DimsHW{ s, s });
|
||||
conv1->setPaddingNd(DimsHW{ p, p });
|
||||
conv1->setNbGroups(g);
|
||||
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn", 1e-3);
|
||||
|
||||
// silu = x * sigmoid
|
||||
auto sig = network->addActivation(*bn1->getOutput(0), ActivationType::kSIGMOID);
|
||||
assert(sig);
|
||||
auto ew = network->addElementWise(*bn1->getOutput(0), *sig->getOutput(0), ElementWiseOperation::kPROD);
|
||||
assert(ew);
|
||||
return ew;
|
||||
}
|
||||
|
||||
IActivationLayer* convBlockLeakRelu(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int outch, int ksize, int s, int p, std::string lname) {
|
||||
Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
|
||||
|
||||
IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[lname + ".conv.weight"], emptywts);
|
||||
assert(conv1);
|
||||
conv1->setName((lname + ".conv").c_str());
|
||||
conv1->setStrideNd(DimsHW{ s, s });
|
||||
conv1->setPaddingNd(DimsHW{ p, p });
|
||||
//conv1->setNbGroups(g);
|
||||
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn", 1e-5);
|
||||
|
||||
auto ew1 = network->addActivation(*bn1->getOutput(0), ActivationType::kLEAKY_RELU);
|
||||
ew1->setAlpha(0.1);
|
||||
return ew1;
|
||||
}
|
||||
|
||||
|
||||
|
||||
ILayer* focus(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int inch, int outch, int ksize, std::string lname) {
|
||||
ISliceLayer *s1 = network->addSlice(input, Dims3{ 0, 0, 0 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 });
|
||||
ISliceLayer *s2 = network->addSlice(input, Dims3{ 0, 1, 0 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 });
|
||||
ISliceLayer *s3 = network->addSlice(input, Dims3{ 0, 0, 1 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 });
|
||||
ISliceLayer *s4 = network->addSlice(input, Dims3{ 0, 1, 1 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 });
|
||||
ITensor* inputTensors[] = { s1->getOutput(0), s2->getOutput(0), s3->getOutput(0), s4->getOutput(0) };
|
||||
auto cat = network->addConcatenation(inputTensors, 4);
|
||||
auto conv = convBlock(network, weightMap, *cat->getOutput(0), outch, ksize, 1, 1, lname + ".conv");
|
||||
return conv;
|
||||
}
|
||||
|
||||
ILayer* bottleneck(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, bool shortcut, int g, float e, std::string lname) {
|
||||
auto cv1 = convBlock(network, weightMap, input, (int)((float)c2 * e), 1, 1, 1, lname + ".cv1");
|
||||
auto cv2 = convBlock(network, weightMap, *cv1->getOutput(0), c2, 3, 1, g, lname + ".cv2");
|
||||
if (shortcut && c1 == c2) {
|
||||
auto ew = network->addElementWise(input, *cv2->getOutput(0), ElementWiseOperation::kSUM);
|
||||
return ew;
|
||||
}
|
||||
return cv2;
|
||||
}
|
||||
|
||||
ILayer* bottleneckCSP(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) {
|
||||
Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
|
||||
int c_ = (int)((float)c2 * e);
|
||||
auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1");
|
||||
auto cv2 = network->addConvolutionNd(input, c_, DimsHW{ 1, 1 }, weightMap[lname + ".cv2.weight"], emptywts);
|
||||
ITensor *y1 = cv1->getOutput(0);
|
||||
for (int i = 0; i < n; i++) {
|
||||
auto b = bottleneck(network, weightMap, *y1, c_, c_, shortcut, g, 1.0, lname + ".m." + std::to_string(i));
|
||||
y1 = b->getOutput(0);
|
||||
}
|
||||
auto cv3 = network->addConvolutionNd(*y1, c_, DimsHW{ 1, 1 }, weightMap[lname + ".cv3.weight"], emptywts);
|
||||
|
||||
ITensor* inputTensors[] = { cv3->getOutput(0), cv2->getOutput(0) };
|
||||
auto cat = network->addConcatenation(inputTensors, 2);
|
||||
|
||||
IScaleLayer* bn = addBatchNorm2d(network, weightMap, *cat->getOutput(0), lname + ".bn", 1e-4);
|
||||
auto lr = network->addActivation(*bn->getOutput(0), ActivationType::kLEAKY_RELU);
|
||||
lr->setAlpha(0.1);
|
||||
|
||||
auto cv4 = convBlock(network, weightMap, *lr->getOutput(0), c2, 1, 1, 1, lname + ".cv4");
|
||||
return cv4;
|
||||
}
|
||||
|
||||
|
||||
|
||||
ILayer* C3(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)((float)c2 * e);
|
||||
auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1");
|
||||
auto cv2 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv2");
|
||||
ITensor *y1 = cv1->getOutput(0);
|
||||
for (int i = 0; i < n; i++) {
|
||||
auto b = bottleneck(network, weightMap, *y1, c_, c_, shortcut, g, 1.0, lname + ".m." + std::to_string(i));
|
||||
y1 = b->getOutput(0);
|
||||
}
|
||||
|
||||
ITensor* inputTensors[] = { y1, cv2->getOutput(0) };
|
||||
auto cat = network->addConcatenation(inputTensors, 2);
|
||||
|
||||
auto cv3 = convBlock(network, weightMap, *cat->getOutput(0), c2, 1, 1, 1, lname + ".cv3");
|
||||
return cv3;
|
||||
}
|
||||
|
||||
ILayer* SPP(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, int k1, int k2, int k3, std::string lname) {
|
||||
int c_ = c1 / 2;
|
||||
auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1");
|
||||
|
||||
auto pool1 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k1, k1 });
|
||||
pool1->setPaddingNd(DimsHW{ k1 / 2, k1 / 2 });
|
||||
pool1->setStrideNd(DimsHW{ 1, 1 });
|
||||
auto pool2 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k2, k2 });
|
||||
pool2->setPaddingNd(DimsHW{ k2 / 2, k2 / 2 });
|
||||
pool2->setStrideNd(DimsHW{ 1, 1 });
|
||||
auto pool3 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k3, k3 });
|
||||
pool3->setPaddingNd(DimsHW{ k3 / 2, k3 / 2 });
|
||||
pool3->setStrideNd(DimsHW{ 1, 1 });
|
||||
|
||||
ITensor* inputTensors[] = { cv1->getOutput(0), pool1->getOutput(0), pool2->getOutput(0), pool3->getOutput(0) };
|
||||
auto cat = network->addConcatenation(inputTensors, 4);
|
||||
|
||||
auto cv2 = convBlock(network, weightMap, *cat->getOutput(0), c2, 1, 1, 1, lname + ".cv2");
|
||||
return cv2;
|
||||
}
|
||||
|
||||
ILayer* SPPF(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, int k, std::string lname) {
|
||||
int c_ = c1 / 2;
|
||||
auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1");
|
||||
|
||||
auto pool1 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k, k });
|
||||
pool1->setPaddingNd(DimsHW{ k / 2, k / 2 });
|
||||
pool1->setStrideNd(DimsHW{ 1, 1 });
|
||||
auto pool2 = network->addPoolingNd(*pool1->getOutput(0), PoolingType::kMAX, DimsHW{ k, k });
|
||||
pool2->setPaddingNd(DimsHW{ k / 2, k / 2 });
|
||||
pool2->setStrideNd(DimsHW{ 1, 1 });
|
||||
auto pool3 = network->addPoolingNd(*pool2->getOutput(0), PoolingType::kMAX, DimsHW{ k, k });
|
||||
pool3->setPaddingNd(DimsHW{ k / 2, k / 2 });
|
||||
pool3->setStrideNd(DimsHW{ 1, 1 });
|
||||
ITensor* inputTensors[] = { cv1->getOutput(0), pool1->getOutput(0), pool2->getOutput(0), pool3->getOutput(0) };
|
||||
auto cat = network->addConcatenation(inputTensors, 4);
|
||||
auto cv2 = convBlock(network, weightMap, *cat->getOutput(0), c2, 1, 1, 1, lname + ".cv2");
|
||||
return cv2;
|
||||
}
|
||||
|
||||
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 = Yolo::CHECK_COUNT * 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;
|
||||
}
|
||||
|
||||
IPluginV2Layer* addYoLoLayer(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, std::string lname, std::vector<IConvolutionLayer*> dets) {
|
||||
auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1");
|
||||
auto anchors = getAnchors(weightMap, lname);
|
||||
float a[3][6] = { {12.0,16.0,19.0,36.0,40.0,28.0},{36.0,75.0,76.0,55.0,72.0,146.0},{142.0,110.0,192.0,243.0,459.0,401.0} };
|
||||
//std::vector<std::vector<float>> anchors = { {12.0,16.0,19.0,36.0,40.0,28.0},{36.0,75.0,76.0,55.0,72.0,146.0},{142.0,110.0,192.0,243.0,459.0,401.0} };
|
||||
/*for (int i = 0; i < 3; i++)
|
||||
{
|
||||
for (int j = 0; j < 6; j++)
|
||||
{
|
||||
anchors[i][j] = a[i][j];
|
||||
|
||||
}
|
||||
}*/
|
||||
// for (int i = 0; i < anchors.size(); i++)
|
||||
// {
|
||||
// for (int j = 0; j < anchors[i].size(); j++)
|
||||
// {
|
||||
// std::cout << anchors[i][j] << " ";
|
||||
//
|
||||
// }
|
||||
// std::cout << std::endl;
|
||||
// }
|
||||
|
||||
PluginField plugin_fields[2];
|
||||
int netinfo[4] = {Yolo::CLASS_NUM, Yolo::INPUT_W, Yolo::INPUT_H, Yolo::MAX_OUTPUT_BBOX_COUNT};
|
||||
plugin_fields[0].data = netinfo;
|
||||
plugin_fields[0].length = 4;
|
||||
plugin_fields[0].name = "netinfo";
|
||||
plugin_fields[0].type = PluginFieldType::kFLOAT32;
|
||||
int scale = 8;
|
||||
|
||||
std::vector<Yolo::YoloKernel> kernels;
|
||||
for (size_t i = 0; i < anchors.size(); i++) {
|
||||
Yolo::YoloKernel kernel;
|
||||
kernel.width = Yolo::INPUT_W / scale;
|
||||
kernel.height = Yolo::INPUT_H / scale;
|
||||
memcpy(kernel.anchors, &anchors[i][0], anchors[i].size() * sizeof(float));
|
||||
kernels.push_back(kernel);
|
||||
scale *= 2;
|
||||
}
|
||||
plugin_fields[1].data = &kernels[0];
|
||||
plugin_fields[1].length = kernels.size();
|
||||
plugin_fields[1].name = "kernels";
|
||||
plugin_fields[1].type = PluginFieldType::kFLOAT32;
|
||||
PluginFieldCollection plugin_data;
|
||||
plugin_data.nbFields = 2;
|
||||
plugin_data.fields = plugin_fields;
|
||||
IPluginV2 *plugin_obj = creator->createPlugin("yololayer", &plugin_data);
|
||||
std::vector<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;
|
||||
}
|
||||
|
||||
|
||||
IPluginV2Layer* addYoLoLayer2(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, std::string lname, std::vector<IShuffleLayer*> dets) {
|
||||
auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1");
|
||||
auto anchors = getAnchors(weightMap, lname);
|
||||
float a[3][6] = { {12.0,16.0,19.0,36.0,40.0,28.0},{36.0,75.0,76.0,55.0,72.0,146.0},{142.0,110.0,192.0,243.0,459.0,401.0} };
|
||||
//std::vector<std::vector<float>> anchors = { {12.0,16.0,19.0,36.0,40.0,28.0},{36.0,75.0,76.0,55.0,72.0,146.0},{142.0,110.0,192.0,243.0,459.0,401.0} };
|
||||
/*for (int i = 0; i < 3; i++)
|
||||
{
|
||||
for (int j = 0; j < 6; j++)
|
||||
{
|
||||
anchors[i][j] = a[i][j];
|
||||
|
||||
}
|
||||
}*/
|
||||
// for (int i = 0; i < anchors.size(); i++)
|
||||
// {
|
||||
// for (int j = 0; j < anchors[i].size(); j++)
|
||||
// {
|
||||
// std::cout << anchors[i][j] << " ";
|
||||
//
|
||||
// }
|
||||
// std::cout << std::endl;
|
||||
// }
|
||||
|
||||
PluginField plugin_fields[2];
|
||||
int netinfo[4] = { Yolo::CLASS_NUM, Yolo::INPUT_W, Yolo::INPUT_H, Yolo::MAX_OUTPUT_BBOX_COUNT };
|
||||
plugin_fields[0].data = netinfo;
|
||||
plugin_fields[0].length = 4;
|
||||
plugin_fields[0].name = "netinfo";
|
||||
plugin_fields[0].type = PluginFieldType::kFLOAT32;
|
||||
int scale = 8;
|
||||
|
||||
std::vector<Yolo::YoloKernel> kernels;
|
||||
for (size_t i = 0; i < anchors.size(); i++) {
|
||||
Yolo::YoloKernel kernel;
|
||||
kernel.width = Yolo::INPUT_W / scale;
|
||||
kernel.height = Yolo::INPUT_H / scale;
|
||||
memcpy(kernel.anchors, &anchors[i][0], anchors[i].size() * sizeof(float));
|
||||
kernels.push_back(kernel);
|
||||
scale *= 2;
|
||||
}
|
||||
plugin_fields[1].data = &kernels[0];
|
||||
plugin_fields[1].length = kernels.size();
|
||||
plugin_fields[1].name = "kernels";
|
||||
plugin_fields[1].type = PluginFieldType::kFLOAT32;
|
||||
PluginFieldCollection plugin_data;
|
||||
plugin_data.nbFields = 2;
|
||||
plugin_data.fields = plugin_fields;
|
||||
IPluginV2* plugin_obj = creator->createPlugin("yololayer", &plugin_data);
|
||||
std::vector<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;
|
||||
}
|
||||
#endif
|
||||
|
||||
18
yolov7/cuda_utils.h
Normal file
18
yolov7/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_
|
||||
|
||||
16
yolov7/gen_wts.py
Normal file
16
yolov7/gen_wts.py
Normal file
@ -0,0 +1,16 @@
|
||||
import struct
|
||||
import torch
|
||||
|
||||
state_dict = torch.load("yolov7.pt", map_location="cpu")
|
||||
state_dict = state_dict["model"].state_dict()
|
||||
|
||||
with open("yolov7.wts", 'w') as f:
|
||||
f.write("{}\n".format(len(state_dict.keys())))
|
||||
|
||||
for k, v in 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')
|
||||
504
yolov7/logging.h
Normal file
504
yolov7/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
|
||||
27
yolov7/macros.h
Normal file
27
yolov7/macros.h
Normal file
@ -0,0 +1,27 @@
|
||||
#ifndef __MACROS_H
|
||||
#define __MACROS_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
|
||||
116
yolov7/preprocess.cu
Normal file
116
yolov7/preprocess.cu
Normal file
@ -0,0 +1,116 @@
|
||||
#include "preprocess.h"
|
||||
#include <opencv2/opencv.hpp>
|
||||
|
||||
__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 preprocess_kernel_img(
|
||||
uint8_t* src, int src_width, int src_height,
|
||||
float* dst, int dst_width, int dst_height,
|
||||
cudaStream_t 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>>>(
|
||||
src, src_width*3, src_width,
|
||||
src_height, dst, dst_width,
|
||||
dst_height, 128, d2s, jobs);
|
||||
|
||||
}
|
||||
16
yolov7/preprocess.h
Normal file
16
yolov7/preprocess.h
Normal file
@ -0,0 +1,16 @@
|
||||
#ifndef __PREPROCESS_H
|
||||
#define __PREPROCESS_H
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cstdint>
|
||||
|
||||
|
||||
struct AffineMatrix{
|
||||
float value[6];
|
||||
};
|
||||
|
||||
|
||||
void preprocess_kernel_img(uint8_t* src, int src_width, int src_height,
|
||||
float* dst, int dst_width, int dst_height,
|
||||
cudaStream_t stream);
|
||||
#endif // __PREPROCESS_H
|
||||
1
yolov7/samples
Normal file
1
yolov7/samples
Normal file
@ -0,0 +1 @@
|
||||
../yolov3-spp/samples/
|
||||
52
yolov7/utils.h
Normal file
52
yolov7/utils.h
Normal file
@ -0,0 +1,52 @@
|
||||
#ifndef TRTX_YOLOV7_UTILS_H_
|
||||
#define TRTX_YOLOV7_UTILS_H_
|
||||
|
||||
#include <dirent.h>
|
||||
#include <opencv2/opencv.hpp>
|
||||
|
||||
static inline 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;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
#endif // TRTX_YOLOV7_UTILS_H_
|
||||
|
||||
328
yolov7/yololayer.cu
Normal file
328
yolov7/yololayer.cu
Normal file
@ -0,0 +1,328 @@
|
||||
#include <assert.h>
|
||||
#include <vector>
|
||||
#include <iostream>
|
||||
#include "yololayer.h"
|
||||
#include "cuda_utils.h"
|
||||
#include <device_launch_parameters.h>
|
||||
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
using namespace Yolo;
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
YoloLayerPlugin::YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, const std::vector<Yolo::YoloKernel>& vYoloKernel)
|
||||
{
|
||||
mClassCount = classCount;
|
||||
mYoloV5NetWidth = netWidth;
|
||||
mYoloV5NetHeight = netHeight;
|
||||
mMaxOutObject = maxOut;
|
||||
/*mYoloKernel.clear();
|
||||
mYoloKernel.push_back(yolo1);
|
||||
mYoloKernel.push_back(yolo2);
|
||||
mYoloKernel.push_back(yolo3);
|
||||
mKernelCount = mYoloKernel.size();*/
|
||||
mYoloKernel = vYoloKernel;
|
||||
mKernelCount = vYoloKernel.size();
|
||||
|
||||
CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*)));
|
||||
size_t AnchorLen = sizeof(float)* CHECK_COUNT * 2;
|
||||
for (int ii = 0; ii < mKernelCount; ii++)
|
||||
{
|
||||
CUDA_CHECK(cudaMalloc(&mAnchor[ii], AnchorLen));
|
||||
const auto& yolo = mYoloKernel[ii];
|
||||
CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice));
|
||||
}
|
||||
}
|
||||
YoloLayerPlugin::~YoloLayerPlugin()
|
||||
{
|
||||
for (int ii = 0; ii < mKernelCount; ii++)
|
||||
{
|
||||
CUDA_CHECK(cudaFree(mAnchor[ii]));
|
||||
}
|
||||
CUDA_CHECK(cudaFreeHost(mAnchor));
|
||||
}
|
||||
|
||||
// create the plugin at runtime from a byte stream
|
||||
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, mKernelCount);
|
||||
read(d, mYoloV5NetWidth);
|
||||
read(d, mYoloV5NetHeight);
|
||||
read(d, mMaxOutObject);
|
||||
mYoloKernel.resize(mKernelCount);
|
||||
auto kernelSize = mKernelCount * sizeof(YoloKernel);
|
||||
memcpy(mYoloKernel.data(), d, kernelSize);
|
||||
d += kernelSize;
|
||||
CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*)));
|
||||
size_t AnchorLen = sizeof(float)* CHECK_COUNT * 2;
|
||||
for (int ii = 0; ii < mKernelCount; ii++)
|
||||
{
|
||||
CUDA_CHECK(cudaMalloc(&mAnchor[ii], AnchorLen));
|
||||
const auto& yolo = mYoloKernel[ii];
|
||||
CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice));
|
||||
}
|
||||
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, mKernelCount);
|
||||
write(d, mYoloV5NetWidth);
|
||||
write(d, mYoloV5NetHeight);
|
||||
write(d, mMaxOutObject);
|
||||
auto kernelSize = mKernelCount * sizeof(YoloKernel);
|
||||
memcpy(d, mYoloKernel.data(), kernelSize);
|
||||
d += kernelSize;
|
||||
|
||||
assert(d == a + getSerializationSize());
|
||||
}
|
||||
|
||||
size_t YoloLayerPlugin::getSerializationSize() const TRT_NOEXCEPT
|
||||
{
|
||||
return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mKernelCount) + sizeof(Yolo::YoloKernel) * mYoloKernel.size() + sizeof(mYoloV5NetWidth) + sizeof(mYoloV5NetHeight) + sizeof(mMaxOutObject);
|
||||
}
|
||||
|
||||
int YoloLayerPlugin::initialize() TRT_NOEXCEPT
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
|
||||
Dims YoloLayerPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRT_NOEXCEPT
|
||||
{
|
||||
//output the result to channel
|
||||
int totalsize = mMaxOutObject * sizeof(Detection) / sizeof(float);
|
||||
|
||||
return Dims3(totalsize + 1, 1, 1);
|
||||
}
|
||||
|
||||
// Set plugin namespace
|
||||
void YoloLayerPlugin::setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT
|
||||
{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
|
||||
const char* YoloLayerPlugin::getPluginNamespace() const TRT_NOEXCEPT
|
||||
{
|
||||
return mPluginNamespace;
|
||||
}
|
||||
|
||||
// Return the DataType of the plugin output at the requested index
|
||||
DataType YoloLayerPlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT
|
||||
{
|
||||
return DataType::kFLOAT;
|
||||
}
|
||||
|
||||
// Return true if output tensor is broadcast across a batch.
|
||||
bool YoloLayerPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
// Return true if plugin can use input that is broadcast across batch without replication.
|
||||
bool YoloLayerPlugin::canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
void YoloLayerPlugin::configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) TRT_NOEXCEPT
|
||||
{
|
||||
}
|
||||
|
||||
// Attach the plugin object to an execution context and grant the plugin the access to some context resource.
|
||||
void YoloLayerPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT
|
||||
{
|
||||
}
|
||||
|
||||
// Detach the plugin object from its execution context.
|
||||
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;
|
||||
}
|
||||
|
||||
// Clone the plugin
|
||||
IPluginV2IOExt* YoloLayerPlugin::clone() const TRT_NOEXCEPT
|
||||
{
|
||||
YoloLayerPlugin* p = new YoloLayerPlugin(mClassCount, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, mYoloKernel);
|
||||
p->setPluginNamespace(mPluginNamespace);
|
||||
return p;
|
||||
}
|
||||
|
||||
__device__ float Logist(float data) { return 1.0f / (1.0f + expf(-data)); };
|
||||
|
||||
__global__ void CalDetection(const float *input, float *output, int noElements,
|
||||
const int netwidth, const int netheight, int maxoutobject, int yoloWidth, int yoloHeight, const float anchors[CHECK_COUNT * 2], int classes, int outputElem)
|
||||
{
|
||||
|
||||
int idx = threadIdx.x + blockDim.x * blockIdx.x;
|
||||
if (idx >= noElements) return;
|
||||
|
||||
int total_grid = yoloWidth * yoloHeight; // 80*80 40*40 20*20
|
||||
int bnIdx = idx / total_grid;
|
||||
idx = idx - total_grid * bnIdx;
|
||||
int info_len_i = 5 + classes;
|
||||
const float* curInput = input + bnIdx * (info_len_i * total_grid * CHECK_COUNT);
|
||||
|
||||
|
||||
for (int k = 0; k < 3; k++) {
|
||||
float box_prob = Logist(curInput[idx + k * info_len_i * total_grid + 4 * total_grid]);
|
||||
if (box_prob < IGNORE_THRESH) continue;
|
||||
int class_id = 0;
|
||||
float max_cls_prob = 0.0;
|
||||
for (int i = 5; i < info_len_i; ++i) {
|
||||
float p = Logist(curInput[idx + k * info_len_i * total_grid + i * total_grid]);
|
||||
if (p > max_cls_prob) {
|
||||
max_cls_prob = p;
|
||||
class_id = i - 5;
|
||||
}
|
||||
}
|
||||
float *res_count = output + bnIdx * outputElem;
|
||||
int count = (int)atomicAdd(res_count, 1);
|
||||
if (count >= maxoutobject) return;
|
||||
char *data = (char*)res_count + sizeof(float) + count * sizeof(Detection);
|
||||
Detection *det = (Detection*)(data);
|
||||
|
||||
int row = idx / yoloWidth;
|
||||
int col = idx % yoloWidth;
|
||||
|
||||
//Location
|
||||
// pytorch:
|
||||
// y = x[i].sigmoid()
|
||||
// y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy
|
||||
// y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
|
||||
// X: (sigmoid(tx) + cx)/FeaturemapW * netwidth
|
||||
det->bbox[0] = (col - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 0 * total_grid])) * netwidth / yoloWidth;
|
||||
det->bbox[1] = (row - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 1 * total_grid])) * netheight / yoloHeight;
|
||||
|
||||
// W: (Pw * e^tw) / FeaturemapW * netwidth
|
||||
// v5: https://github.com/ultralytics/yolov7/issues/471
|
||||
//float box_w = ((row[2] * 2)*(row[2] * 2)) * float(anchors[a][c][0]) * scale;
|
||||
//float box_h = ((row[3] * 2) * (row[3] * 2)) * float(anchors[a][c][1]) * scale;
|
||||
det->bbox[2] = 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 2 * total_grid]);
|
||||
det->bbox[2] = det->bbox[2] * det->bbox[2] * anchors[2*k];
|
||||
det->bbox[3] = 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 3 * total_grid]);
|
||||
det->bbox[3] = det->bbox[3] * det->bbox[3] * anchors[2*k + 1];
|
||||
det->conf = box_prob * max_cls_prob;
|
||||
det->class_id = class_id;
|
||||
//printf("ok");
|
||||
}
|
||||
}
|
||||
|
||||
void YoloLayerPlugin::forwardGpu(const float* const* inputs, float *output, cudaStream_t stream, int batchSize)
|
||||
{
|
||||
int outputElem = 1 + mMaxOutObject * sizeof(Detection) / sizeof(float);
|
||||
for (int idx = 0; idx < batchSize; ++idx) {
|
||||
CUDA_CHECK(cudaMemsetAsync(output + idx * outputElem, 0, sizeof(float), stream));
|
||||
}
|
||||
int numElem = 0;
|
||||
|
||||
|
||||
|
||||
for (unsigned int i = 0; i < mYoloKernel.size(); ++i) {
|
||||
const auto& yolo = mYoloKernel[i];
|
||||
numElem = yolo.width * yolo.height * batchSize;
|
||||
if (numElem < mThreadCount) mThreadCount = numElem;
|
||||
|
||||
|
||||
CalDetection << < (numElem + mThreadCount - 1) / mThreadCount, mThreadCount, 0, stream >> >
|
||||
(inputs[i], output, numElem, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, yolo.width, yolo.height, (float*)mAnchor[i], mClassCount, outputElem);
|
||||
/* CalDetection << < (numElem + mThreadCount - 1) / mThreadCount2, mThreadCount2, 0, stream >> >
|
||||
(inputs[i], output, numElem, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, yolo.width, yolo.height, (float*)mAnchor[i], mClassCount, outputElem);*/
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
int YoloLayerPlugin::enqueue(int batchSize, const void* const* inputs, void* TRT_CONST_ENQUEUE* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT
|
||||
{
|
||||
forwardGpu((const float* const*)inputs, (float*)outputs[0], stream, batchSize);
|
||||
return 0;
|
||||
}
|
||||
|
||||
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 == 2);
|
||||
assert(strcmp(fc->fields[0].name, "netinfo") == 0);
|
||||
assert(strcmp(fc->fields[1].name, "kernels") == 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];
|
||||
std::vector<Yolo::YoloKernel> kernels(fc->fields[1].length);
|
||||
memcpy(&kernels[0], fc->fields[1].data, kernels.size() * sizeof(Yolo::YoloKernel));
|
||||
YoloLayerPlugin* obj = new YoloLayerPlugin(class_count, input_w, input_h, max_output_object_count, kernels);
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
157
yolov7/yololayer.h
Normal file
157
yolov7/yololayer.h
Normal file
@ -0,0 +1,157 @@
|
||||
#ifndef _YOLO_LAYER_H
|
||||
#define _YOLO_LAYER_H
|
||||
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <NvInfer.h>
|
||||
#include "macros.h"
|
||||
|
||||
namespace Yolo
|
||||
{
|
||||
static constexpr int CHECK_COUNT = 3;
|
||||
static constexpr float IGNORE_THRESH = 0.5f;
|
||||
struct YoloKernel
|
||||
{
|
||||
int width;
|
||||
int height;
|
||||
float anchors[CHECK_COUNT * 2];
|
||||
};
|
||||
static constexpr int MAX_OUTPUT_BBOX_COUNT = 1000;
|
||||
static constexpr int CLASS_NUM = 80;
|
||||
static constexpr int INPUT_H = 640; // yolov7's input height and width must be divisible by 32.
|
||||
static constexpr int INPUT_W = 640;
|
||||
|
||||
static constexpr int LOCATIONS = 4;
|
||||
struct alignas(float) Detection {
|
||||
//center_x center_y w h
|
||||
float bbox[LOCATIONS];
|
||||
float conf; // bbox_conf * cls_conf
|
||||
float class_id;
|
||||
};
|
||||
|
||||
static constexpr YoloKernel yolo1 = {
|
||||
INPUT_W / 8,
|
||||
INPUT_H / 8,
|
||||
{12,16, 19,36, 40,28}
|
||||
};
|
||||
static constexpr YoloKernel yolo2 = {
|
||||
INPUT_W / 16,
|
||||
INPUT_H / 16,
|
||||
{36,75, 76,55, 72,146}
|
||||
};
|
||||
static constexpr YoloKernel yolo3 = {
|
||||
INPUT_W / 32,
|
||||
INPUT_H / 32,
|
||||
{142,110, 192,243, 459,401}
|
||||
};
|
||||
|
||||
|
||||
}
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
class API YoloLayerPlugin : public IPluginV2IOExt
|
||||
{
|
||||
public:
|
||||
YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, const std::vector<Yolo::YoloKernel>& vYoloKernel);
|
||||
YoloLayerPlugin(const void* data, size_t length);
|
||||
~YoloLayerPlugin();
|
||||
|
||||
int getNbOutputs() const TRT_NOEXCEPT override
|
||||
{
|
||||
return 1;
|
||||
}
|
||||
|
||||
Dims getOutputDimensions(int index, const 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;
|
||||
|
||||
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT override;
|
||||
|
||||
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(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) TRT_NOEXCEPT override;
|
||||
|
||||
void detachFromContext() TRT_NOEXCEPT override;
|
||||
|
||||
private:
|
||||
void forwardGpu(const float* const* inputs, float *output, cudaStream_t stream, int batchSize = 1);
|
||||
int mThreadCount = 256;
|
||||
int mThreadCount2 = (Yolo::CLASS_NUM+5)*3+1;
|
||||
const char* mPluginNamespace;
|
||||
int mKernelCount;
|
||||
int mClassCount;
|
||||
int mYoloV5NetWidth;
|
||||
int mYoloV5NetHeight;
|
||||
int mMaxOutObject;
|
||||
std::vector<Yolo::YoloKernel> mYoloKernel;
|
||||
void** mAnchor;
|
||||
};
|
||||
|
||||
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 PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override;
|
||||
|
||||
IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) TRT_NOEXCEPT override;
|
||||
|
||||
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);
|
||||
};
|
||||
|
||||
#endif // _YOLO_LAYER_H
|
||||
2480
yolov7/yolov7.cpp
Normal file
2480
yolov7/yolov7.cpp
Normal file
File diff suppressed because it is too large
Load Diff
450
yolov7/yolov7_trt.py
Normal file
450
yolov7/yolov7_trt.py
Normal file
@ -0,0 +1,450 @@
|
||||
"""
|
||||
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 YoLov5 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 YoLov5TRT(object):
|
||||
"""
|
||||
description: A YOLOv5 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 * 6001: (i + 1) * 6001], 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] - x[:, 2] / 2
|
||||
y[:, 2] = x[:, 0] + x[:, 2] / 2
|
||||
y[:, 1] = x[:, 1] - x[:, 3] / 2 - (self.input_h - r_w * origin_h) / 2
|
||||
y[:, 3] = x[:, 1] + x[:, 3] / 2 - (self.input_h - r_w * origin_h) / 2
|
||||
y /= r_w
|
||||
else:
|
||||
y[:, 0] = x[:, 0] - x[:, 2] / 2 - (self.input_w - r_h * origin_w) / 2
|
||||
y[:, 2] = x[:, 0] + x[:, 2] / 2 - (self.input_w - r_h * origin_w) / 2
|
||||
y[:, 1] = x[:, 1] - x[:, 3] / 2
|
||||
y[:, 3] = x[:, 1] + x[:, 3] / 2
|
||||
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, 6))[: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, yolov7_wrapper, image_path_batch):
|
||||
threading.Thread.__init__(self)
|
||||
self.yolov7_wrapper = yolov7_wrapper
|
||||
self.image_path_batch = image_path_batch
|
||||
|
||||
def run(self):
|
||||
batch_image_raw, use_time = self.yolov7_wrapper.infer(self.yolov7_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, yolov7_wrapper):
|
||||
threading.Thread.__init__(self)
|
||||
self.yolov7_wrapper = yolov7_wrapper
|
||||
|
||||
def run(self):
|
||||
batch_image_raw, use_time = self.yolov7_wrapper.infer(self.yolov7_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 = "yolov7.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 YoLov5TRT instance
|
||||
yolov7_wrapper = YoLov5TRT(engine_file_path)
|
||||
try:
|
||||
print('batch size is', yolov7_wrapper.batch_size)
|
||||
|
||||
image_dir = "samples/"
|
||||
image_path_batches = get_img_path_batches(yolov7_wrapper.batch_size, image_dir)
|
||||
|
||||
for i in range(10):
|
||||
# create a new thread to do warm_up
|
||||
thread1 = warmUpThread(yolov7_wrapper)
|
||||
thread1.start()
|
||||
thread1.join()
|
||||
for batch in image_path_batches:
|
||||
# create a new thread to do inference
|
||||
thread1 = inferThread(yolov7_wrapper, batch)
|
||||
thread1.start()
|
||||
thread1.join()
|
||||
finally:
|
||||
# destroy the instance
|
||||
yolov7_wrapper.destroy()
|
||||
Loading…
Reference in New Issue
Block a user