model.h

Go to the documentation of this file.

// Copyright (c) 2024, DaoAI Robotics. All rights reserved.
// Version: 0.0.0
 
#pragma once
#include "API_EXPORT.h"
#include <string>
#include <map>
#include <filesystem>
#include "common.h"
#include "prediction.h"
#include "config.h"
#include "component_memory.h"
namespace DaoAI
{
 
 namespace DeepLearning
 {
 enum class ModelType
 {
 Object_Detection, Instance_Segmentation, Keypoint_Detection, Classification, OCR, Supervised_Defect_Segmentation, Auto_Segmentation,
 
#ifdef INDUSTRIAL
 Presence_Checking, Positioning, Unsupervised_Defect_Segmentation
#endif // INDUSTRIAL
 };
 
 // Initialize DaoAI Deep Learning SDK
 DAOAI_API void initialize(const bool& use_shared_memory = false, const size_t& reserved_vram_size = 500000000);
 
 DAOAI_API void clearGPUMemory();
 
 DAOAI_API int getNumCUDADevices();
 
 struct ModelInfo
 {
 ModelType model_type;
 DeviceType device;
 std::vector<std::string> class_labels;
 std::vector<int> num_keypoints;
 std::vector<std::string> keypoint_labels;
 std::string extra_infos = "";
 };
 
 class ModelInstance;
 
 class Model
 {
 public:
 Model() = default;
 Model(const std::filesystem::path& model_config_path, const std::filesystem::path& model_weights_path, const DaoAI::DeepLearning::DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 Model(const std::filesystem::path& model_path, const DaoAI::DeepLearning::DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 
 DAOAI_API virtual ~Model();
 
 DAOAI_API ModelInfo getModelInfo();
 
 /*
 * @brief Set the batch size for inference
 * @param batch_size: the batch size for inference
 */
 DAOAI_API void setBatchSize(const size_t& batch_size);
 
 /*
 * @brief Get the batch size for inference
 */
 DAOAI_API size_t getBatchSize() const;
 
 protected:
 std::shared_ptr<ModelInstance> model_instance_ = nullptr;
 };
 
 class UnsupervisedDefectSegmentationImpl;
 
 namespace Vision
 {
 
 
 class Detection : public Model
 {
 public:
 using Model::Model;
 
 /*
 * @brief Set the threshold for how confident the model should be in its predictions
 * @param threshold: the threshold for how confident the model should be in its predictions
 */
 DAOAI_API void setConfidenceThreshold(const float& threshold);
 
 /*
 * @brief Get the threshold for how confident the model should be in its predictions
 */
 DAOAI_API float getConfidenceThreshold() const;
 
 /*
 * @brief Set the threshold for the IOU between the predicted bounding box and the ground truth bounding box
 * @param threshold: the threshold for the IOU between the predicted bounding box and the ground truth bounding box
 */
 DAOAI_API void setIOUThreshold(const float& threshold);
 
 /*
 * @brief Get the threshold for the IOU between the predicted bounding box and the ground truth bounding box
 */
 DAOAI_API float getIOUThreshold() const;
 };
 
 class ObjectDetection final : public Detection
 {
 public:
 DAOAI_API ObjectDetection(const std::filesystem::path& model_file, const DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 DAOAI_API ObjectDetection(const std::filesystem::path& model_config_path, const std::filesystem::path& model_weights_path, const DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 DAOAI_API ObjectDetectionResult inference(const Image& image);
 DAOAI_API std::vector<ObjectDetectionResult> inference(const std::vector<Image>& images);
 };
 class InstanceSegmentation final : public Detection
 {
 public:
 DAOAI_API InstanceSegmentation(const std::filesystem::path& model_file, const DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 DAOAI_API InstanceSegmentation(const std::filesystem::path& model_config_path, const std::filesystem::path& model_weights_path, const DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 DAOAI_API InstanceSegmentationResult inference(const Image& image);
 DAOAI_API std::vector<InstanceSegmentationResult> inference(const std::vector<Image>& images);
 };
 class KeypointDetection final : public Detection
 {
 public:
 DAOAI_API KeypointDetection(const std::filesystem::path& model_file, const DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 DAOAI_API KeypointDetection(const std::filesystem::path& model_config_path, const std::filesystem::path& model_weights_path, const DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 DAOAI_API KeypointDetectionResult inference(const Image& image);
 DAOAI_API std::vector<KeypointDetectionResult> inference(const std::vector<Image>& images);
 };
 class OCR final : public Model
 {
 public:
 DAOAI_API OCR(const std::filesystem::path& model_file, const DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 DAOAI_API OCR(const std::filesystem::path& model_config_path, const std::filesystem::path& model_weights_path, const DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 DAOAI_API OCRResult inference(const Image& image);
 DAOAI_API std::vector<OCRResult> inference(const std::vector<Image>& images);
 };
 class Classification final : public Model
 {
 public:
 DAOAI_API Classification(const std::filesystem::path& model_file, const DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 DAOAI_API Classification(const std::filesystem::path& model_config_path, const std::filesystem::path& model_weights_path, const DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 DAOAI_API ClassificationResult inference(const Image& image);
 DAOAI_API std::vector<ClassificationResult> inference(const std::vector<Image>& images);
 };
 class SupervisedDefectSegmentation final : public Model
 {
 public:
 DAOAI_API SupervisedDefectSegmentation(const std::filesystem::path& model_file, const DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 DAOAI_API SupervisedDefectSegmentation(const std::filesystem::path& model_config_path, const std::filesystem::path& model_weights_path, const DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 DAOAI_API SupervisedDefectSegmentationResult inference(const Image& image);
 DAOAI_API std::vector<SupervisedDefectSegmentationResult> inference(const std::vector<Image>& images);
 };
 class AutoSegmentation final : public Model
 {
 public:
 DAOAI_API AutoSegmentation(const std::filesystem::path& model_file, const DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 DAOAI_API std::vector<ImageEmbedding> generateImageEmbeddings(const std::vector<Image>& images);
 DAOAI_API ImageEmbedding generateImageEmbeddings(const Image& image);
 DAOAI_API std::vector<AutoSegmentationResult> inference(const std::vector<ImageEmbedding>& embeddings, const std::vector<std::vector<Box>>& boxes, const std::vector<std::vector<Point>>& points = std::vector<std::vector<Point>>());
 DAOAI_API std::vector<AutoSegmentationResult> inference(const std::vector<ImageEmbedding>& embeddings, const std::vector<std::vector<Point>>& points);
 DAOAI_API AutoSegmentationResult inference(const ImageEmbedding& embedding, const std::vector<Box>& boxes, const std::vector<Point>& points = std::vector<Point>());
 DAOAI_API AutoSegmentationResult inference(const ImageEmbedding& embedding, const std::vector<Point>& points);
 };
 
#ifdef INDUSTRIAL
 class PresenceChecking final : public Detection
 {
 public:
 DAOAI_API PresenceChecking(const std::filesystem::path& model_file, const DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 DAOAI_API PresenceChecking(const std::filesystem::path& model_config_path, const std::filesystem::path& model_weights_path, const DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 DAOAI_API PresenceCheckingResult inference(const Image& image);
 DAOAI_API std::vector<PresenceCheckingResult> inference(const std::vector<Image>& images);
 };
 
 class Positioning final : public Detection
 {
 public:
 DAOAI_API Positioning(const std::filesystem::path& model_file, const DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 DAOAI_API Positioning(const std::filesystem::path& model_config_path, const std::filesystem::path& model_weights_path, const DeviceType& device = DeviceType::GPU, const int& device_idx = -1);
 DAOAI_API PositioningResult inference(const Image& image);
 DAOAI_API std::vector<PositioningResult> inference(const std::vector<Image>& images);
 };
 
 
 class UnsupervisedDefectSegmentation
 {
 public:
 DAOAI_API UnsupervisedDefectSegmentation(const DeviceType& device);
 
 DAOAI_API ~UnsupervisedDefectSegmentation();
 
 DAOAI_API UnsupervisedDefectSegmentationResult inference(
 const Image& image, const std::string& component_name, float sensitivity);
 
 DAOAI_API std::vector<UnsupervisedDefectSegmentationResult> inference(
 const std::vector<Image>& images,
 const std::vector<std::string>& component_names,
 const std::vector<float>& sensitivities);
 
 DAOAI_API std::vector<UnsupervisedDefectSegmentationResult> inference(const std::vector<Image>& images);
 
 DAOAI_API UnsupervisedDefectSegmentationResult inference(const Image& image);
 
 DAOAI_API void setBatchSize(const size_t& batch_size);
 
 DAOAI_API size_t getBatchSize() const;
 
 DAOAI_API void addComponentMemory(const std::string& component_name, const std::filesystem::path& component_memory_file_path);
 
 DAOAI_API void addComponentMemory(const std::string& component_name, const ComponentMemory& component_memory);
 
 DAOAI_API void addComponentArchive(const std::string& archive_path);
 
 DAOAI_API void removeComponentMemory(const std::string& component_name);
 
 DAOAI_API std::unordered_map<std::string, ComponentMemory> listComponentMemory();
 
 DAOAI_API void setDetectionLevel(const DetectionLevel& value);
 
 DAOAI_API ComponentMemory createComponentMemory(const std::string& component_name, const std::vector<Image>& good_images, const std::vector<Image>& bad_images = {}, const std::vector<Image>& masks = {}, const bool& append_memory = false);
 
 private:
 std::shared_ptr<UnsupervisedDefectSegmentationImpl> impl_ = nullptr;
 };
#endif // INDUSTRIAL
 }
 }
}