From a8b2cea21137986000664bf74508dbf271995ab8 Mon Sep 17 00:00:00 2001 From: kuba Date: Sun, 22 Sep 2024 14:53:53 -0400 Subject: [PATCH 1/2] Add model loading support, Python bindings, and improve OpenVINO const correctness. --- CMakeLists.txt | 5 + bindings/python/src/modelzoo/ZooBindings.cpp | 46 ++- .../pipeline/node/NeuralNetworkBindings.cpp | 5 + examples/python/RVC2/NNArchive/nn_archive.py | 13 +- examples/python/t | 268 ++++++++++++++++++ include/depthai/models/ModelLoader.hpp | 52 ++++ include/depthai/models/ModelOptimizer.hpp | 0 include/depthai/models/ModelSerializer.hpp | 0 include/depthai/models/ModelZoo.hpp | 18 ++ include/depthai/models/Models.hpp | 101 +++++++ include/depthai/openvino/OpenVINO.hpp | 10 +- .../pipeline/node/DetectionNetwork.hpp | 17 ++ .../depthai/pipeline/node/NeuralNetwork.hpp | 15 + src/models/ModelLoader.cpp | 204 +++++++++++++ src/models/Models.cpp | 12 + src/models/main.cpp | 15 + src/openvino/OpenVINO.cpp | 10 +- 17 files changed, 779 insertions(+), 12 deletions(-) create mode 100644 examples/python/t create mode 100644 include/depthai/models/ModelLoader.hpp create mode 100644 include/depthai/models/ModelOptimizer.hpp create mode 100644 include/depthai/models/ModelSerializer.hpp create mode 100644 include/depthai/models/ModelZoo.hpp create mode 100644 include/depthai/models/Models.hpp create mode 100644 src/models/ModelLoader.cpp create mode 100644 src/models/Models.cpp create mode 100644 src/models/main.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 7720632358..9dd7026d6e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -382,8 +382,13 @@ set(TARGET_CORE_SOURCES src/nn_archive/NNArchiveConfig.cpp src/modelzoo/NNModelDescription.cpp src/modelzoo/Zoo.cpp + src/models/Models.cpp + src/models/ModelLoader.cpp ) +add_executable(model_test src/models/main.cpp) +target_link_libraries(model_test PRIVATE ${TARGET_CORE_NAME}) + set(TARGET_OPENCV_SOURCES src/opencv/ImgFrame.cpp src/pipeline/node/host/Display.cpp diff --git a/bindings/python/src/modelzoo/ZooBindings.cpp b/bindings/python/src/modelzoo/ZooBindings.cpp index 0209060748..508534dbcb 100644 --- a/bindings/python/src/modelzoo/ZooBindings.cpp +++ b/bindings/python/src/modelzoo/ZooBindings.cpp @@ -2,6 +2,11 @@ // depthai #include "depthai/modelzoo/Zoo.hpp" +#include +#include +#include + +// Model bindings void ZooBindings::bind(pybind11::module& m, void* pCallstack) { using namespace dai; @@ -32,4 +37,43 @@ void ZooBindings::bind(pybind11::module& m, void* pCallstack) { py::arg("cacheDirectory") = MODEL_ZOO_DEFAULT_CACHE_DIRECTORY, py::arg("apiKey") = "", DOC(dai, downloadModelsFromZoo)); -} \ No newline at end of file + + // Create a new submodule + py::module m_model = m.def_submodule("model"); + py::module m_zoo = m_model.def_submodule("zoo"); + + // Bind Model + py::class_>(m_model, "Model"); + py::enum_(m_model, "ModelType") + .value("BLOB", depthai::model::ModelType::BLOB) + .value("SUPERBLOB", depthai::model::ModelType::SUPERBLOB) + .value("DLC", depthai::model::ModelType::DLC) + .value("NNARCHIVE", depthai::model::ModelType::NNARCHIVE) + .export_values(); + + // Bind ModelSettings + py::class_>(m_model, "ModelSettings"); + + // Bind BlobSettings + py::class_>(m_model, "BlobSettings"); + + // Bind SuperBlobSettings + py::class_>(m_model, "SuperBlobSettings"); + + // Bind DlcSettings + py::class_>(m_model, "DlcSettings"); + + // Bind BlobModel + py::class_>(m_model, "BlobModel"); + + // Bind SuperBlobModel + py::class_>(m_model, "SuperBlobModel"); + + // Bind DlcModel + py::class_>(m_model, "DlcModel"); + + // Bind load function + m_model.def("load", py::overload_cast(&depthai::model::load), py::arg("path")); + m_zoo.def("load", &depthai::model::zoo::load, py::arg("ModelDescription")); + +} diff --git a/bindings/python/src/pipeline/node/NeuralNetworkBindings.cpp b/bindings/python/src/pipeline/node/NeuralNetworkBindings.cpp index c6288efc70..04028d3277 100644 --- a/bindings/python/src/pipeline/node/NeuralNetworkBindings.cpp +++ b/bindings/python/src/pipeline/node/NeuralNetworkBindings.cpp @@ -5,6 +5,8 @@ #include "depthai/pipeline/Node.hpp" #include "depthai/pipeline/node/NeuralNetwork.hpp" +#include "depthai/models/Models.hpp" + void bind_neuralnetwork(pybind11::module& m, void* pCallstack){ @@ -44,6 +46,9 @@ void bind_neuralnetwork(pybind11::module& m, void* pCallstack){ .def("setBlobPath", &NeuralNetwork::setBlobPath, py::arg("path"), DOC(dai, node, NeuralNetwork, setBlobPath)) .def("setNumPoolFrames", &NeuralNetwork::setNumPoolFrames, py::arg("numFrames"), DOC(dai, node, NeuralNetwork, setNumPoolFrames)) .def("setNumInferenceThreads", &NeuralNetwork::setNumInferenceThreads, py::arg("numThreads"), DOC(dai, node, NeuralNetwork, setNumInferenceThreads)) + .def("setModel", py::overload_cast(&NeuralNetwork::setModel), py::arg("model")) + .def("setModel", py::overload_cast(&NeuralNetwork::setModel), py::arg("model")) + .def("setModel", py::overload_cast(&NeuralNetwork::setModel), py::arg("model")) .def("setNumNCEPerInferenceThread", &NeuralNetwork::setNumNCEPerInferenceThread, py::arg("numNCEPerThread"), diff --git a/examples/python/RVC2/NNArchive/nn_archive.py b/examples/python/RVC2/NNArchive/nn_archive.py index a1d62aff79..02e9f9adf9 100644 --- a/examples/python/RVC2/NNArchive/nn_archive.py +++ b/examples/python/RVC2/NNArchive/nn_archive.py @@ -6,7 +6,6 @@ import time # Get argument first -modelDescription = dai.NNModelDescription(modelSlug="yolov6-nano", platform="RVC2") archivePath = dai.getModelFromZoo(modelDescription, useCached=True) # Create pipeline @@ -21,6 +20,18 @@ nnArchive = dai.NNArchive(archivePath) h, w = nnArchive.getConfig().getConfigV1().model.inputs[0].shape[-2:] camRgb.setPreviewSize(w, h) + + """ + # load model + modelDescription = dai.NNModelDescription(modelSlug="yolov6-nano", platform="RVC2") + model = dai.model.zoo.load(modelDescription) + settings = model.settings() + + # Set model + detectionNetwork.setModel + """ + + detectionNetwork = pipeline.create(dai.node.DetectionNetwork).build( camRgb.preview, nnArchive ) diff --git a/examples/python/t b/examples/python/t new file mode 100644 index 0000000000..dba38ed4f9 --- /dev/null +++ b/examples/python/t @@ -0,0 +1,268 @@ +diff --git a/CMakeLists.txt b/CMakeLists.txt +index 772063235..9dd7026d6 100644 +--- a/CMakeLists.txt ++++ b/CMakeLists.txt +@@ -382,8 +382,13 @@ set(TARGET_CORE_SOURCES + src/nn_archive/NNArchiveConfig.cpp + src/modelzoo/NNModelDescription.cpp + src/modelzoo/Zoo.cpp ++ src/models/Models.cpp ++ src/models/ModelLoader.cpp + ) + ++add_executable(model_test src/models/main.cpp) ++target_link_libraries(model_test PRIVATE ${TARGET_CORE_NAME}) ++ + set(TARGET_OPENCV_SOURCES + src/opencv/ImgFrame.cpp + src/pipeline/node/host/Display.cpp +diff --git a/bindings/python/src/modelzoo/ZooBindings.cpp b/bindings/python/src/modelzoo/ZooBindings.cpp +index 020906074..508534dbc 100644 +--- a/bindings/python/src/modelzoo/ZooBindings.cpp ++++ b/bindings/python/src/modelzoo/ZooBindings.cpp +@@ -2,6 +2,11 @@ + + // depthai + #include "depthai/modelzoo/Zoo.hpp" ++#include ++#include ++#include ++ ++// Model bindings + + void ZooBindings::bind(pybind11::module& m, void* pCallstack) { + using namespace dai; +@@ -32,4 +37,43 @@ void ZooBindings::bind(pybind11::module& m, void* pCallstack) { + py::arg("cacheDirectory") = MODEL_ZOO_DEFAULT_CACHE_DIRECTORY, + py::arg("apiKey") = "", + DOC(dai, downloadModelsFromZoo)); +-} +\ No newline at end of file ++ ++ // Create a new submodule ++ py::module m_model = m.def_submodule("model"); ++ py::module m_zoo = m_model.def_submodule("zoo"); ++ ++ // Bind Model ++ py::class_>(m_model, "Model"); ++ py::enum_(m_model, "ModelType") ++ .value("BLOB", depthai::model::ModelType::BLOB) ++ .value("SUPERBLOB", depthai::model::ModelType::SUPERBLOB) ++ .value("DLC", depthai::model::ModelType::DLC) ++ .value("NNARCHIVE", depthai::model::ModelType::NNARCHIVE) ++ .export_values(); ++ ++ // Bind ModelSettings ++ py::class_>(m_model, "ModelSettings"); ++ ++ // Bind BlobSettings ++ py::class_>(m_model, "BlobSettings"); ++ ++ // Bind SuperBlobSettings ++ py::class_>(m_model, "SuperBlobSettings"); ++ ++ // Bind DlcSettings ++ py::class_>(m_model, "DlcSettings"); ++ ++ // Bind BlobModel ++ py::class_>(m_model, "BlobModel"); ++ ++ // Bind SuperBlobModel ++ py::class_>(m_model, "SuperBlobModel"); ++ ++ // Bind DlcModel ++ py::class_>(m_model, "DlcModel"); ++ ++ // Bind load function ++ m_model.def("load", py::overload_cast(&depthai::model::load), py::arg("path")); ++ m_zoo.def("load", &depthai::model::zoo::load, py::arg("ModelDescription")); ++ ++} +diff --git a/bindings/python/src/pipeline/node/NeuralNetworkBindings.cpp b/bindings/python/src/pipeline/node/NeuralNetworkBindings.cpp +index c6288efc7..04028d327 100644 +--- a/bindings/python/src/pipeline/node/NeuralNetworkBindings.cpp ++++ b/bindings/python/src/pipeline/node/NeuralNetworkBindings.cpp +@@ -5,6 +5,8 @@ + #include "depthai/pipeline/Node.hpp" + #include "depthai/pipeline/node/NeuralNetwork.hpp" + ++#include "depthai/models/Models.hpp" ++ + + void bind_neuralnetwork(pybind11::module& m, void* pCallstack){ + +@@ -44,6 +46,9 @@ void bind_neuralnetwork(pybind11::module& m, void* pCallstack){ + .def("setBlobPath", &NeuralNetwork::setBlobPath, py::arg("path"), DOC(dai, node, NeuralNetwork, setBlobPath)) + .def("setNumPoolFrames", &NeuralNetwork::setNumPoolFrames, py::arg("numFrames"), DOC(dai, node, NeuralNetwork, setNumPoolFrames)) + .def("setNumInferenceThreads", &NeuralNetwork::setNumInferenceThreads, py::arg("numThreads"), DOC(dai, node, NeuralNetwork, setNumInferenceThreads)) ++ .def("setModel", py::overload_cast(&NeuralNetwork::setModel), py::arg("model")) ++ .def("setModel", py::overload_cast(&NeuralNetwork::setModel), py::arg("model")) ++ .def("setModel", py::overload_cast(&NeuralNetwork::setModel), py::arg("model")) + .def("setNumNCEPerInferenceThread", + &NeuralNetwork::setNumNCEPerInferenceThread, + py::arg("numNCEPerThread"), +diff --git a/examples/python/RVC2/NNArchive/nn_archive.py b/examples/python/RVC2/NNArchive/nn_archive.py +index a1d62aff7..02e9f9adf 100644 +--- a/examples/python/RVC2/NNArchive/nn_archive.py ++++ b/examples/python/RVC2/NNArchive/nn_archive.py +@@ -6,7 +6,6 @@ import numpy as np + import time + + # Get argument first +-modelDescription = dai.NNModelDescription(modelSlug="yolov6-nano", platform="RVC2") + archivePath = dai.getModelFromZoo(modelDescription, useCached=True) + + # Create pipeline +@@ -21,6 +20,18 @@ with dai.Pipeline() as pipeline: + nnArchive = dai.NNArchive(archivePath) + h, w = nnArchive.getConfig().getConfigV1().model.inputs[0].shape[-2:] + camRgb.setPreviewSize(w, h) ++ ++ """ ++ # load model ++ modelDescription = dai.NNModelDescription(modelSlug="yolov6-nano", platform="RVC2") ++ model = dai.model.zoo.load(modelDescription) ++ settings = model.settings() ++ ++ # Set model ++ detectionNetwork.setModel ++ """ ++ ++ + detectionNetwork = pipeline.create(dai.node.DetectionNetwork).build( + camRgb.preview, nnArchive + ) +diff --git a/include/depthai/openvino/OpenVINO.hpp b/include/depthai/openvino/OpenVINO.hpp +index 6e4218301..49dc69854 100644 +--- a/include/depthai/openvino/OpenVINO.hpp ++++ b/include/depthai/openvino/OpenVINO.hpp +@@ -82,7 +82,7 @@ class OpenVINO { + * @param numShaves: Number of shaves to generate the blob for. Must be between 1 and NUMBER_OF_PATCHES. + * @return dai::OpenVINO::Blob: Blob compiled for the specified number of shaves + */ +- dai::OpenVINO::Blob getBlobWithNumShaves(int numShaves); ++ dai::OpenVINO::Blob getBlobWithNumShaves(int numShaves) const; + + private: + // A header in the superblob containing metadata about the blob and patches +@@ -99,16 +99,16 @@ class OpenVINO { + std::vector readSuperBlobFile(const std::string& path); + + // Get a pointer to the first byte of the blob data +- const uint8_t* getBlobDataPointer(); ++ const uint8_t* getBlobDataPointer() const; + + // Get the size in bytes of the blob data +- int64_t getBlobDataSize(); ++ int64_t getBlobDataSize() const; + + // Get a pointer to the first byte of the patch data for a specific number of shaves +- const uint8_t* getPatchDataPointer(int numShaves); ++ const uint8_t* getPatchDataPointer(int numShaves) const; + + // Get the size in bytes of the patch data for a specific number of shaves +- int64_t getPatchDataSize(int numShaves); ++ int64_t getPatchDataSize(int numShaves) const; + + // Load header - throw if an error occurs + void loadAndCheckHeader(); +diff --git a/include/depthai/pipeline/node/DetectionNetwork.hpp b/include/depthai/pipeline/node/DetectionNetwork.hpp +index c2ee2ce35..ddd2dcf35 100644 +--- a/include/depthai/pipeline/node/DetectionNetwork.hpp ++++ b/include/depthai/pipeline/node/DetectionNetwork.hpp +@@ -28,6 +28,23 @@ class DetectionNetwork : public DeviceNodeGroup { + return networkPtr; + } + ++ ++ void setModel(const depthai::model::ModelVariant& model) { ++ std::visit([this](auto &&p){this->setModel(p);}, model); ++ } ++ ++ void setModel(const depthai::model::BlobModel& model) { ++ neuralNetwork->setModel(model); ++ } ++ ++ void setModel(const depthai::model::SuperBlobModel& model) { ++ neuralNetwork->setModel(model); ++ } ++ ++ void setModel(const depthai::model::DlcModel& model) { ++ neuralNetwork->setModel(model); ++ } ++ + std::shared_ptr build(Node::Output& input, const NNArchive& nnArchive); + std::shared_ptr build(std::shared_ptr input, dai::NNModelDescription modelDesc, float fps = 30.0f); + +diff --git a/include/depthai/pipeline/node/NeuralNetwork.hpp b/include/depthai/pipeline/node/NeuralNetwork.hpp +index 1cdba3298..32428eddc 100644 +--- a/include/depthai/pipeline/node/NeuralNetwork.hpp ++++ b/include/depthai/pipeline/node/NeuralNetwork.hpp +@@ -1,5 +1,6 @@ + #pragma once + ++#include + #include + #include + #include +@@ -74,6 +75,20 @@ class NeuralNetwork : public DeviceNodeCRTPsetModel(p);}, model); ++ } ++ ++ void setModel(const depthai::model::BlobModel& model) { ++ setBlob(model.model()); ++ } ++ void setModel(const depthai::model::SuperBlobModel& model) { ++ setBlob(model.model().getBlobWithNumShaves(model.settings().numShaves)); ++ } ++ void setModel(const depthai::model::DlcModel& model) { ++ // pass ++ } ++ + /** + * @brief Set NNArchive for this Node, throws if the archive's type is not SUPERBLOB + * +diff --git a/src/openvino/OpenVINO.cpp b/src/openvino/OpenVINO.cpp +index e52553b56..834181a00 100644 +--- a/src/openvino/OpenVINO.cpp ++++ b/src/openvino/OpenVINO.cpp +@@ -99,7 +99,7 @@ OpenVINO::SuperBlob::SuperBlob(std::vector data) { + validateSuperblob(); + } + +-dai::OpenVINO::Blob OpenVINO::SuperBlob::getBlobWithNumShaves(int numShaves) { ++dai::OpenVINO::Blob OpenVINO::SuperBlob::getBlobWithNumShaves(int numShaves) const{ + if(numShaves < 1 || numShaves > static_cast(OpenVINO::SuperBlob::NUMBER_OF_PATCHES)) { + throw std::out_of_range("Invalid number of shaves: " + std::to_string(numShaves) + " (expected 1 to " + + std::to_string(OpenVINO::SuperBlob::NUMBER_OF_PATCHES) + ")"); +@@ -162,22 +162,22 @@ OpenVINO::SuperBlob::SuperBlobHeader OpenVINO::SuperBlob::SuperBlobHeader::fromD + return header; + } + +-const uint8_t* OpenVINO::SuperBlob::getBlobDataPointer() { ++const uint8_t* OpenVINO::SuperBlob::getBlobDataPointer() const { + const uint64_t offset = SuperBlobHeader::HEADER_SIZE; + return data.data() + offset; + } + +-int64_t OpenVINO::SuperBlob::getBlobDataSize() { ++int64_t OpenVINO::SuperBlob::getBlobDataSize() const { + return header.blobSize; + } + +-const uint8_t* OpenVINO::SuperBlob::getPatchDataPointer(int numShaves) { ++const uint8_t* OpenVINO::SuperBlob::getPatchDataPointer(int numShaves) const { + const uint64_t offset = + SuperBlobHeader::HEADER_SIZE + header.blobSize + std::accumulate(header.patchSizes.begin(), header.patchSizes.begin() + numShaves - 1, 0); + return data.data() + offset; + } + +-int64_t OpenVINO::SuperBlob::getPatchDataSize(int numShaves) { ++int64_t OpenVINO::SuperBlob::getPatchDataSize(int numShaves) const { + return header.patchSizes[numShaves - 1]; + } + diff --git a/include/depthai/models/ModelLoader.hpp b/include/depthai/models/ModelLoader.hpp new file mode 100644 index 0000000000..91bca810a5 --- /dev/null +++ b/include/depthai/models/ModelLoader.hpp @@ -0,0 +1,52 @@ +#pragma once + +#include +#include +#include + +#include + +namespace depthai { +namespace model { + +class ModelLoader { + public: + + ModelLoader() = default; + + // setters + void loadModelFromPath(const std::string &path); + void loadModelFromBytes(std::vector bytes, ModelType type); + + bool isModelLoaded() const; + + ModelType getModelType() const; + ModelVariant getModelVariant() const; + + static ModelType getModelTypeFromFilePath(const std::string &path); // Opens NNArchive and checks + static std::unordered_map> getSupportedModelFileExtensions(); + + private: + + void loadBlobModel(std::vector data); + void loadBlobModel(const std::string &path); + + void loadSuperblobModel(std::vector data); + void loadSuperblobModel(const std::string &path); + + void loadDlcModel(std::vector data); + void loadDlcModel(const std::string &path); + + void loadNNArchive(std::vector data); + void loadNNArchive(const std::string &path); + + ModelType modelType_; + std::optional modelVariant_; + +}; + +ModelVariant load(const std::string& modelPath); +ModelVariant load(const std::vector& data, ModelType type); + +} +} \ No newline at end of file diff --git a/include/depthai/models/ModelOptimizer.hpp b/include/depthai/models/ModelOptimizer.hpp new file mode 100644 index 0000000000..e69de29bb2 diff --git a/include/depthai/models/ModelSerializer.hpp b/include/depthai/models/ModelSerializer.hpp new file mode 100644 index 0000000000..e69de29bb2 diff --git a/include/depthai/models/ModelZoo.hpp b/include/depthai/models/ModelZoo.hpp new file mode 100644 index 0000000000..c498d58d27 --- /dev/null +++ b/include/depthai/models/ModelZoo.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include +#include +#include + +namespace depthai { +namespace model { +namespace zoo { + +ModelVariant load(const dai::NNModelDescription& description) { + std::string path = getModelFromZoo(description); + return depthai::model::load(path); +} + +} // namespace zoo +} // namespace model +} // namespace depthai \ No newline at end of file diff --git a/include/depthai/models/Models.hpp b/include/depthai/models/Models.hpp new file mode 100644 index 0000000000..1763eda697 --- /dev/null +++ b/include/depthai/models/Models.hpp @@ -0,0 +1,101 @@ +#pragma once + +#include +#include +#include + +namespace depthai { +namespace model { + +enum class ModelType { + UNKNOWN = 0, // Unknown model type + BLOB, + SUPERBLOB, + DLC, + NNARCHIVE // Helper type for NNArchive containing one of the above models +}; + +struct ModelSettings { + std::string modelName = ""; + + virtual bool isValid() const { + return true; + } + virtual ~ModelSettings() = default; +}; + +struct BlobSettings : ModelSettings {}; + +struct SuperBlobSettings : BlobSettings { + int numShaves = 6; +}; + +struct DlcSettings : ModelSettings {}; + +class Model { + public: + virtual ModelType type() const = 0; + virtual ~Model() = default; +}; + +class BlobModel : public Model { + public: + ModelType type() const override { + return ModelType::BLOB; + } + BlobModel(dai::OpenVINO::Blob model, BlobSettings settings) : model_(model), settings_(settings) {} + + const dai::OpenVINO::Blob& model() const { + return model_; + } + + const BlobSettings& settings() const { + return settings_; + } + + private: + dai::OpenVINO::Blob model_; + BlobSettings settings_; +}; + +class SuperBlobModel : public Model { + public: + ModelType type() const override { + return ModelType::SUPERBLOB; + } + + SuperBlobModel(dai::OpenVINO::SuperBlob model, SuperBlobSettings settings) : model_(model), settings_(settings) {} + + const dai::OpenVINO::SuperBlob& model() const { + return model_; + } + + const SuperBlobSettings& settings() const { + return settings_; + } + + private: + dai::OpenVINO::SuperBlob model_; + SuperBlobSettings settings_; +}; + +class DlcModel : public Model { + public: + ModelType type() const override { + return ModelType::DLC; + } + + DlcModel(const std::string& modelPath, DlcSettings settings) : modelPath_(modelPath), settings_(settings) {} + + private: + std::string modelPath_; + DlcSettings settings_; +}; + +using ModelVariant = std::variant; + +// Helper functions +ModelType getModelType(const ModelVariant& model); + +} // namespace model +} // namespace depthai diff --git a/include/depthai/openvino/OpenVINO.hpp b/include/depthai/openvino/OpenVINO.hpp index 6e42183014..49dc698545 100644 --- a/include/depthai/openvino/OpenVINO.hpp +++ b/include/depthai/openvino/OpenVINO.hpp @@ -82,7 +82,7 @@ class OpenVINO { * @param numShaves: Number of shaves to generate the blob for. Must be between 1 and NUMBER_OF_PATCHES. * @return dai::OpenVINO::Blob: Blob compiled for the specified number of shaves */ - dai::OpenVINO::Blob getBlobWithNumShaves(int numShaves); + dai::OpenVINO::Blob getBlobWithNumShaves(int numShaves) const; private: // A header in the superblob containing metadata about the blob and patches @@ -99,16 +99,16 @@ class OpenVINO { std::vector readSuperBlobFile(const std::string& path); // Get a pointer to the first byte of the blob data - const uint8_t* getBlobDataPointer(); + const uint8_t* getBlobDataPointer() const; // Get the size in bytes of the blob data - int64_t getBlobDataSize(); + int64_t getBlobDataSize() const; // Get a pointer to the first byte of the patch data for a specific number of shaves - const uint8_t* getPatchDataPointer(int numShaves); + const uint8_t* getPatchDataPointer(int numShaves) const; // Get the size in bytes of the patch data for a specific number of shaves - int64_t getPatchDataSize(int numShaves); + int64_t getPatchDataSize(int numShaves) const; // Load header - throw if an error occurs void loadAndCheckHeader(); diff --git a/include/depthai/pipeline/node/DetectionNetwork.hpp b/include/depthai/pipeline/node/DetectionNetwork.hpp index c2ee2ce351..ddd2dcf355 100644 --- a/include/depthai/pipeline/node/DetectionNetwork.hpp +++ b/include/depthai/pipeline/node/DetectionNetwork.hpp @@ -28,6 +28,23 @@ class DetectionNetwork : public DeviceNodeGroup { return networkPtr; } + + void setModel(const depthai::model::ModelVariant& model) { + std::visit([this](auto &&p){this->setModel(p);}, model); + } + + void setModel(const depthai::model::BlobModel& model) { + neuralNetwork->setModel(model); + } + + void setModel(const depthai::model::SuperBlobModel& model) { + neuralNetwork->setModel(model); + } + + void setModel(const depthai::model::DlcModel& model) { + neuralNetwork->setModel(model); + } + std::shared_ptr build(Node::Output& input, const NNArchive& nnArchive); std::shared_ptr build(std::shared_ptr input, dai::NNModelDescription modelDesc, float fps = 30.0f); diff --git a/include/depthai/pipeline/node/NeuralNetwork.hpp b/include/depthai/pipeline/node/NeuralNetwork.hpp index 1cdba3298f..32428eddcb 100644 --- a/include/depthai/pipeline/node/NeuralNetwork.hpp +++ b/include/depthai/pipeline/node/NeuralNetwork.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -74,6 +75,20 @@ class NeuralNetwork : public DeviceNodeCRTPsetModel(p);}, model); + } + + void setModel(const depthai::model::BlobModel& model) { + setBlob(model.model()); + } + void setModel(const depthai::model::SuperBlobModel& model) { + setBlob(model.model().getBlobWithNumShaves(model.settings().numShaves)); + } + void setModel(const depthai::model::DlcModel& model) { + // pass + } + /** * @brief Set NNArchive for this Node, throws if the archive's type is not SUPERBLOB * diff --git a/src/models/ModelLoader.cpp b/src/models/ModelLoader.cpp new file mode 100644 index 0000000000..d61ed6072f --- /dev/null +++ b/src/models/ModelLoader.cpp @@ -0,0 +1,204 @@ +#include +#include + +#include "../utility/ErrorMacros.hpp" +#include "models/Models.hpp" + +namespace depthai { +namespace model { + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +std::unordered_map> ModelLoader::getSupportedModelFileExtensions() { + return {{ModelType::UNKNOWN, {}}, + {ModelType::BLOB, {".blob"}}, + {ModelType::SUPERBLOB, {".superblob"}}, + {ModelType::DLC, {".dlc"}}, + {ModelType::NNARCHIVE, {".tar", ".tar.gz", ".tar.xz"}}}; +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +ModelType ModelLoader::getModelTypeFromFilePath(const std::string& path) { + auto endsWith = [](const std::string& str, const std::string& suffix) { + return str.size() >= suffix.size() && 0 == str.compare(str.size() - suffix.size(), suffix.size(), suffix); + }; + for(const auto& [modelType, extensions] : getSupportedModelFileExtensions()) { + for(const auto& extension : extensions) { + if(endsWith(path, extension)) { + return modelType; + } + } + } + return ModelType::UNKNOWN; +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +ModelType ModelLoader::getModelType() const { + return modelType_; +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +ModelVariant ModelLoader::getModelVariant() const { + DAI_CHECK_V(isModelLoaded(), "Model not loaded"); + return modelVariant_.value(); +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +bool ModelLoader::isModelLoaded() const { + return modelVariant_.has_value(); +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +void ModelLoader::loadModelFromPath(const std::string& path) { + ModelType modelType = getModelTypeFromFilePath(path); + switch(modelType) { + case ModelType::BLOB: + loadBlobModel(path); + break; + case ModelType::SUPERBLOB: + loadSuperblobModel(path); + break; + case ModelType::DLC: + loadDlcModel(path); + break; + case ModelType::NNARCHIVE: + loadNNArchive(path); + break; + case ModelType::UNKNOWN: + DAI_CHECK_V(false, "Unknown model type. Check your path: " + path); + } +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +void ModelLoader::loadModelFromBytes(std::vector bytes, ModelType type) { + switch(type) { + case ModelType::BLOB: + loadBlobModel(bytes); + break; + case ModelType::SUPERBLOB: + loadSuperblobModel(bytes); + break; + case ModelType::DLC: + loadDlcModel(bytes); + break; + case ModelType::NNARCHIVE: + loadNNArchive(bytes); + break; + case ModelType::UNKNOWN: + DAI_CHECK_V(false, "Unknown model type"); + } +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +void ModelLoader::loadBlobModel(std::vector data) { + DAI_CHECK_V(false, "Not implemented"); +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +void ModelLoader::loadBlobModel(const std::string& path) { + modelType_ = ModelType::BLOB; + dai::OpenVINO::Blob blob(path); + depthai::model::BlobSettings settings; + modelVariant_ = BlobModel(blob, settings); +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +void ModelLoader::loadSuperblobModel(std::vector data) { + DAI_CHECK_V(false, "Not implemented"); +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +void ModelLoader::loadSuperblobModel(const std::string& path) { + modelType_ = ModelType::SUPERBLOB; + dai::OpenVINO::SuperBlob superblob(path); + depthai::model::SuperBlobSettings settings; + modelVariant_ = SuperBlobModel(superblob, settings); +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +void ModelLoader::loadDlcModel(std::vector data) { + DAI_CHECK_V(false, "Not implemented"); +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +void ModelLoader::loadDlcModel(const std::string& path) { + modelType_ = ModelType::DLC; + modelVariant_ = DlcModel(path, DlcSettings()); +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +void ModelLoader::loadNNArchive(std::vector data) { + DAI_CHECK_V(false, "Not implemented"); +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +void ModelLoader::loadNNArchive(const std::string& path) { + + + dai::NNArchive archive(path); + switch(archive.getModelType()) { + case dai::model::ModelType::BLOB: + modelType_ = ModelType::BLOB; + modelVariant_ = BlobModel(archive.getBlob().value(), BlobSettings()); + break; + case dai::model::ModelType::SUPERBLOB: + modelType_ = ModelType::SUPERBLOB; + modelVariant_ = SuperBlobModel(archive.getSuperBlob().value(), SuperBlobSettings()); + break; + case dai::model::ModelType::DLC: + case dai::model::ModelType::OTHER: + case dai::model::ModelType::NNARCHIVE: + DAI_CHECK_V(false, "Unsupported model type"); + } +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +ModelVariant load(const std::string& modelPath) { + ModelLoader loader; + loader.loadModelFromPath(modelPath); + return loader.getModelVariant(); +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +ModelVariant load(const std::vector& data, ModelType type) { + ModelLoader loader; + loader.loadModelFromBytes(data, type); + DAI_CHECK_V(loader.isModelLoaded(), "Model not loaded"); + return loader.getModelVariant(); +} + +} // namespace model +} // namespace depthai \ No newline at end of file diff --git a/src/models/Models.cpp b/src/models/Models.cpp new file mode 100644 index 0000000000..dfa8e10b06 --- /dev/null +++ b/src/models/Models.cpp @@ -0,0 +1,12 @@ +#include + +namespace depthai { +namespace model { + +ModelType getModelType(const ModelVariant& model) { + return std::visit([](auto&& arg) -> ModelType { return arg.type(); }, model); +} + + +} +} \ No newline at end of file diff --git a/src/models/main.cpp b/src/models/main.cpp new file mode 100644 index 0000000000..e05c63a9b9 --- /dev/null +++ b/src/models/main.cpp @@ -0,0 +1,15 @@ +#include +#include +#include +#include "depthai/models/Models.hpp" +#include + +int main(int argc, char *argv[]) { + + depthai::model::ModelVariant model = depthai::model::zoo::load({"yolov6-nano", "RVC2"}); + depthai::model::SuperBlobModel superblob = std::get(model); + std::cout << "Superblob: " << static_cast(superblob.type()) << std::endl; + std::cout << "Number of shaves: " << superblob.settings().numShaves << std::endl; + + return 0; +} \ No newline at end of file diff --git a/src/openvino/OpenVINO.cpp b/src/openvino/OpenVINO.cpp index e52553b566..834181a000 100644 --- a/src/openvino/OpenVINO.cpp +++ b/src/openvino/OpenVINO.cpp @@ -99,7 +99,7 @@ OpenVINO::SuperBlob::SuperBlob(std::vector data) { validateSuperblob(); } -dai::OpenVINO::Blob OpenVINO::SuperBlob::getBlobWithNumShaves(int numShaves) { +dai::OpenVINO::Blob OpenVINO::SuperBlob::getBlobWithNumShaves(int numShaves) const{ if(numShaves < 1 || numShaves > static_cast(OpenVINO::SuperBlob::NUMBER_OF_PATCHES)) { throw std::out_of_range("Invalid number of shaves: " + std::to_string(numShaves) + " (expected 1 to " + std::to_string(OpenVINO::SuperBlob::NUMBER_OF_PATCHES) + ")"); @@ -162,22 +162,22 @@ OpenVINO::SuperBlob::SuperBlobHeader OpenVINO::SuperBlob::SuperBlobHeader::fromD return header; } -const uint8_t* OpenVINO::SuperBlob::getBlobDataPointer() { +const uint8_t* OpenVINO::SuperBlob::getBlobDataPointer() const { const uint64_t offset = SuperBlobHeader::HEADER_SIZE; return data.data() + offset; } -int64_t OpenVINO::SuperBlob::getBlobDataSize() { +int64_t OpenVINO::SuperBlob::getBlobDataSize() const { return header.blobSize; } -const uint8_t* OpenVINO::SuperBlob::getPatchDataPointer(int numShaves) { +const uint8_t* OpenVINO::SuperBlob::getPatchDataPointer(int numShaves) const { const uint64_t offset = SuperBlobHeader::HEADER_SIZE + header.blobSize + std::accumulate(header.patchSizes.begin(), header.patchSizes.begin() + numShaves - 1, 0); return data.data() + offset; } -int64_t OpenVINO::SuperBlob::getPatchDataSize(int numShaves) { +int64_t OpenVINO::SuperBlob::getPatchDataSize(int numShaves) const { return header.patchSizes[numShaves - 1]; } From 2ae8d1f78016f8731ea873b10e2fe3b55875827f Mon Sep 17 00:00:00 2001 From: kuba Date: Wed, 25 Sep 2024 21:36:21 -0400 Subject: [PATCH 2/2] poc --- bindings/python/src/modelzoo/ZooBindings.cpp | 21 ++-- .../node/DetectionNetworkBindings.cpp | 17 +++ examples/python/RVC2/NNArchive/nn_archive2.py | 119 ++++++++++++++++++ include/depthai/models/Models.hpp | 62 ++++++--- include/depthai/models/README.md | 27 ++++ include/depthai/openvino/OpenVINO.hpp | 7 ++ .../pipeline/node/DetectionNetwork.hpp | 8 ++ .../depthai/pipeline/node/DetectionParser.hpp | 25 ++++ .../depthai/pipeline/node/NeuralNetwork.hpp | 17 ++- src/models/ModelLoader.cpp | 32 +++-- src/models/main.cpp | 9 ++ 11 files changed, 305 insertions(+), 39 deletions(-) create mode 100644 examples/python/RVC2/NNArchive/nn_archive2.py create mode 100644 include/depthai/models/README.md diff --git a/bindings/python/src/modelzoo/ZooBindings.cpp b/bindings/python/src/modelzoo/ZooBindings.cpp index 508534dbcb..7a9fec3285 100644 --- a/bindings/python/src/modelzoo/ZooBindings.cpp +++ b/bindings/python/src/modelzoo/ZooBindings.cpp @@ -43,7 +43,10 @@ void ZooBindings::bind(pybind11::module& m, void* pCallstack) { py::module m_zoo = m_model.def_submodule("zoo"); // Bind Model - py::class_>(m_model, "Model"); + py::class_> model(m_model, "Model"); + + + py::enum_(m_model, "ModelType") .value("BLOB", depthai::model::ModelType::BLOB) .value("SUPERBLOB", depthai::model::ModelType::SUPERBLOB) @@ -52,25 +55,29 @@ void ZooBindings::bind(pybind11::module& m, void* pCallstack) { .export_values(); // Bind ModelSettings - py::class_>(m_model, "ModelSettings"); + py::class_> modelSettings(m_model, "ModelSettings"); + modelSettings.def_readwrite("nnArchiveConfig", &depthai::model::ModelSettings::nnArchiveConfig, DOC(depthai::model, ModelSettings, settings)); // Bind BlobSettings py::class_>(m_model, "BlobSettings"); // Bind SuperBlobSettings - py::class_>(m_model, "SuperBlobSettings"); + py::class_> superblobSettings(m_model, "SuperBlobSettings"); + superblobSettings.def_readwrite("nnArchiveConfig", &depthai::model::SuperBlobSettings::nnArchiveConfig, DOC(depthai::model, SuperBlobSettings, settings)); // Bind DlcSettings py::class_>(m_model, "DlcSettings"); // Bind BlobModel - py::class_>(m_model, "BlobModel"); - + py::class_> blobmodel(m_model, "BlobModel"); + blobmodel.def("settings", &depthai::model::BlobModel::settings, DOC(depthai::model, BlobModel, settings)); // Bind SuperBlobModel - py::class_>(m_model, "SuperBlobModel"); + py::class_> superblobmodel(m_model, "SuperBlobModel"); + superblobmodel.def("settings", &depthai::model::SuperBlobModel::settings, DOC(depthai::model, SuperBlobModel, settings)); // Bind DlcModel - py::class_>(m_model, "DlcModel"); + py::class_> dlcmodel(m_model, "DlcModel"); + dlcmodel.def("settings", &depthai::model::DlcModel::settings, DOC(depthai::model, DlcModel, settings)); // Bind load function m_model.def("load", py::overload_cast(&depthai::model::load), py::arg("path")); diff --git a/bindings/python/src/pipeline/node/DetectionNetworkBindings.cpp b/bindings/python/src/pipeline/node/DetectionNetworkBindings.cpp index 8fcb198601..8c994ac07b 100644 --- a/bindings/python/src/pipeline/node/DetectionNetworkBindings.cpp +++ b/bindings/python/src/pipeline/node/DetectionNetworkBindings.cpp @@ -1,4 +1,6 @@ #include +#include + #include #include "Common.hpp" @@ -10,6 +12,11 @@ #include "depthai/pipeline/Pipeline.hpp" #include "depthai/pipeline/node/DetectionNetwork.hpp" +#include "depthai/models/Models.hpp" + +PYBIND11_MAKE_OPAQUE(depthai::model::ModelVariant); + + void bind_detectionnetwork(pybind11::module& m, void* pCallstack) { using namespace dai; using namespace dai::node; @@ -58,6 +65,16 @@ void bind_detectionnetwork(pybind11::module& m, void* pCallstack) { DETECTION_NETWORK_BUILD_PYARGS, DETECTION_NETWORK_PYARGS) .def("build", py::overload_cast, NNModelDescription, float>(&DetectionNetwork::build), py::arg("input"), py::arg("model"), py::arg("fps") = 30.0f) + // .def("build", py::overload_cast(&DetectionNetwork::build), py::arg("input"), py::arg("model")) + .def("build", [](DetectionNetwork& self, Node::Output& input, const depthai::model::BlobModel& model) { + return self.build(input, model); + }, py::arg("input"), py::arg("model")) + .def("build", [](DetectionNetwork& self, Node::Output& input, const depthai::model::SuperBlobModel& model) { + return self.build(input, model); + }, py::arg("input"), py::arg("model")) + .def("build", [](DetectionNetwork& self, Node::Output& input, const depthai::model::DlcModel& model) { + return self.build(input, model); + }, py::arg("input"), py::arg("model")) .def(py::init([](DETECTION_NETWORK_BUILD_ARGS, DETECTION_NETWORK_ARGS) { auto self = getImplicitPipeline()->create(); self->build(input, nnArchive); diff --git a/examples/python/RVC2/NNArchive/nn_archive2.py b/examples/python/RVC2/NNArchive/nn_archive2.py new file mode 100644 index 0000000000..5077f0cfc9 --- /dev/null +++ b/examples/python/RVC2/NNArchive/nn_archive2.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python3 + +import cv2 +import depthai as dai +import numpy as np +import time + +# Create pipeline +with dai.Pipeline() as pipeline: + # Define sources and outputs + camRgb = pipeline.create(dai.node.ColorCamera) + # Properties + camRgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P) + camRgb.setInterleaved(False) + camRgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR) + camRgb.setFps(15) + + + """ + # load model + model = dai.model.zoo.load(modelDescription) + settings = model.settings() + + # Set model + detectionNetwork.setModel + """ + modelDescription = dai.NNModelDescription(modelSlug="yolov6-nano", platform="RVC2") + model = dai.model.zoo.load(modelDescription) + + + detectionNetwork = pipeline.create(dai.node.DetectionNetwork).build( + camRgb.preview, model + ) + detectionNetwork.setNumInferenceThreads(2) + + + # detectionNetwork.setModel(model) + h,w = model.settings().nnArchiveConfig.getConfigV1().model.inputs[0].shape[-2:] + camRgb.setPreviewSize(w, h) + + + # If needed, you can set the NNArchive by yourself + # detectionNetwork.setNNArchive(nnArchive) + + # If nnArchive.getModelType() == dai.ModelType.SUPERBLOB + # you can specify the number of shaves + # detectionNetwork.setNNArchive(nnArchive, numShaves=9) + # When ^^^ is used and the archive type is not SUPERBLOB, an exception will be thrown + + qRgb = detectionNetwork.passthrough.createOutputQueue() + qDet = detectionNetwork.out.createOutputQueue() + + labelMap = detectionNetwork.getClasses() + + pipeline.start() + + frame = None + detections = [] + startTime = time.monotonic() + counter = 0 + color2 = (255, 255, 255) + + # nn data, being the bounding box locations, are in <0..1> range - they need to be normalized with frame width/height + def frameNorm(frame, bbox): + normVals = np.full(len(bbox), frame.shape[0]) + normVals[::2] = frame.shape[1] + return (np.clip(np.array(bbox), 0, 1) * normVals).astype(int) + + def displayFrame(name, frame): + color = (255, 0, 0) + for detection in detections: + bbox = frameNorm( + frame, + (detection.xmin, detection.ymin, detection.xmax, detection.ymax), + ) + cv2.putText( + frame, + labelMap[detection.label], + (bbox[0] + 10, bbox[1] + 20), + cv2.FONT_HERSHEY_TRIPLEX, + 0.5, + 255, + ) + cv2.putText( + frame, + f"{int(detection.confidence * 100)}%", + (bbox[0] + 10, bbox[1] + 40), + cv2.FONT_HERSHEY_TRIPLEX, + 0.5, + 255, + ) + cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2) + # Show the frame + cv2.imshow(name, frame) + + while pipeline.isRunning(): + inRgb: dai.ImgFrame = qRgb.get() + inDet: dai.ImgDetections = qDet.get() + if inRgb is not None: + frame = inRgb.getCvFrame() + cv2.putText( + frame, + "NN fps: {:.2f}".format(counter / (time.monotonic() - startTime)), + (2, frame.shape[0] - 4), + cv2.FONT_HERSHEY_TRIPLEX, + 0.4, + color2, + ) + + if inDet is not None: + detections = inDet.detections + counter += 1 + + if frame is not None: + displayFrame("rgb", frame) + + if cv2.waitKey(1) == ord("q"): + pipeline.stop() + break diff --git a/include/depthai/models/Models.hpp b/include/depthai/models/Models.hpp index 1763eda697..743eab5243 100644 --- a/include/depthai/models/Models.hpp +++ b/include/depthai/models/Models.hpp @@ -1,8 +1,13 @@ #pragma once +#include #include +#include #include #include +#include +#include +#include namespace depthai { namespace model { @@ -17,14 +22,28 @@ enum class ModelType { struct ModelSettings { std::string modelName = ""; + std::vector supportedPlatforms = {}; virtual bool isValid() const { return true; } + + bool isSupported(dai::Platform platform) const { + return std::find(supportedPlatforms.begin(), supportedPlatforms.end(), platform) != supportedPlatforms.end(); + } + + bool isSupported(const dai::Device& device) const { + return isSupported(device.getPlatform()); + } + virtual ~ModelSettings() = default; + + // Archive config for NNArchive models + std::optional nnArchiveConfig; }; -struct BlobSettings : ModelSettings {}; +struct BlobSettings : ModelSettings { +}; struct SuperBlobSettings : BlobSettings { int numShaves = 6; @@ -43,19 +62,20 @@ class BlobModel : public Model { ModelType type() const override { return ModelType::BLOB; } - BlobModel(dai::OpenVINO::Blob model, BlobSettings settings) : model_(model), settings_(settings) {} - const dai::OpenVINO::Blob& model() const { - return model_; + BlobModel(std::shared_ptr model, std::shared_ptr settings) : modelPtr_(model), settingsPtr_(settings) {} + + inline const dai::OpenVINO::Blob& model() const { + return *modelPtr_; } - const BlobSettings& settings() const { - return settings_; + inline const BlobSettings& settings() const { + return *settingsPtr_; } private: - dai::OpenVINO::Blob model_; - BlobSettings settings_; + std::shared_ptr modelPtr_; + std::shared_ptr settingsPtr_; }; class SuperBlobModel : public Model { @@ -64,19 +84,19 @@ class SuperBlobModel : public Model { return ModelType::SUPERBLOB; } - SuperBlobModel(dai::OpenVINO::SuperBlob model, SuperBlobSettings settings) : model_(model), settings_(settings) {} + SuperBlobModel(std::shared_ptr model, std::shared_ptr settings) : modelPtr_(model), settingsPtr_(settings) {} - const dai::OpenVINO::SuperBlob& model() const { - return model_; + inline const dai::OpenVINO::SuperBlob& model() const { + return *modelPtr_; } - const SuperBlobSettings& settings() const { - return settings_; + inline const SuperBlobSettings& settings() const { + return *settingsPtr_; } private: - dai::OpenVINO::SuperBlob model_; - SuperBlobSettings settings_; + std::shared_ptr modelPtr_; + std::shared_ptr settingsPtr_; }; class DlcModel : public Model { @@ -85,11 +105,19 @@ class DlcModel : public Model { return ModelType::DLC; } - DlcModel(const std::string& modelPath, DlcSettings settings) : modelPath_(modelPath), settings_(settings) {} + DlcModel(const std::string& modelPath, std::shared_ptr settings) : modelPath_(modelPath), settingsPtr_(settings) {} + + inline const std::string& modelPath() const { + return modelPath_; + } + + inline const DlcSettings& settings() const { + return *settingsPtr_; + } private: std::string modelPath_; - DlcSettings settings_; + std::shared_ptr settingsPtr_; }; using ModelVariant = std::variant; diff --git a/include/depthai/models/README.md b/include/depthai/models/README.md new file mode 100644 index 0000000000..b832f1e65c --- /dev/null +++ b/include/depthai/models/README.md @@ -0,0 +1,27 @@ +## Modules + +### Models +A wrapper around a model and its settings. + +### ModelSerializer +Serializing and Deserializing models. This is useful for both sending models from and to device as well as storing models in memory + +### ModelLoader +Everything related to model loading and management. Depends on Model and ModelSerializer. + +### ModelZoo +A collection of functions for downloading a model from a remote repository. Depends on ModelLoader and ModelSerializer (for interpreting bytes as a model + metadata). + +### ModelArchive +A helper class for working with NNArchives. Depends on Models, ModelSerializer and ModelLoader + +### ModelOptimizer +A set of functions for optimizing model settings (e.g. number of shaves in case of a Superblob) for a specific node. + + +Ideally we can provide smart default settings for each model (i.e. without head) for each of the supported model types (Blob, SuperBlob, +Dlc, Onnx, Pytorch, etc.) + +I think NNArchive is a neat thing and we should keep it, as playing around with parameters settable in a json file is much easier than fiddling +with bits stored in a data structure containing the model and the settings themselves ([ModelBytes; SettingsBytes]); + diff --git a/include/depthai/openvino/OpenVINO.hpp b/include/depthai/openvino/OpenVINO.hpp index 49dc698545..a1cb5460a0 100644 --- a/include/depthai/openvino/OpenVINO.hpp +++ b/include/depthai/openvino/OpenVINO.hpp @@ -76,6 +76,13 @@ class OpenVINO { */ SuperBlob(const std::string& pathToSuperBlobFile); + /** + * @brief Copy constructor + * + * @param other: Other superblob to copy + */ + SuperBlob(const SuperBlob& other) = default; + /** * @brief Generate a blob with a specific number of shaves * diff --git a/include/depthai/pipeline/node/DetectionNetwork.hpp b/include/depthai/pipeline/node/DetectionNetwork.hpp index ddd2dcf355..52cb6f2f32 100644 --- a/include/depthai/pipeline/node/DetectionNetwork.hpp +++ b/include/depthai/pipeline/node/DetectionNetwork.hpp @@ -35,17 +35,25 @@ class DetectionNetwork : public DeviceNodeGroup { void setModel(const depthai::model::BlobModel& model) { neuralNetwork->setModel(model); + detectionParser->setModel(model); } void setModel(const depthai::model::SuperBlobModel& model) { neuralNetwork->setModel(model); + detectionParser->setModel(model); } void setModel(const depthai::model::DlcModel& model) { neuralNetwork->setModel(model); + detectionParser->setModel(model); } std::shared_ptr build(Node::Output& input, const NNArchive& nnArchive); + std::shared_ptr build(Node::Output& input, const depthai::model::ModelVariant& model) { + setModel(model); + input.link(this->input); + return std::static_pointer_cast(shared_from_this()); + } std::shared_ptr build(std::shared_ptr input, dai::NNModelDescription modelDesc, float fps = 30.0f); Subnode neuralNetwork{*this, "neuralNetwork"}; diff --git a/include/depthai/pipeline/node/DetectionParser.hpp b/include/depthai/pipeline/node/DetectionParser.hpp index 09b3532f57..390da547b0 100644 --- a/include/depthai/pipeline/node/DetectionParser.hpp +++ b/include/depthai/pipeline/node/DetectionParser.hpp @@ -3,6 +3,8 @@ #include #include +#include + // standard #include @@ -53,6 +55,29 @@ class DetectionParser : public DeviceNodeCRTPsetModel(p); }, model); + } + + void setModel(const depthai::model::BlobModel& model) { + const auto &settings = model.settings(); + if(settings.nnArchiveConfig.has_value()) { + setConfig(settings.nnArchiveConfig.value()); + } + } + + void setModel(const depthai::model::SuperBlobModel& model) { + const auto &settings = model.settings(); + if(settings.nnArchiveConfig.has_value()) { + setConfig(settings.nnArchiveConfig.value()); + } + } + + void setModel(const depthai::model::DlcModel& model) { + // pass + } + + /** * @brief Set NNArchive for this Node. If the archive's type is SUPERBLOB, use default number of shaves. * diff --git a/include/depthai/pipeline/node/NeuralNetwork.hpp b/include/depthai/pipeline/node/NeuralNetwork.hpp index ca2ac07523..718ae059a1 100644 --- a/include/depthai/pipeline/node/NeuralNetwork.hpp +++ b/include/depthai/pipeline/node/NeuralNetwork.hpp @@ -39,8 +39,14 @@ class NeuralNetwork : public DeviceNodeCRTP build(Node::Output& input, const NNArchive& nnArchive); - std::shared_ptr build(const NNArchiveConfig& nnArchiveConfig, std::shared_ptr input, dai::NNModelDescription modelDesc, float fps = 30.0f); - std::shared_ptr build(const NNArchiveVersionedConfig& nnArchiveConfig, std::shared_ptr input, dai::NNModelDescription modelDesc, float fps = 30.0f); + std::shared_ptr build(const NNArchiveConfig& nnArchiveConfig, + std::shared_ptr input, + dai::NNModelDescription modelDesc, + float fps = 30.0f); + std::shared_ptr build(const NNArchiveVersionedConfig& nnArchiveConfig, + std::shared_ptr input, + dai::NNModelDescription modelDesc, + float fps = 30.0f); /** * Input message with data to be inferred upon @@ -78,15 +84,18 @@ class NeuralNetwork : public DeviceNodeCRTPsetModel(p);}, model); + std::visit([this](auto&& p) { this->setModel(p); }, model); } void setModel(const depthai::model::BlobModel& model) { setBlob(model.model()); } + void setModel(const depthai::model::SuperBlobModel& model) { - setBlob(model.model().getBlobWithNumShaves(model.settings().numShaves)); + int numShaves = model.settings().numShaves; + setBlob(model.model().getBlobWithNumShaves(numShaves)); } + void setModel(const depthai::model::DlcModel& model) { // pass } diff --git a/src/models/ModelLoader.cpp b/src/models/ModelLoader.cpp index d61ed6072f..73459bf9c8 100644 --- a/src/models/ModelLoader.cpp +++ b/src/models/ModelLoader.cpp @@ -114,9 +114,9 @@ void ModelLoader::loadBlobModel(std::vector data) { /**********************************************************************************************************************/ void ModelLoader::loadBlobModel(const std::string& path) { modelType_ = ModelType::BLOB; - dai::OpenVINO::Blob blob(path); - depthai::model::BlobSettings settings; - modelVariant_ = BlobModel(blob, settings); + auto blobPtr = std::make_shared(path); + auto settingsPtr = std::make_shared(); + modelVariant_ = BlobModel(blobPtr, settingsPtr); } /**********************************************************************************************************************/ @@ -131,9 +131,9 @@ void ModelLoader::loadSuperblobModel(std::vector data) { /**********************************************************************************************************************/ void ModelLoader::loadSuperblobModel(const std::string& path) { modelType_ = ModelType::SUPERBLOB; - dai::OpenVINO::SuperBlob superblob(path); - depthai::model::SuperBlobSettings settings; - modelVariant_ = SuperBlobModel(superblob, settings); + auto superblobPtr = std::make_shared(path); + auto settingsPtr = std::make_shared(); + modelVariant_ = SuperBlobModel(superblobPtr, settingsPtr); } /**********************************************************************************************************************/ @@ -148,7 +148,8 @@ void ModelLoader::loadDlcModel(std::vector data) { /**********************************************************************************************************************/ void ModelLoader::loadDlcModel(const std::string& path) { modelType_ = ModelType::DLC; - modelVariant_ = DlcModel(path, DlcSettings()); + auto settingsPtr = std::make_shared(); + modelVariant_ = DlcModel(path, settingsPtr); } /**********************************************************************************************************************/ @@ -162,17 +163,26 @@ void ModelLoader::loadNNArchive(std::vector data) { /**********************************************************************************************************************/ /**********************************************************************************************************************/ void ModelLoader::loadNNArchive(const std::string& path) { - - + // TODO: Remove NNArchive type completely == Load data from NNArchive but return a ModelVariant, not an NNArchive object dai::NNArchive archive(path); switch(archive.getModelType()) { case dai::model::ModelType::BLOB: modelType_ = ModelType::BLOB; - modelVariant_ = BlobModel(archive.getBlob().value(), BlobSettings()); + { + auto blobPtr = std::make_shared(archive.getBlob().value().data); + auto settingsPtr = std::make_shared(); + settingsPtr->nnArchiveConfig = archive.getVersionedConfig(); + modelVariant_ = BlobModel(blobPtr, settingsPtr); + } break; case dai::model::ModelType::SUPERBLOB: modelType_ = ModelType::SUPERBLOB; - modelVariant_ = SuperBlobModel(archive.getSuperBlob().value(), SuperBlobSettings()); + { + auto superblobPtr = std::make_shared(archive.getSuperBlob().value()); + auto settingsPtr = std::make_shared(); + settingsPtr->nnArchiveConfig = archive.getVersionedConfig(); + modelVariant_ = SuperBlobModel(superblobPtr, settingsPtr); + } break; case dai::model::ModelType::DLC: case dai::model::ModelType::OTHER: diff --git a/src/models/main.cpp b/src/models/main.cpp index e05c63a9b9..e6709f4380 100644 --- a/src/models/main.cpp +++ b/src/models/main.cpp @@ -3,11 +3,20 @@ #include #include "depthai/models/Models.hpp" #include +#include int main(int argc, char *argv[]) { + // Load model from ModelZoo + auto start = std::chrono::high_resolution_clock::now(); depthai::model::ModelVariant model = depthai::model::zoo::load({"yolov6-nano", "RVC2"}); + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration elapsed = end - start; + std::cout << "Time taken to load model: " << elapsed.count() << " seconds" << std::endl; depthai::model::SuperBlobModel superblob = std::get(model); + auto end2 = std::chrono::high_resolution_clock::now(); + std::chrono::duration elapsed2 = end2 - end; + std::cout << "Time taken to get superblob: " << elapsed2.count() << " seconds" << std::endl; std::cout << "Superblob: " << static_cast(superblob.type()) << std::endl; std::cout << "Number of shaves: " << superblob.settings().numShaves << std::endl;