diff --git a/README.md b/README.md index 381deb1..d836e02 100644 --- a/README.md +++ b/README.md @@ -1,18 +1,67 @@ -# YOLOv6、 YOLOX、 YOLOV5、 TensorRT Python/C++ API -## Update 2022.7.3 support TRT int8 post-training quantization +# YOLO Series TensorRT Python/C++ +## Support +YOLOv7、YOLOv6、 YOLOX、 YOLOV5、 -## Prepare TRT Python +## Update +- 2022.7.8 support YOLOV7 +- 2022.7.3 support TRT int8 post-training quantization +## Prepare TRT Env +`Python` ``` pip install --upgrade setuptools pip --user pip install nvidia-pyindex pip install --upgrade nvidia-tensorrt pip install pycuda ``` +`C++` +[By Docker](https://github.com/NVIDIA/TensorRT/blob/main/docker/ubuntu-20.04.Dockerfile) + +## Quick Start Here is a Python Demo mybe help you quickly understand this repo [Link](https://aistudio.baidu.com/aistudio/projectdetail/4263301?contributionType=1&shared=1) + +## YOLOv7 [C++, Python Support] + +![](yolov7/3_yolov7.jpg) + +```shell +https://github.com/WongKinYiu/yolov7.git +``` +修改代码:将 yolo.py 对应行修改如下： +https://github.com/WongKinYiu/yolov7/blob/5f1b78ad614b45c5a98e7afdd295e20033d5ad3c/models/yolo.py#L57 + +```python +return x if self.training else (torch.cat(z, 1), ) if not self.export else (torch.cat(z, 1), x) +``` + +### 导出onnx +```shell +python models/export.py --weights ../yolov7.pt --grid +``` + +### 转化为TensorRT Engine + +``` +python export.py -o onnx-name -e trt-name -p fp32/16/int8 +``` +### 测试 + +``` +cd yolov7 +python trt.py +``` + +### C++ + +C++ [Demo](yolov7/cpp/README.md) + + + + + ## YOLOv6 [C++, Python Support] | model | input | | FPS | Device | Language | @@ -22,7 +71,7 @@ Here is a Python Demo mybe help you quickly understand this repo [Link](https:// | yolov6s | 640*640 | FP32 | 330FPS | 1080Ti | C++ | | yolov6s | 640*640 | FP32 | 300FPS | 1080Ti | Python | -[bilibili](https://www.bilibili.com/video/BV1x3411w7T6?share_source=copy_web) +[YOLOv6 bilibili](https://www.bilibili.com/video/BV1x3411w7T6?share_source=copy_web) ![](yolov6/3_yolov6.jpg) ```shell @@ -117,4 +166,4 @@ python export.py -o onnx-name -e trt-name -p fp32/16/int8 ``` cd yolov5 python trt.py -``` +``` \ No newline at end of file diff --git a/yolov7/3_yolov7.jpg b/yolov7/3_yolov7.jpg new file mode 100644 index 0000000..58c505b Binary files /dev/null and b/yolov7/3_yolov7.jpg differ diff --git a/yolov7/cpp/CMakeLists.txt b/yolov7/cpp/CMakeLists.txt new file mode 100644 index 0000000..9614e5f --- /dev/null +++ b/yolov7/cpp/CMakeLists.txt @@ -0,0 +1,32 @@ +cmake_minimum_required(VERSION 2.6) + +project(yolov7) + +add_definitions(-std=c++11) + +option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) +set(CMAKE_CXX_STANDARD 11) +set(CMAKE_BUILD_TYPE Debug) +# add_definitions("-Wall -g") +find_package(CUDA REQUIRED) + +include_directories(${PROJECT_SOURCE_DIR}/include) +# include and link dirs of cuda and tensorrt, you need adapt them if yours are different +# cuda +include_directories(/usr/local/cuda/include) +link_directories(/usr/local/cuda/lib64) +# tensorrt +include_directories(/usr/include/x86_64-linux-gnu/) +link_directories(/usr/lib/x86_64-linux-gnu/) + +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED") + +find_package(OpenCV) +include_directories(${OpenCV_INCLUDE_DIRS}) + +add_executable(yolov7 ${PROJECT_SOURCE_DIR}/yolov7.cpp) +target_link_libraries(yolov7 nvinfer) +target_link_libraries(yolov7 cudart) +target_link_libraries(yolov7 ${OpenCV_LIBS}) + +add_definitions(-O2 -pthread) \ No newline at end of file diff --git a/yolov7/cpp/README.md b/yolov7/cpp/README.md new file mode 100644 index 0000000..b3af667 --- /dev/null +++ b/yolov7/cpp/README.md @@ -0,0 +1,58 @@ +# YOLOV7-TensorRT in C++ + +## Step 1: 准备TRT序列化引擎 + +```shell +https://github.com/WongKinYiu/yolov7.git +``` +### 导出onnx +```shell +python models/export.py --weights ../yolov7.pt --grid +``` + +### 转化为TensorRT Engine + +``` +python export.py -o onnx-name -e trt-name -p fp32/16/int8 +``` +### 测试 + +``` +cd yolov7 +python trt.py +``` + +## Step 2: C++ + +Please follow the [TensorRT Installation Guide](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html) to install TensorRT. + +And you should set the TensorRT path and CUDA path in CMakeLists.txt. + +If you train your custom dataset, you may need to modify the value of `num_class`. + +```c++ +const int num_class = 80; +``` + +Install opencv with ```sudo apt-get install libopencv-dev``` (we don't need a higher version of opencv like v3.3+). + +build the demo: + +```shell +mkdir build +cd build +cmake .. +make +``` + +Then run the demo: + +```shell +./yolov7 ../model_trt.engine -i ../../../../assets/dog.jpg +``` + +or + +```shell +./yolov7 -i +``` diff --git a/yolov7/cpp/det_res.jpg b/yolov7/cpp/det_res.jpg new file mode 100644 index 0000000..5d49caf Binary files /dev/null and b/yolov7/cpp/det_res.jpg differ diff --git a/yolov7/cpp/logging.h b/yolov7/cpp/logging.h new file mode 100644 index 0000000..e04857e --- /dev/null +++ b/yolov7/cpp/logging.h @@ -0,0 +1,503 @@ +/* + * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef TENSORRT_LOGGING_H +#define TENSORRT_LOGGING_H + +#include "NvInferRuntimeCommon.h" +#include +#include +#include +#include +#include +#include +#include + +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) noexcept + { + LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl; + } + + //! + //! \brief Method for controlling the verbosity of logging output + //! + //! \param severity The logger will only emit messages that have severity of this level or higher. + //! + void setReportableSeverity(Severity severity) + { + mReportableSeverity = severity; + } + + //! + //! \brief Opaque handle that holds logging information for a particular test + //! + //! This object is an opaque handle to information used by the Logger to print test results. + //! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used + //! with Logger::reportTest{Start,End}(). + //! + class TestAtom + { + public: + TestAtom(TestAtom&&) = default; + + private: + friend class Logger; + + TestAtom(bool started, const std::string& name, const std::string& cmdline) + : mStarted(started) + , mName(name) + , mCmdline(cmdline) + { + } + + bool mStarted; + std::string mName; + std::string mCmdline; + }; + + //! + //! \brief Define a test for logging + //! + //! \param[in] name The name of the test. This should be a string starting with + //! "TensorRT" and containing dot-separated strings containing + //! the characters [A-Za-z0-9_]. + //! For example, "TensorRT.sample_googlenet" + //! \param[in] cmdline The command line used to reproduce the test + // + //! \return a TestAtom that can be used in Logger::reportTest{Start,End}(). + //! + static TestAtom defineTest(const std::string& name, const std::string& cmdline) + { + return TestAtom(false, name, cmdline); + } + + //! + //! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments + //! as input + //! + //! \param[in] name The name of the test + //! \param[in] argc The number of command-line arguments + //! \param[in] argv The array of command-line arguments (given as C strings) + //! + //! \return a TestAtom that can be used in Logger::reportTest{Start,End}(). + static TestAtom defineTest(const std::string& name, int argc, char const* const* argv) + { + auto cmdline = genCmdlineString(argc, argv); + return defineTest(name, cmdline); + } + + //! + //! \brief Report that a test has started. + //! + //! \pre reportTestStart() has not been called yet for the given testAtom + //! + //! \param[in] testAtom The handle to the test that has started + //! + static void reportTestStart(TestAtom& testAtom) + { + reportTestResult(testAtom, TestResult::kRUNNING); + assert(!testAtom.mStarted); + testAtom.mStarted = true; + } + + //! + //! \brief Report that a test has ended. + //! + //! \pre reportTestStart() has been called for the given testAtom + //! + //! \param[in] testAtom The handle to the test that has ended + //! \param[in] result The result of the test. Should be one of TestResult::kPASSED, + //! TestResult::kFAILED, TestResult::kWAIVED + //! + static void reportTestEnd(const TestAtom& testAtom, TestResult result) + { + assert(result != TestResult::kRUNNING); + assert(testAtom.mStarted); + reportTestResult(testAtom, result); + } + + static int reportPass(const TestAtom& testAtom) + { + reportTestEnd(testAtom, TestResult::kPASSED); + return EXIT_SUCCESS; + } + + static int reportFail(const TestAtom& testAtom) + { + reportTestEnd(testAtom, TestResult::kFAILED); + return EXIT_FAILURE; + } + + static int reportWaive(const TestAtom& testAtom) + { + reportTestEnd(testAtom, TestResult::kWAIVED); + return EXIT_SUCCESS; + } + + static int reportTest(const TestAtom& testAtom, bool pass) + { + return pass ? reportPass(testAtom) : reportFail(testAtom); + } + + Severity getReportableSeverity() const + { + return mReportableSeverity; + } + +private: + //! + //! \brief returns an appropriate string for prefixing a log message with the given severity + //! + static const char* severityPrefix(Severity severity) + { + switch (severity) + { + case Severity::kINTERNAL_ERROR: return "[F] "; + case Severity::kERROR: return "[E] "; + case Severity::kWARNING: return "[W] "; + case Severity::kINFO: return "[I] "; + case Severity::kVERBOSE: return "[V] "; + default: assert(0); return ""; + } + } + + //! + //! \brief returns an appropriate string for prefixing a test result message with the given result + //! + static const char* testResultString(TestResult result) + { + switch (result) + { + case TestResult::kRUNNING: return "RUNNING"; + case TestResult::kPASSED: return "PASSED"; + case TestResult::kFAILED: return "FAILED"; + case TestResult::kWAIVED: return "WAIVED"; + default: assert(0); return ""; + } + } + + //! + //! \brief returns an appropriate output stream (cout or cerr) to use with the given severity + //! + static std::ostream& severityOstream(Severity severity) + { + return severity >= Severity::kINFO ? std::cout : std::cerr; + } + + //! + //! \brief method that implements logging test results + //! + static void reportTestResult(const TestAtom& testAtom, TestResult result) + { + severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # " + << testAtom.mCmdline << std::endl; + } + + //! + //! \brief generate a command line string from the given (argc, argv) values + //! + static std::string genCmdlineString(int argc, char const* const* argv) + { + std::stringstream ss; + for (int i = 0; i < argc; i++) + { + if (i > 0) + ss << " "; + ss << argv[i]; + } + return ss.str(); + } + + Severity mReportableSeverity; +}; + +namespace +{ + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE +//! +//! Example usage: +//! +//! LOG_VERBOSE(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) +{ + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE); +} + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO +//! +//! Example usage: +//! +//! LOG_INFO(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_INFO(const Logger& logger) +{ + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO); +} + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING +//! +//! Example usage: +//! +//! LOG_WARN(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_WARN(const Logger& logger) +{ + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING); +} + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR +//! +//! Example usage: +//! +//! LOG_ERROR(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_ERROR(const Logger& logger) +{ + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR); +} + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR +// ("fatal" severity) +//! +//! Example usage: +//! +//! LOG_FATAL(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_FATAL(const Logger& logger) +{ + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR); +} + +} // anonymous namespace + +#endif // TENSORRT_LOGGING_H diff --git a/yolov7/cpp/yolov7.cpp b/yolov7/cpp/yolov7.cpp new file mode 100644 index 0000000..a48ea5d --- /dev/null +++ b/yolov7/cpp/yolov7.cpp @@ -0,0 +1,499 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include "NvInfer.h" +#include "cuda_runtime_api.h" +#include "logging.h" + +#define CHECK(status) \ + do\ + {\ + auto ret = (status);\ + if (ret != 0)\ + {\ + std::cerr << "Cuda failure: " << ret << std::endl;\ + abort();\ + }\ + } while (0) + +#define DEVICE 0 // GPU id +#define NMS_THRESH 0.45 +#define BBOX_CONF_THRESH 0.2 + +using namespace nvinfer1; + +// stuff we know about the network and the input/output blobs +static const int INPUT_W = 640; +static const int INPUT_H = 640; +const char* INPUT_BLOB_NAME = "image_arrays"; +const char* OUTPUT_BLOB_NAME = "outputs"; +static Logger gLogger; + + +cv::Mat static_resize(cv::Mat& img) { + float r = std::min(INPUT_W / (img.cols*1.0), INPUT_H / (img.rows*1.0)); + int unpad_w = r * img.cols; + int unpad_h = r * img.rows; + cv::Mat re(unpad_h, unpad_w, CV_8UC3); + cv::resize(img, re, re.size()); + cv::Mat out(INPUT_W, INPUT_H, CV_8UC3, cv::Scalar(114, 114, 114)); + re.copyTo(out(cv::Rect(0, 0, re.cols, re.rows))); + return out; +} + +struct Object +{ + cv::Rect_ rect; + int label; + float prob; +}; + + +static inline float intersection_area(const Object& a, const Object& b) +{ + cv::Rect_ inter = a.rect & b.rect; + return inter.area(); +} + +static void qsort_descent_inplace(std::vector& faceobjects, int left, int right) +{ + int i = left; + int j = right; + float p = faceobjects[(left + right) / 2].prob; + + while (i <= j) + { + while (faceobjects[i].prob > p) + i++; + + while (faceobjects[j].prob < p) + j--; + + if (i <= j) + { + // swap + std::swap(faceobjects[i], faceobjects[j]); + + i++; + j--; + } + } + + #pragma omp parallel sections + { + #pragma omp section + { + if (left < j) qsort_descent_inplace(faceobjects, left, j); + } + #pragma omp section + { + if (i < right) qsort_descent_inplace(faceobjects, i, right); + } + } +} + +static void qsort_descent_inplace(std::vector& objects) +{ + if (objects.empty()) + return; + + qsort_descent_inplace(objects, 0, objects.size() - 1); +} + +static void nms_sorted_bboxes(const std::vector& faceobjects, std::vector& picked, float nms_threshold) +{ + picked.clear(); + + const int n = faceobjects.size(); + + std::vector areas(n); + for (int i = 0; i < n; i++) + { + areas[i] = faceobjects[i].rect.area(); + } + + for (int i = 0; i < n; i++) + { + const Object& a = faceobjects[i]; + + int keep = 1; + for (int j = 0; j < (int)picked.size(); j++) + { + const Object& b = faceobjects[picked[j]]; + + // intersection over union + float inter_area = intersection_area(a, b); + float union_area = areas[i] + areas[picked[j]] - inter_area; + // float IoU = inter_area / union_area + if (inter_area / union_area > nms_threshold) + keep = 0; + } + + if (keep) + picked.push_back(i); + } +} + + +static void generate_yolo_proposals(float* feat_blob, int output_size, float prob_threshold, std::vector& objects) +{ + const int num_class = 80; + auto dets = output_size / (num_class + 5); + for (int boxs_idx = 0; boxs_idx < dets; boxs_idx++) + { + const int basic_pos = boxs_idx *(num_class + 5); + float x_center = feat_blob[basic_pos+0]; + float y_center = feat_blob[basic_pos+1]; + float w = feat_blob[basic_pos+2]; + float h = feat_blob[basic_pos+3]; + float x0 = x_center - w * 0.5f; + float y0 = y_center - h * 0.5f; + float box_objectness = feat_blob[basic_pos+4]; + // std::cout<<*feat_blob< prob_threshold) + { + Object obj; + obj.rect.x = x0; + obj.rect.y = y0; + obj.rect.width = w; + obj.rect.height = h; + obj.label = class_idx; + obj.prob = box_prob; + + objects.push_back(obj); + } + + } // class loop + } + +} + +float* blobFromImage(cv::Mat& img){ + cv::cvtColor(img, img, cv::COLOR_BGR2RGB); + + float* blob = new float[img.total()*3]; + int channels = 3; + int img_h = img.rows; + int img_w = img.cols; + for (size_t c = 0; c < channels; c++) + { + for (size_t h = 0; h < img_h; h++) + { + for (size_t w = 0; w < img_w; w++) + { + blob[c * img_w * img_h + h * img_w + w] = + (((float)img.at(h, w)[c]) / 255.0f); + } + } + } + return blob; +} + + +static void decode_outputs(float* prob, int output_size, std::vector& objects, float scale, const int img_w, const int img_h) { + std::vector proposals; + generate_yolo_proposals(prob, output_size, BBOX_CONF_THRESH, proposals); + std::cout << "num of boxes before nms: " << proposals.size() << std::endl; + + qsort_descent_inplace(proposals); + + std::vector picked; + nms_sorted_bboxes(proposals, picked, NMS_THRESH); + + + int count = picked.size(); + + std::cout << "num of boxes: " << count << std::endl; + + objects.resize(count); + for (int i = 0; i < count; i++) + { + objects[i] = proposals[picked[i]]; + + // adjust offset to original unpadded + float x0 = (objects[i].rect.x) / scale; + float y0 = (objects[i].rect.y) / scale; + float x1 = (objects[i].rect.x + objects[i].rect.width) / scale; + float y1 = (objects[i].rect.y + objects[i].rect.height) / scale; + + // clip + x0 = std::max(std::min(x0, (float)(img_w - 1)), 0.f); + y0 = std::max(std::min(y0, (float)(img_h - 1)), 0.f); + x1 = std::max(std::min(x1, (float)(img_w - 1)), 0.f); + y1 = std::max(std::min(y1, (float)(img_h - 1)), 0.f); + + objects[i].rect.x = x0; + objects[i].rect.y = y0; + objects[i].rect.width = x1 - x0; + objects[i].rect.height = y1 - y0; + } +} + +const float color_list[80][3] = +{ + {0.000, 0.447, 0.741}, + {0.850, 0.325, 0.098}, + {0.929, 0.694, 0.125}, + {0.494, 0.184, 0.556}, + {0.466, 0.674, 0.188}, + {0.301, 0.745, 0.933}, + {0.635, 0.078, 0.184}, + {0.300, 0.300, 0.300}, + {0.600, 0.600, 0.600}, + {1.000, 0.000, 0.000}, + {1.000, 0.500, 0.000}, + {0.749, 0.749, 0.000}, + {0.000, 1.000, 0.000}, + {0.000, 0.000, 1.000}, + {0.667, 0.000, 1.000}, + {0.333, 0.333, 0.000}, + {0.333, 0.667, 0.000}, + {0.333, 1.000, 0.000}, + {0.667, 0.333, 0.000}, + {0.667, 0.667, 0.000}, + {0.667, 1.000, 0.000}, + {1.000, 0.333, 0.000}, + {1.000, 0.667, 0.000}, + {1.000, 1.000, 0.000}, + {0.000, 0.333, 0.500}, + {0.000, 0.667, 0.500}, + {0.000, 1.000, 0.500}, + {0.333, 0.000, 0.500}, + {0.333, 0.333, 0.500}, + {0.333, 0.667, 0.500}, + {0.333, 1.000, 0.500}, + {0.667, 0.000, 0.500}, + {0.667, 0.333, 0.500}, + {0.667, 0.667, 0.500}, + {0.667, 1.000, 0.500}, + {1.000, 0.000, 0.500}, + {1.000, 0.333, 0.500}, + {1.000, 0.667, 0.500}, + {1.000, 1.000, 0.500}, + {0.000, 0.333, 1.000}, + {0.000, 0.667, 1.000}, + {0.000, 1.000, 1.000}, + {0.333, 0.000, 1.000}, + {0.333, 0.333, 1.000}, + {0.333, 0.667, 1.000}, + {0.333, 1.000, 1.000}, + {0.667, 0.000, 1.000}, + {0.667, 0.333, 1.000}, + {0.667, 0.667, 1.000}, + {0.667, 1.000, 1.000}, + {1.000, 0.000, 1.000}, + {1.000, 0.333, 1.000}, + {1.000, 0.667, 1.000}, + {0.333, 0.000, 0.000}, + {0.500, 0.000, 0.000}, + {0.667, 0.000, 0.000}, + {0.833, 0.000, 0.000}, + {1.000, 0.000, 0.000}, + {0.000, 0.167, 0.000}, + {0.000, 0.333, 0.000}, + {0.000, 0.500, 0.000}, + {0.000, 0.667, 0.000}, + {0.000, 0.833, 0.000}, + {0.000, 1.000, 0.000}, + {0.000, 0.000, 0.167}, + {0.000, 0.000, 0.333}, + {0.000, 0.000, 0.500}, + {0.000, 0.000, 0.667}, + {0.000, 0.000, 0.833}, + {0.000, 0.000, 1.000}, + {0.000, 0.000, 0.000}, + {0.143, 0.143, 0.143}, + {0.286, 0.286, 0.286}, + {0.429, 0.429, 0.429}, + {0.571, 0.571, 0.571}, + {0.714, 0.714, 0.714}, + {0.857, 0.857, 0.857}, + {0.000, 0.447, 0.741}, + {0.314, 0.717, 0.741}, + {0.50, 0.5, 0} +}; + + +static void draw_objects(const cv::Mat& bgr, const std::vector& objects, std::string f) +{ + static const char* class_names[] = { + "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" + }; + + cv::Mat image = bgr.clone(); + + for (size_t i = 0; i < objects.size(); i++) + { + const Object& obj = objects[i]; + + fprintf(stderr, "%d = %.5f at %.2f %.2f %.2f x %.2f\n", obj.label, obj.prob, + obj.rect.x, obj.rect.y, obj.rect.width, obj.rect.height); + + cv::Scalar color = cv::Scalar(color_list[obj.label][0], color_list[obj.label][1], color_list[obj.label][2]); + float c_mean = cv::mean(color)[0]; + cv::Scalar txt_color; + if (c_mean > 0.5){ + txt_color = cv::Scalar(0, 0, 0); + }else{ + txt_color = cv::Scalar(255, 255, 255); + } + + cv::rectangle(image, obj.rect, color * 255, 2); + + char text[256]; + sprintf(text, "%s %.1f%%", class_names[obj.label], obj.prob * 100); + + int baseLine = 0; + cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine); + + cv::Scalar txt_bk_color = color * 0.7 * 255; + + int x = obj.rect.x; + int y = obj.rect.y + 1; + //int y = obj.rect.y - label_size.height - baseLine; + if (y > image.rows) + y = image.rows; + //if (x + label_size.width > image.cols) + //x = image.cols - label_size.width; + + cv::rectangle(image, cv::Rect(cv::Point(x, y), cv::Size(label_size.width, label_size.height + baseLine)), + txt_bk_color, -1); + + cv::putText(image, text, cv::Point(x, y + label_size.height), + cv::FONT_HERSHEY_SIMPLEX, 0.4, txt_color, 1); + } + + cv::imwrite("det_res.jpg", image); + fprintf(stderr, "save vis file\n"); + /* cv::imshow("image", image); */ + /* cv::waitKey(0); */ +} + + +void doInference(IExecutionContext& context, float* input, float* output, const int output_size, cv::Size input_shape) { + const ICudaEngine& engine = context.getEngine(); + + // Pointers to input and output device buffers to pass to engine. + // Engine requires exactly IEngine::getNbBindings() number of buffers. + assert(engine.getNbBindings() == 2); + void* buffers[2]; + + // In order to bind the buffers, we need to know the names of the input and output tensors. + // Note that indices are guaranteed to be less than IEngine::getNbBindings() + const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME); + + assert(engine.getBindingDataType(inputIndex) == nvinfer1::DataType::kFLOAT); + const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME); + assert(engine.getBindingDataType(outputIndex) == nvinfer1::DataType::kFLOAT); + int mBatchSize = engine.getMaxBatchSize(); + + // Create GPU buffers on device + CHECK(cudaMalloc(&buffers[inputIndex], 3 * input_shape.height * input_shape.width * sizeof(float))); + CHECK(cudaMalloc(&buffers[outputIndex], output_size*sizeof(float))); + + // Create stream + cudaStream_t stream; + CHECK(cudaStreamCreate(&stream)); + + // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host + CHECK(cudaMemcpyAsync(buffers[inputIndex], input, 3 * input_shape.height * input_shape.width * sizeof(float), cudaMemcpyHostToDevice, stream)); + context.enqueue(1, buffers, stream, nullptr); + CHECK(cudaMemcpyAsync(output, buffers[outputIndex], output_size * sizeof(float), cudaMemcpyDeviceToHost, stream)); + cudaStreamSynchronize(stream); + + // Release stream and buffers + cudaStreamDestroy(stream); + CHECK(cudaFree(buffers[inputIndex])); + CHECK(cudaFree(buffers[outputIndex])); +} + +int main(int argc, char** argv) { + cudaSetDevice(DEVICE); + // create a model using the API directly and serialize it to a stream + char *trtModelStream{nullptr}; + size_t size{0}; + + if (argc == 4 && std::string(argv[2]) == "-i") { + const std::string engine_file_path {argv[1]}; + std::ifstream file(engine_file_path, std::ios::binary); + if (file.good()) { + file.seekg(0, file.end); + size = file.tellg(); + file.seekg(0, file.beg); + trtModelStream = new char[size]; + assert(trtModelStream); + file.read(trtModelStream, size); + file.close(); + } + } else { + std::cerr << "arguments not right!" << std::endl; + std::cerr << "./yolov6 ../model_trt.engine -i ../*.jpg // deserialize file and run inference" << std::endl; + return -1; + } + const std::string input_image_path {argv[3]}; + + //std::vector file_names; + //if (read_files_in_dir(argv[2], file_names) < 0) { + //std::cout << "read_files_in_dir failed." << std::endl; + //return -1; + //} + + IRuntime* runtime = createInferRuntime(gLogger); + assert(runtime != nullptr); + ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size); + assert(engine != nullptr); + IExecutionContext* context = engine->createExecutionContext(); + assert(context != nullptr); + delete[] trtModelStream; + auto out_dims = engine->getBindingDimensions(1); + auto output_size = 1; + for(int j=0;j(end - start).count() << "ms" << std::endl; + + std::vector objects; + decode_outputs(prob, output_size, objects, scale, img_w, img_h); + draw_objects(img, objects, input_image_path); + // delete the pointer to the float + delete blob; + // destroy the engine + context->destroy(); + engine->destroy(); + runtime->destroy(); + return 0; +} + diff --git a/yolov7/trt.py b/yolov7/trt.py new file mode 100644 index 0000000..fb136ab --- /dev/null +++ b/yolov7/trt.py @@ -0,0 +1,34 @@ +import sys +sys.path.append('../') +from utils.utils import preproc, vis +from utils.utils import BaseEngine +import numpy as np +import cv2 +import time +import os + + +class Predictor(BaseEngine): + def __init__(self, engine_path , imgsz=(640,640)): + super(Predictor, self).__init__(engine_path) + self.imgsz = imgsz + self.n_classes = 80 + self.class_names = [ '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 __name__ == '__main__': + pred = Predictor(engine_path='../../yolov7-tiny.trt') + img_path = '../src/3.jpg' + origin_img = pred.inference(img_path, conf=0.3) + cv2.imwrite("%s_yolov7.jpg" % os.path.splitext( + os.path.split(img_path)[-1])[0], origin_img) + pred.detect_video('../src/video1.mp4') # set 0 use a webcam + pred.get_fps()