model.h

#pragma once
#include "API_EXPORT.h"
#include <string>
#include "common.h"
#include "prediction.h"
 
namespace DaoAI
{
    
    namespace DeepLearning
    {
        enum class ModelType
        {
            Object_Detection, Classification, Anomaly_Detection, OCR, Semantic_Segmentation
        };
 
        enum DetectionTask
        {
            Unknown = 0, Bounding_Box = 1, Instance_Segmentation = 2, Keypoint = 4, Rotated_Box = 8
        };
 
        class BaseModel;
 
        // Initialize DaoAI Deep Learning SDK
        DAOAI_API void initialize();
 
        DAOAI_API bool checkDaoAIModelValidity(const std::string& model_config_path, const std::string& model_weights_path);
 
        DAOAI_API bool checkDaoAIModelValidity(const std::string& model_path);
 
        DAOAI_API int getNumCUDADevices();
 
        class ModelInfo
        {
        public:
            ModelType model_type;
            Device_Type device;
            std::vector<std::string> class_labels;
            std::vector<int> num_keypoints;
            DetectionTask tasks;
        };
 
        class Model
        {
        public:
            DAOAI_API Model();
            DAOAI_API ~Model();
 
            DAOAI_API void load(const std::string& model_config_path, const std::string& model_weights_path, const DaoAI::DeepLearning::Device_Type& device = Device_Type::GPU, const int& device_idx = -1);
 
            DAOAI_API void load(const std::string& model_directory, const DaoAI::DeepLearning::Device_Type& device = Device_Type::GPU, const int& device_idx = -1);
            
            DAOAI_API void loadZip(const std::string& model_zip_path, const DaoAI::DeepLearning::Device_Type& device = Device_Type::GPU, const int& device_idx = -1);
 
            DAOAI_API void loadNestedZip(const std::string& model_zip_path, const DaoAI::DeepLearning::Device_Type& device = Device_Type::GPU, const int& device_idx = -1);
 
            DAOAI_API void unload();
 
            DAOAI_API Prediction inference(const Image& input_image);
 
            DAOAI_API std::string inferenceJSON(const Image& input_image);
 
            DAOAI_API ModelInfo getModelInfo();
 
 
        private:
            std::string model_root_path_ = "";
            std::string model_config_path_ = "";
            std::string model_weights_path_ = "";
            Device_Type device_ = Device_Type::GPU;
            int gpu_index_ = -1;
            int instance_index_;
            std::unique_ptr<BaseModel> model_ptr_ = nullptr;
        };
    }
}