diff --git a/README.md b/README.md
index cfdb9ec..c993955 100644
--- a/README.md
+++ b/README.md
@@ -113,6 +113,7 @@ Some tricky operations encountered in these models, already solved, but might ha
|-|-|:-:|:-:|:-:|:-:|
| YOLOv3-tiny | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 333 |
| YOLOv3(darknet53) | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 39.2 |
+| YOLOv3(darknet53) | Xeon E5-2620/GTX1080 | 1 | INT8 | 608x608 | 71.4 |
| YOLOv3-spp(darknet53) | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 38.5 |
| YOLOv4(CSPDarknet53) | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 35.7 |
| YOLOv4(CSPDarknet53) | Xeon E5-2620/GTX1080 | 4 | FP32 | 608x608 | 40.9 |
diff --git a/yolov3/CMakeLists.txt b/yolov3/CMakeLists.txt
index d583485..380978f 100644
--- a/yolov3/CMakeLists.txt
+++ b/yolov3/CMakeLists.txt
@@ -13,16 +13,13 @@ find_package(CUDA REQUIRED)
set(CUDA_NVCC_PLAGS ${CUDA_NVCC_PLAGS};-std=c++11;-g;-G;-gencode;arch=compute_30;code=sm_30)
include_directories(${PROJECT_SOURCE_DIR}/include)
-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")
- include_directories(/usr/local/cuda/include)
- link_directories(/usr/local/cuda/lib64)
-endif()
-
+# 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")
@@ -33,7 +30,7 @@ target_link_libraries(yololayer nvinfer cudart)
find_package(OpenCV)
include_directories(OpenCV_INCLUDE_DIRS)
-add_executable(yolov3 ${PROJECT_SOURCE_DIR}/yolov3.cpp)
+add_executable(yolov3 ${PROJECT_SOURCE_DIR}/calibrator.cpp ${PROJECT_SOURCE_DIR}/yolov3.cpp)
target_link_libraries(yolov3 nvinfer)
target_link_libraries(yolov3 cudart)
target_link_libraries(yolov3 yololayer)
diff --git a/yolov3/README.md b/yolov3/README.md
index cab174a..a891f02 100644
--- a/yolov3/README.md
+++ b/yolov3/README.md
@@ -4,33 +4,54 @@ The Pytorch implementation is [ultralytics/yolov3](https://github.com/ultralytic
This branch is using tensorrt7 API, there is also a yolov3 implementation using tensorrt4 API, go to [branch trt4/yolov3](https://github.com/wang-xinyu/tensorrtx/tree/trt4/yolov3), which is using [ayooshkathuria/pytorch-yolo-v3](https://github.com/ayooshkathuria/pytorch-yolo-v3).
-## Excute:
+## Config
+
+- Input shape defined in yololayer.h
+- Number of classes defined in yololayer.h
+- INT8/FP16/FP32 can be selected by the macro in yolov3.cpp
+- GPU id can be selected by the macro in yolov3.cpp
+- NMS thresh in yolov3.cpp
+- BBox confidence thresh in yolov3.cpp
+
+## How to run
-```
1. generate yolov3.wts from pytorch implementation with yolov3.cfg and yolov3.weights, or download .wts from model zoo
+```
git clone https://github.com/wang-xinyu/tensorrtx.git
git clone https://github.com/ultralytics/yolov3.git
// download its weights 'yolov3.pt' or 'yolov3.weights'
-cd yolov3
-cp ../tensorrtx/yolov3/gen_wts.py .
+cp {tensorrtx}/yolov3/gen_wts.py {ultralytics/yolov3/}
+cd {ultralytics/yolov3/}
python gen_wts.py yolov3.weights
// a file 'yolov3.wts' will be generated.
// the master branch of yolov3 should work, if not, you can checkout cf7a4d31d37788023a9186a1a143a2dab0275ead
+```
2. put yolov3.wts into tensorrtx/yolov3, build and run
-mv yolov3.wts ../tensorrtx/yolov3/
-cd ../tensorrtx/yolov3
+```
+mv yolov3.wts {tensorrtx}/yolov3/
+cd {tensorrtx}/yolov3
mkdir build
cd build
cmake ..
make
-sudo ./yolov3 -s // serialize model to plan file i.e. 'yolov3.engine'
-sudo ./yolov3 -d ../../yolov3-spp/samples // deserialize plan file and run inference, the images in samples will be processed.
+sudo ./yolov3 -s // serialize model to plan file i.e. 'yolov3.engine'
+sudo ./yolov3 -d ../../yolov3-spp/samples // deserialize plan file and run inference, the images in samples will be processed.
+```
3. check the images generated, as follows. _zidane.jpg and _bus.jpg
-```
+
+# 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 [BaiduPan](https://pan.baidu.com/s/1GOm_-JobpyLMAqZWCDUhKg) pwd: a9wh
+
+2. unzip it in yolov3/build
+
+3. set the macro `USE_INT8` in yolov3.cpp and make
+
+4. serialize the model and test
@@ -40,15 +61,6 @@ sudo ./yolov3 -d ../../yolov3-spp/samples // deserialize plan file and run infe
-## Config
-
-- Input shape defined in yololayer.h
-- Number of classes defined in yololayer.h
-- FP16/FP32 can be selected by the macro in yolov3.cpp
-- GPU id can be selected by the macro in yolov3.cpp
-- NMS thresh in yolov3.cpp
-- BBox confidence thresh in yolov3.cpp
-
## More Information
See the readme in [home page.](https://github.com/wang-xinyu/tensorrtx)
diff --git a/yolov3/Utils.h b/yolov3/Utils.h
deleted file mode 100644
index 0de663c..0000000
--- a/yolov3/Utils.h
+++ /dev/null
@@ -1,94 +0,0 @@
-#ifndef __TRT_UTILS_H_
-#define __TRT_UTILS_H_
-
-#include
-#include
-#include
-#include
-
-#ifndef CUDA_CHECK
-
-#define CUDA_CHECK(callstr) \
- { \
- cudaError_t error_code = callstr; \
- if (error_code != cudaSuccess) { \
- std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__; \
- assert(0); \
- } \
- }
-
-#endif
-
-namespace Tn
-{
- class Profiler : public nvinfer1::IProfiler
- {
- public:
- void printLayerTimes(int itrationsTimes)
- {
- float totalTime = 0;
- for (size_t i = 0; i < mProfile.size(); i++)
- {
- printf("%-40.40s %4.3fms\n", mProfile[i].first.c_str(), mProfile[i].second / itrationsTimes);
- totalTime += mProfile[i].second;
- }
- printf("Time over all layers: %4.3f\n", totalTime / itrationsTimes);
- }
- private:
- typedef std::pair Record;
- std::vector mProfile;
-
- virtual void reportLayerTime(const char* layerName, float ms)
- {
- auto record = std::find_if(mProfile.begin(), mProfile.end(), [&](const Record& r){ return r.first == layerName; });
- if (record == mProfile.end())
- mProfile.push_back(std::make_pair(layerName, ms));
- else
- record->second += ms;
- }
- };
-
- //Logger for TensorRT info/warning/errors
- class Logger : public nvinfer1::ILogger
- {
- public:
-
- Logger(): Logger(Severity::kWARNING) {}
-
- Logger(Severity severity): reportableSeverity(severity) {}
-
- void log(Severity severity, const char* msg) override
- {
- // suppress messages with severity enum value greater than the reportable
- if (severity > reportableSeverity) return;
-
- switch (severity)
- {
- case Severity::kINTERNAL_ERROR: std::cerr << "INTERNAL_ERROR: "; break;
- case Severity::kERROR: std::cerr << "ERROR: "; break;
- case Severity::kWARNING: std::cerr << "WARNING: "; break;
- case Severity::kINFO: std::cerr << "INFO: "; break;
- default: std::cerr << "UNKNOWN: "; break;
- }
- std::cerr << msg << std::endl;
- }
-
- Severity reportableSeverity{Severity::kWARNING};
- };
-
- template
- void write(char*& buffer, const T& val)
- {
- *reinterpret_cast(buffer) = val;
- buffer += sizeof(T);
- }
-
- template
- void read(const char*& buffer, T& val)
- {
- val = *reinterpret_cast(buffer);
- buffer += sizeof(T);
- }
-}
-
-#endif
\ No newline at end of file
diff --git a/yolov3/calibrator.cpp b/yolov3/calibrator.cpp
new file mode 100644
index 0000000..67b5fb2
--- /dev/null
+++ b/yolov3/calibrator.cpp
@@ -0,0 +1,80 @@
+#include
+#include
+#include
+#include
+#include "calibrator.h"
+#include "cuda_runtime_api.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
+{
+ return batchsize_;
+}
+
+bool Int8EntropyCalibrator2::getBatch(void* bindings[], const char* names[], int nbBindings)
+{
+ if (img_idx_ + batchsize_ > (int)img_files_.size()) {
+ return false;
+ }
+
+ std::vector input_imgs_;
+ for (int i = img_idx_; i < img_idx_ + batchsize_; i++) {
+ std::cout << img_files_[i] << " " << i << std::endl;
+ cv::Mat temp = cv::imread(img_dir_ + img_files_[i]);
+ if (temp.empty()){
+ std::cerr << "Fatal error: image cannot open!" << std::endl;
+ return false;
+ }
+ cv::Mat pr_img = preprocess_img(temp, input_w_, input_h_);
+ input_imgs_.push_back(pr_img);
+ }
+ img_idx_ += batchsize_;
+ cv::Mat blob = cv::dnn::blobFromImages(input_imgs_, 1.0 / 255.0, cv::Size(input_w_, input_h_), cv::Scalar(0, 0, 0), true, false);
+
+ CUDA_CHECK(cudaMemcpy(device_input_, blob.ptr(0), input_count_ * sizeof(float), cudaMemcpyHostToDevice));
+ assert(!strcmp(names[0], input_blob_name_));
+ bindings[0] = device_input_;
+ return true;
+}
+
+const void* Int8EntropyCalibrator2::readCalibrationCache(size_t& length)
+{
+ std::cout << "reading calib cache: " << calib_table_name_ << std::endl;
+ calib_cache_.clear();
+ std::ifstream input(calib_table_name_, std::ios::binary);
+ input >> std::noskipws;
+ if (read_cache_ && input.good())
+ {
+ std::copy(std::istream_iterator(input), std::istream_iterator(), std::back_inserter(calib_cache_));
+ }
+ length = calib_cache_.size();
+ return length ? calib_cache_.data() : nullptr;
+}
+
+void Int8EntropyCalibrator2::writeCalibrationCache(const void* cache, size_t length)
+{
+ std::cout << "writing calib cache: " << calib_table_name_ << " size: " << length << std::endl;
+ std::ofstream output(calib_table_name_, std::ios::binary);
+ output.write(reinterpret_cast(cache), length);
+}
+
diff --git a/yolov3/calibrator.h b/yolov3/calibrator.h
new file mode 100644
index 0000000..1cc9dbb
--- /dev/null
+++ b/yolov3/calibrator.h
@@ -0,0 +1,39 @@
+#ifndef ENTROPY_CALIBRATOR_H
+#define ENTROPY_CALIBRATOR_H
+
+#include "NvInfer.h"
+#include
+#include
+
+//! \class Int8EntropyCalibrator2
+//!
+//! \brief Implements Entropy calibrator 2.
+//! CalibrationAlgoType is kENTROPY_CALIBRATION_2.
+//!
+class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2
+{
+public:
+ Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache = true);
+
+ virtual ~Int8EntropyCalibrator2();
+ int getBatchSize() const override;
+ bool getBatch(void* bindings[], const char* names[], int nbBindings) override;
+ const void* readCalibrationCache(size_t& length) override;
+ void writeCalibrationCache(const void* cache, size_t length) override;
+
+private:
+ int batchsize_;
+ int input_w_;
+ int input_h_;
+ int img_idx_;
+ std::string img_dir_;
+ std::vector img_files_;
+ size_t input_count_;
+ std::string calib_table_name_;
+ const char* input_blob_name_;
+ bool read_cache_;
+ void* device_input_;
+ std::vector calib_cache_;
+};
+
+#endif // ENTROPY_CALIBRATOR_H
diff --git a/yolov3/utils.h b/yolov3/utils.h
new file mode 100644
index 0000000..b0ddfb3
--- /dev/null
+++ b/yolov3/utils.h
@@ -0,0 +1,85 @@
+#ifndef __TRT_UTILS_H_
+#define __TRT_UTILS_H_
+
+#include
+#include
+#include
+#include
+#include
+#include
+
+#ifndef CUDA_CHECK
+
+#define CUDA_CHECK(callstr) \
+ { \
+ cudaError_t error_code = callstr; \
+ if (error_code != cudaSuccess) { \
+ std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__; \
+ assert(0); \
+ } \
+ }
+
+#endif
+
+namespace Tn
+{
+ template
+ void write(char*& buffer, const T& val)
+ {
+ *reinterpret_cast(buffer) = val;
+ buffer += sizeof(T);
+ }
+
+ template
+ void read(const char*& buffer, T& val)
+ {
+ val = *reinterpret_cast(buffer);
+ buffer += sizeof(T);
+ }
+}
+
+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 &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
diff --git a/yolov3/yololayer.cu b/yolov3/yololayer.cu
index 722cb27..ee90aac 100644
--- a/yolov3/yololayer.cu
+++ b/yolov3/yololayer.cu
@@ -1,4 +1,6 @@
#include "yololayer.h"
+#include "utils.h"
+#include
using namespace Yolo;
diff --git a/yolov3/yololayer.h b/yolov3/yololayer.h
index 1f080cb..d1173c3 100644
--- a/yolov3/yololayer.h
+++ b/yolov3/yololayer.h
@@ -1,13 +1,9 @@
#ifndef _YOLO_LAYER_H
#define _YOLO_LAYER_H
-#include
-#include
-#include
-#include
-#include "NvInfer.h"
-#include "Utils.h"
#include
+#include
+#include "NvInfer.h"
namespace Yolo
{
@@ -51,7 +47,6 @@ namespace Yolo
};
}
-
namespace nvinfer1
{
class YoloLayerPlugin: public IPluginV2IOExt
diff --git a/yolov3/yolov3.cpp b/yolov3/yolov3.cpp
index 2bd5a8f..c1c9a03 100644
--- a/yolov3/yolov3.cpp
+++ b/yolov3/yolov3.cpp
@@ -4,26 +4,14 @@
#include
#include
#include
-#include
-#include
#include "NvInfer.h"
#include "cuda_runtime_api.h"
+#include "utils.h"
#include "logging.h"
#include "yololayer.h"
+#include "calibrator.h"
-#define CHECK(status) \
- do\
- {\
- auto ret = (status);\
- if (ret != 0)\
- {\
- std::cerr << "Cuda failure: " << ret << std::endl;\
- abort();\
- }\
- } while (0)
-
-
-#define USE_FP16 // comment out this if want to use FP32
+#define USE_FP16 // set USE_INT8 or USE_FP16 or USE_FP32
#define DEVICE 0 // GPU id
#define NMS_THRESH 0.4
#define BBOX_CONF_THRESH 0.5
@@ -33,33 +21,12 @@ using namespace nvinfer1;
// stuff we know about the network and the input/output blobs
static const int INPUT_H = Yolo::INPUT_H;
static const int INPUT_W = Yolo::INPUT_W;
-static const int OUTPUT_SIZE = 1000 * 7 + 1; // we assume the yololayer outputs no more than 1000 boxes that conf >= 0.1
+static const int DETECTION_SIZE = sizeof(Yolo::Detection) / sizeof(float);
+static const int OUTPUT_SIZE = Yolo::MAX_OUTPUT_BBOX_COUNT * DETECTION_SIZE + 1; // we assume the yololayer outputs no more than MAX_OUTPUT_BBOX_COUNT boxes that conf >= 0.1
const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "prob";
static Logger gLogger;
-cv::Mat preprocess_img(cv::Mat& img) {
- 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_CUBIC);
- 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;
-}
-
cv::Rect get_rect(cv::Mat& img, float bbox[4]) {
int l, r, t, b;
float r_w = INPUT_W / (img.cols * 1.0);
@@ -370,9 +337,16 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilder
// Build engine
builder->setMaxBatchSize(maxBatchSize);
config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB
-#ifdef USE_FP16
+#if defined(USE_FP16)
config->setFlag(BuilderFlag::kFP16);
+#elif defined(USE_INT8)
+ std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl;
+ assert(builder->platformHasFastInt8());
+ config->setFlag(BuilderFlag::kINT8);
+ Int8EntropyCalibrator2 *calibrator = new Int8EntropyCalibrator2(1, INPUT_W, INPUT_H, "./coco_calib/", "int8calib.table", INPUT_BLOB_NAME);
+ config->setInt8Calibrator(calibrator);
#endif
+
std::cout << "Building engine, please wait for a while..." << std::endl;
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
std::cout << "Build engine successfully!" << std::endl;
@@ -420,45 +394,23 @@ void doInference(IExecutionContext& context, float* input, float* output, int ba
const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
// Create GPU buffers on device
- CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
- CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));
+ CUDA_CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
+ CUDA_CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));
// Create stream
cudaStream_t stream;
- CHECK(cudaStreamCreate(&stream));
+ CUDA_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, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
+ CUDA_CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
context.enqueue(batchSize, buffers, stream, nullptr);
- CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
+ CUDA_CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
cudaStreamDestroy(stream);
- CHECK(cudaFree(buffers[inputIndex]));
- CHECK(cudaFree(buffers[outputIndex]));
-}
-
-int read_files_in_dir(const char *p_dir_name, std::vector &file_names) {
- DIR *p_dir = opendir(p_dir_name);
- if (p_dir == nullptr) {
- return -1;
- }
-
- struct dirent* p_file = nullptr;
- while ((p_file = readdir(p_dir)) != nullptr) {
- if (strcmp(p_file->d_name, ".") != 0 &&
- strcmp(p_file->d_name, "..") != 0) {
- //std::string cur_file_name(p_dir_name);
- //cur_file_name += "/";
- //cur_file_name += p_file->d_name;
- std::string cur_file_name(p_file->d_name);
- file_names.push_back(cur_file_name);
- }
- }
-
- closedir(p_dir);
- return 0;
+ CUDA_CHECK(cudaFree(buffers[inputIndex]));
+ CUDA_CHECK(cudaFree(buffers[outputIndex]));
}
int main(int argc, char** argv) {
@@ -522,7 +474,7 @@ int main(int argc, char** argv) {
std::cout << fcount << " " << f << std::endl;
cv::Mat img = cv::imread(std::string(argv[2]) + "/" + f);
if (img.empty()) continue;
- cv::Mat pr_img = preprocess_img(img);
+ cv::Mat pr_img = preprocess_img(img, INPUT_W, INPUT_H);
for (int i = 0; i < INPUT_H * INPUT_W; i++) {
data[i] = pr_img.at(i)[2] / 255.0;
data[i + INPUT_H * INPUT_W] = pr_img.at(i)[1] / 255.0;
@@ -536,16 +488,7 @@ int main(int argc, char** argv) {
std::cout << std::chrono::duration_cast(end - start).count() << "ms" << std::endl;
std::vector res;
nms(res, prob);
- for (int i=0; i<20; i++) {
- std::cout << prob[i] << ",";
- }
- std::cout << res.size() << std::endl;
for (size_t j = 0; j < res.size(); j++) {
- float *p = (float*)&res[j];
- for (size_t k = 0; k < 7; k++) {
- std::cout << p[k] << ", ";
- }
- std::cout << std::endl;
cv::Rect r = get_rect(img, res[j].bbox);
cv::rectangle(img, r, cv::Scalar(0x27, 0xC1, 0x36), 2);
cv::putText(img, std::to_string((int)res[j].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2);