thebears/deepstream_yolo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add YOLOv8 support

Browse Source
This commit is contained in:
Marcos Luciano
2023-01-27 15:56:00 -03:00
parent f1cd701247
commit f9c7a4dfca
59 changed files with 3260 additions and 2763 deletions
Download Patch File Download Diff File Expand all files Collapse all files

85
nvdsinfer_custom_impl_Yolo/calibrator.h
View File

@@ -6,57 +6,50 @@
#ifndef CALIBRATOR_H
#define CALIBRATOR_H
#include "opencv2/opencv.hpp"
#include "cuda_runtime.h"
#include "NvInfer.h"
#include <vector>
#include <string>
#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
#include "NvInfer.h"
#include "opencv2/opencv.hpp"
namespace nvinfer1 {
class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2 {
public:
Int8EntropyCalibrator2(const int &batchsize,
const int &channels,
const int &height,
const int &width,
const int &letterbox,
const std::string &imgPath,
const std::string &calibTablePath);
virtual ~Int8EntropyCalibrator2();
int getBatchSize() const noexcept override;
bool getBatch(void* bindings[], const char* names[], int nbBindings) noexcept override;
const void* readCalibrationCache(std::size_t& length) noexcept override;
void writeCalibrationCache(const void* cache, size_t length) noexcept override;
private:
int batchSize;
int inputC;
int inputH;
int inputW;
int letterBox;
std::string calibTablePath;
size_t imageIndex;
size_t inputCount;
std::vector<std::string> imgPaths;
float *batchData{ nullptr };
void *deviceInput{ nullptr };
bool readCache;
std::vector<char> calibrationCache;
};
#define CUDA_CHECK(status) { \
if (status != 0) { \
std::cout << "CUDA failure: " << cudaGetErrorString(status) << " in file " << __FILE__ << " at line " << __LINE__ << \
std::endl; \
abort(); \
} \
}
class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2 {
public:
Int8EntropyCalibrator2(const int& batchsize, const int& channels, const int& height, const int& width,
const int& letterbox, const std::string& imgPath, const std::string& calibTablePath);
virtual ~Int8EntropyCalibrator2();
int getBatchSize() const noexcept override;
bool getBatch(void* bindings[], const char* names[], int nbBindings) noexcept override;
const void* readCalibrationCache(std::size_t& length) noexcept override;
void writeCalibrationCache(const void* cache, size_t length) noexcept override;
private:
int batchSize;
int inputC;
int inputH;
int inputW;
int letterBox;
std::string calibTablePath;
size_t imageIndex;
size_t inputCount;
std::vector<std::string> imgPaths;
float* batchData {nullptr};
void* deviceInput {nullptr};
bool readCache;
std::vector<char> calibrationCache;
};
std::vector<float> prepareImage(cv::Mat& img, int input_c, int input_h, int input_w, int letter_box);
#endif //CALIBRATOR_H