mirror of
https://github.com/JDAI-CV/fast-reid.git
synced 2025-06-03 14:50:47 +08:00
40 lines
1.1 KiB
C
40 lines
1.1 KiB
C
|
#ifndef ENTROPY_CALIBRATOR_H
|
||
|
|
#define ENTROPY_CALIBRATOR_H
|
||
|
|
|
||
|
|
#include "NvInfer.h"
|
||
|
|
#include <string>
|
||
|
|
#include <vector>
|
||
|
|
|
||
|
|
//! \class Int8EntropyCalibrator2
|
||
|
|
//!
|
||
|
|
//! \brief Implements Entropy calibrator 2.
|
||
|
|
//! CalibrationAlgoType is kENTROPY_CALIBRATION_2.
|
||
|
|
//!
|
||
|
|
class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2
|
||
|
|
{
|
||
|
|
public:
|
||
|
|
Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache = true);
|
||
|
|
|
||
|
|
virtual ~Int8EntropyCalibrator2();
|
||
|
|
int getBatchSize() const 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<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
|