HeadCountAlgo/detector/components/deps/tensorRT_Pro/include/application/yolocrowd.hpp

#ifndef YOLO_CROWD_HPP
#define YOLO_CROWD_HPP

#include "trt_tensor.hpp"
#include "object_detector.hpp"

#include <opencv2/opencv.hpp>

#include <vector>
#include <memory>
#include <string>
#include <future>

namespace YoloCrowd
{
    using namespace std;
    using namespace ObjectDetector;

    enum class Type : int
    {
        V5 = 0,
        X = 1,
        V3 = 2,
        V7 = 3
    };

    enum class NMSMethod : int
    {
        CPU = 0,    // General, for estimate mAP
        FastGPU = 1 // Fast NMS with a small loss of accuracy in corner cases
    };

    void image_to_tensor(const cv::Mat &image, shared_ptr<TRT::Tensor> &tensor, Type type, int ibatch);

    class Infer
    {
    public:
        virtual shared_future<BoxArray> commit(const cv::Mat &image) = 0;
        virtual std::vector<shared_future<BoxArray>> commits(const std::vector<cv::Mat> &images) = 0;
    };

    shared_ptr<Infer> create_infer(
        const string &engine_file, Type type, int gpuid,
        float confidence_threshold = 0.25f, float nms_threshold = 0.5f,
        NMSMethod nms_method = NMSMethod::FastGPU, int max_objects = 1024,
        bool use_multi_preprocess_stream = false);
    const char *type_name(Type type);

}; // namespace YoloCrowd

#endif // YOLO_CROWD_HPP