ANSCORE/modules/ANSOCR/ANSONNXOCR/ONNXOCRClassifier.h

#pragma once

#include "ONNXOCRTypes.h"
#include "ONNXEngine.h"
#include <vector>
#include <mutex>

namespace ANSCENTER {
namespace onnxocr {

class ONNXOCRClassifier : public BasicOrtHandler {
public:
    explicit ONNXOCRClassifier(const std::string& onnx_path, unsigned int num_threads = 1);
    explicit ONNXOCRClassifier(const std::string& onnx_path,
                               const OrtHandlerOptions& options,
                               unsigned int num_threads = 1);
    ~ONNXOCRClassifier() override = default;

    // Classify text orientation for a list of cropped images
    // Returns vector of (cls_label, cls_score) pairs
    // cls_label: 0 = normal, 1 = rotated 180 degrees
    void Classify(std::vector<cv::Mat>& img_list,
                  std::vector<int>& cls_labels,
                  std::vector<float>& cls_scores,
                  float cls_thresh = kClsThresh);

    // Pre-warm cuDNN/TRT for the classifier's fixed [1,3,80,160] shape.
    // Idempotent — no-op after the first call.
    void Warmup();

private:
    bool _warmedUp = false;
    Ort::Value transform(const cv::Mat& mat) override;
    Ort::Value transformBatch(const std::vector<cv::Mat>& images) override;

    std::mutex _mutex;
};

} // namespace onnxocr
} // namespace ANSCENTER