42 lines
1.1 KiB
C++
42 lines
1.1 KiB
C++
#pragma once
|
|
|
|
#include "RTOCRTypes.h"
|
|
#include "engine.h"
|
|
#include "engine/EnginePoolManager.h"
|
|
#include <memory>
|
|
#include <mutex>
|
|
|
|
namespace ANSCENTER {
|
|
namespace rtocr {
|
|
|
|
class RTOCRRecognizer {
|
|
public:
|
|
RTOCRRecognizer() = default;
|
|
~RTOCRRecognizer();
|
|
RTOCRRecognizer(const RTOCRRecognizer&) = delete;
|
|
RTOCRRecognizer& operator=(const RTOCRRecognizer&) = delete;
|
|
|
|
bool Initialize(const std::string& onnxPath, const std::string& dictPath,
|
|
int gpuId = 0, const std::string& engineCacheDir = "");
|
|
|
|
TextLine Recognize(const cv::Mat& croppedImage);
|
|
std::vector<TextLine> RecognizeBatch(const std::vector<cv::Mat>& croppedImages);
|
|
|
|
void SetRecImageHeight(int h) { imgH_ = h; }
|
|
void SetRecImageMaxWidth(int w) { imgMaxW_ = w; }
|
|
|
|
private:
|
|
TextLine CTCDecode(const float* outputData, int seqLen, int numClasses);
|
|
|
|
std::shared_ptr<Engine<float>> m_engine = nullptr;
|
|
EnginePoolManager<float>::PoolKey m_poolKey;
|
|
bool m_usingSharedPool = false;
|
|
std::vector<std::string> keys_;
|
|
int imgH_ = kRecImgH;
|
|
int imgMaxW_ = kRecImgMaxW;
|
|
std::mutex _mutex;
|
|
};
|
|
|
|
} // namespace rtocr
|
|
} // namespace ANSCENTER
|