DepthGuide/src/ai/zmq_client.h

#ifndef DEPTHGUIDE_ZMQ_CLIENT_H
#define DEPTHGUIDE_ZMQ_CLIENT_H

#include "core/cuda_helper.hpp"
#include "core/image_utility_v2.h"

#include <boost/container/static_vector.hpp>

#include <nlohmann/json.hpp>

#include <memory>

class zmq_client {
public:

    static constexpr auto max_img_cnt = 2;
    using img_list_type =
            boost::container::static_vector<image_ptr, max_img_cnt>;

    struct request_type {
        img_list_type img_list;
        std::string model_name;
        nlohmann::json extra;
    };

    struct reply_type {
        img_list_type img_list;
        nlohmann::json extra;
    };

    // reply will be call in another thread
    using reply_cb_type = std::function<void(reply_type)>;

    struct create_config {
        std::string python_interpreter;
        std::string server_working_dir;
        std::string server_script_path;
        std::string serv_addr; // like "ipc://fast_sam_v1"
        reply_cb_type reply_cb;
        CUcontext *cuda_ctx = nullptr;
    };

    explicit zmq_client(const create_config &conf);

    ~zmq_client();

    // requests will be dropped if the previous request has not been finished.
    void request(const request_type &req);

    bool is_ideal() const;

    bool is_running() const;

private:
    struct impl;
    std::unique_ptr<impl> pimpl;
};

#endif //DEPTHGUIDE_ZMQ_CLIENT_H