37 lines
1006 B
C++
37 lines
1006 B
C++
#pragma once
|
|
|
|
#include "RTOCRTypes.h"
|
|
#include "engine.h"
|
|
#include "engine/EnginePoolManager.h"
|
|
#include <memory>
|
|
#include <mutex>
|
|
|
|
namespace ANSCENTER {
|
|
namespace rtocr {
|
|
|
|
class RTOCRClassifier {
|
|
public:
|
|
RTOCRClassifier() = default;
|
|
~RTOCRClassifier();
|
|
RTOCRClassifier(const RTOCRClassifier&) = delete;
|
|
RTOCRClassifier& operator=(const RTOCRClassifier&) = delete;
|
|
|
|
bool Initialize(const std::string& onnxPath, int gpuId = 0,
|
|
const std::string& engineCacheDir = "");
|
|
|
|
// Classify a batch of text images
|
|
// Returns vector of (cls_label, cls_score) per image
|
|
// cls_label: 0 = normal, 1 = rotated 180 degrees
|
|
std::vector<std::pair<int, float>> Classify(
|
|
const std::vector<cv::Mat>& images, float clsThresh = kClsThresh);
|
|
|
|
private:
|
|
std::shared_ptr<Engine<float>> m_engine = nullptr;
|
|
EnginePoolManager<float>::PoolKey m_poolKey;
|
|
bool m_usingSharedPool = false;
|
|
std::mutex _mutex;
|
|
};
|
|
|
|
} // namespace rtocr
|
|
} // namespace ANSCENTER
|