File size: 1,569 Bytes
dc79b9f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | #pragma once
#include <cstddef>
#include <string>
#include <unordered_map>
#include <vector>
#include "ax_engine_api.h"
class EngineWrapper {
public:
EngineWrapper() = default;
~EngineWrapper();
EngineWrapper(const EngineWrapper &) = delete;
EngineWrapper &operator=(const EngineWrapper &) = delete;
int Init(const std::string &model_path);
void Release();
int SetInputByName(const std::string &name, const void *data,
std::size_t size = 0);
int ZeroInputByName(const std::string &name);
int RunSync();
int GetOutputByName(const std::string &name, void *data,
std::size_t size = 0) const;
int CopyOutputToInputByName(const std::string &output_name,
const std::string &input_name);
int GetInputSizeByName(const std::string &name) const;
int GetOutputSizeByName(const std::string &name) const;
const std::string &InputName(std::size_t index) const;
const std::string &OutputName(std::size_t index) const;
std::size_t InputCount() const { return input_names_.size(); }
std::size_t OutputCount() const { return output_names_.size(); }
private:
int InputIndex(const std::string &name) const;
int OutputIndex(const std::string &name) const;
bool initialized_ = false;
AX_ENGINE_HANDLE handle_ = nullptr;
AX_ENGINE_IO_INFO_T *info_ = nullptr;
AX_ENGINE_IO_T io_{};
std::unordered_map<std::string, int> inputs_;
std::unordered_map<std::string, int> outputs_;
std::vector<std::string> input_names_;
std::vector<std::string> output_names_;
};
|