InferenceEngine.hpp

Go to the documentation of this file.

#pragma once
 
#include "FAST/Data/DataTypes.hpp"
#include <FAST/Data/Tensor.hpp>
#include <FAST/Data/TensorShape.hpp>
 
// This is a macro for creating a load function for a given inference engine
// Need C linkage here (extern "C" to avoid mangled names of the load function on windows, see https://stackoverflow.com/questions/19422550/why-getprocaddress-is-not-working
// Export statement is needed for Windows
#define DEFINE_INFERENCE_ENGINE(classType, exportStatement)                                   \
extern "C" exportStatement                                                          \
InferenceEngine* load() {                                                                  \
      return new classType();                                                       \
}                                                                                      \
 
namespace fast {
 
enum class ImageOrdering {
    ChannelFirst,
    ChannelLast
};
 
enum class NodeType {
    IMAGE,
    TENSOR,
};
 
enum class InferenceDeviceType {
    ANY,
    CPU,
    GPU,
    VPU,
    OTHER,
};
 
struct InferenceDeviceInfo {
    std::string name;
    InferenceDeviceType type;
    int index;
};
 
enum class ModelFormat {
    PROTOBUF,
    SAVEDMODEL,
    ONNX,
    OPENVINO,
    UFF
};
 
FAST_EXPORT std::string getModelFileExtension(ModelFormat format);
 
FAST_EXPORT ModelFormat getModelFormat(std::string filename);
 
FAST_EXPORT std::string getModelFormatName(ModelFormat format);
 
class FAST_EXPORT InferenceEngine : public Object {
    public:
        typedef std::shared_ptr<InferenceEngine> pointer;
        struct NetworkNode {
            uint portID;
            NodeType type;
            TensorShape shape;
            std::shared_ptr<Tensor> data;
        };
        virtual void setFilename(std::string filename);
        virtual void setModelAndWeights(std::vector<uint8_t> model, std::vector<uint8_t> weights);
        virtual std::string getFilename() const;
        virtual void run() = 0;
        virtual void addInputNode(uint portID, std::string name, NodeType type = NodeType::IMAGE, TensorShape shape = {});
        virtual void addOutputNode(uint portID, std::string name, NodeType type = NodeType::IMAGE, TensorShape shape = {});
        virtual void setInputNodeShape(std::string name, TensorShape shape);
        virtual void setOutputNodeShape(std::string name, TensorShape shape);
        virtual NetworkNode getInputNode(std::string name) const;
        virtual NetworkNode getOutputNode(std::string name) const;
        virtual std::unordered_map<std::string, NetworkNode> getOutputNodes() const;
        virtual std::unordered_map<std::string, NetworkNode> getInputNodes() const;
        virtual void setInputData(std::string inputNodeName, std::shared_ptr<Tensor> tensor);
        virtual std::shared_ptr<Tensor> getOutputData(std::string inputNodeName);
        virtual void load() = 0;
        virtual bool isLoaded() const;
        virtual ImageOrdering getPreferredImageOrdering() const = 0;
        virtual std::string getName() const = 0;
        virtual std::vector<ModelFormat> getSupportedModelFormats() const = 0;
        virtual ModelFormat getPreferredModelFormat() const = 0;
        virtual bool isModelFormatSupported(ModelFormat format);
        virtual void setDeviceType(InferenceDeviceType type);
        virtual void setDevice(int index = -1, InferenceDeviceType type = InferenceDeviceType::ANY);
        virtual std::vector<InferenceDeviceInfo> getDeviceList();
 
        virtual int getMaxBatchSize();
        virtual void setMaxBatchSize(int size);
    protected:
        virtual void setIsLoaded(bool loaded);
 
        std::unordered_map<std::string, NetworkNode> mInputNodes;
        std::unordered_map<std::string, NetworkNode> mOutputNodes;
 
        int m_deviceIndex = -1;
        InferenceDeviceType m_deviceType = InferenceDeviceType::ANY;
        int m_maxBatchSize = 1;
 
        std::vector<uint8_t> m_model;
        std::vector<uint8_t> m_weights;
    private:
        std::string m_filename = "";
        bool m_isLoaded = false;
};
 
}