[POC] Neural Network Models

CMakeLists.txt

@@ -380,11 +380,16 @@ set(TARGET_CORE_SOURCES
     src/device/DeviceGate.cpp
     src/utility/ArchiveUtil.cpp
     src/nn_archive/NNArchive.cpp
-    src/nn_archive/NNArchiveConfig.cpp
+    src/nn_archive/NNArchiveVersionedConfig.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

bindings/python/src/modelzoo/ZooBindings.cpp

@@ -2,6 +2,11 @@
 
 // depthai
 #include "depthai/modelzoo/Zoo.hpp"
+#include <depthai/models/Models.hpp>
+#include <depthai/models/ModelLoader.hpp>
+#include <depthai/models/ModelZoo.hpp>
+
+// Model bindings
 
 void ZooBindings::bind(pybind11::module& m, void* pCallstack) {
     using namespace dai;
@@ -32,4 +37,50 @@ void ZooBindings::bind(pybind11::module& m, void* pCallstack) {
           py::arg("cacheDirectory") = MODEL_ZOO_DEFAULT_CACHE_DIRECTORY,
           py::arg("apiKey") = "",
           DOC(dai, downloadModelsFromZoo));
-}
+
+    // 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_<depthai::model::Model, std::shared_ptr<depthai::model::Model>> model(m_model, "Model");
+
+
+
+    py::enum_<depthai::model::ModelType>(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_<depthai::model::ModelSettings, std::shared_ptr<depthai::model::ModelSettings>> modelSettings(m_model, "ModelSettings");
+    modelSettings.def_readwrite("nnArchiveConfig", &depthai::model::ModelSettings::nnArchiveConfig, DOC(depthai::model, ModelSettings, settings));
+
+    // Bind BlobSettings
+    py::class_<depthai::model::BlobSettings, std::shared_ptr<depthai::model::BlobSettings>>(m_model, "BlobSettings");
+
+    // Bind SuperBlobSettings
+    py::class_<depthai::model::SuperBlobSettings, std::shared_ptr<depthai::model::SuperBlobSettings>> superblobSettings(m_model, "SuperBlobSettings");
+    superblobSettings.def_readwrite("nnArchiveConfig", &depthai::model::SuperBlobSettings::nnArchiveConfig, DOC(depthai::model, SuperBlobSettings, settings));
+
+    // Bind DlcSettings
+    py::class_<depthai::model::DlcSettings, std::shared_ptr<depthai::model::DlcSettings>>(m_model, "DlcSettings");
+
+    // Bind BlobModel
+    py::class_<depthai::model::BlobModel, std::shared_ptr<depthai::model::BlobModel>> blobmodel(m_model, "BlobModel");
+    blobmodel.def("settings", &depthai::model::BlobModel::settings, DOC(depthai::model, BlobModel, settings));
+    // Bind SuperBlobModel
+    py::class_<depthai::model::SuperBlobModel, std::shared_ptr<depthai::model::SuperBlobModel>> superblobmodel(m_model, "SuperBlobModel");
+    superblobmodel.def("settings", &depthai::model::SuperBlobModel::settings, DOC(depthai::model, SuperBlobModel, settings));
+
+    // Bind DlcModel
+    py::class_<depthai::model::DlcModel, std::shared_ptr<depthai::model::DlcModel>> 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<const std::string&>(&depthai::model::load), py::arg("path"));
+    m_zoo.def("load", &depthai::model::zoo::load, py::arg("ModelDescription"));
+
+}

bindings/python/src/nn_archive/NNArchiveBindings.cpp

@@ -5,7 +5,7 @@
 #include <pybind11/stl_bind.h>
 
 #include "depthai/nn_archive/NNArchive.hpp"
-#include "depthai/nn_archive/NNArchiveConfig.hpp"
+#include "depthai/nn_archive/NNArchiveVersionedConfig.hpp"
 #include "depthai/nn_archive/NNArchiveEntry.hpp"
 
 // v1 nn_archive bindings
@@ -32,8 +32,9 @@ void NNArchiveBindings::bind(pybind11::module& m, void* pCallstack) {
     py::class_<NNArchive> nnArchive(m, "NNArchive", DOC(dai, NNArchive));
     py::class_<NNArchiveOptions> nnArchiveOptions(m, "NNArchiveOptions", DOC(dai, NNArchiveOptions));
 
-    // NNArchiveConfig
-    py::class_<NNArchiveConfig> nnArchiveConfig(m, "NNArchiveConfig", DOC(dai, NNArchiveConfig));
+    // NNArchiveVersionedConfig
+    py::class_<NNArchiveVersionedConfig> nnArchiveVersionedConfig(m, "NNArchiveVersionedConfig", DOC(dai, NNArchiveVersionedConfig));
+    py::enum_<NNArchiveConfigVersion> nnArchiveConfigVersion(m, "NNArchiveConfigVersion", DOC(dai, NNArchiveConfigVersion));
 
     // NNArchiveEntry
     py::class_<NNArchiveEntry> nnArchiveEntry(m, "NNArchiveEntry", DOC(dai, NNArchiveEntry));
@@ -85,7 +86,8 @@ void NNArchiveBindings::bind(pybind11::module& m, void* pCallstack) {
     nnArchive.def("getBlob", &NNArchive::getBlob, DOC(dai, NNArchive, getBlob));
     nnArchive.def("getSuperBlob", &NNArchive::getSuperBlob, DOC(dai, NNArchive, getBlob));
     nnArchive.def("getModelPath", &NNArchive::getModelPath, DOC(dai, NNArchive, getModelPath));
-    nnArchive.def("getConfig", &NNArchive::getConfig, DOC(dai, NNArchive, getConfig));
+    nnArchive.def("getConfig", &NNArchive::getConfig<NNArchiveConfig>, DOC(dai, NNArchive, getConfig));
+    nnArchive.def("getConfigV1", &NNArchive::getConfig<v1::Config>, DOC(dai, NNArchive, getConfig));
     nnArchive.def("getModelType", &NNArchive::getModelType, DOC(dai, NNArchive, getModelType));
 
     // Bind NNArchive options
@@ -99,25 +101,32 @@ void NNArchiveBindings::bind(pybind11::module& m, void* pCallstack) {
         [](const NNArchiveOptions& opt) { return opt.extractFolder(); },
         [](NNArchiveOptions& opt, const std::string& extractFolder) { opt.extractFolder(extractFolder); });
 
-    // Bind NNArchiveConfig
-    nnArchiveConfig.def(py::init<const dai::Path&, NNArchiveEntry::Compression>(),
+    // Bind NNArchiveVersionedConfig
+    nnArchiveVersionedConfig.def(py::init<const dai::Path&, NNArchiveEntry::Compression>(),
                         py::arg("path"),
                         py::arg("compression") = NNArchiveEntry::Compression::AUTO,
-                        DOC(dai, NNArchiveConfig, NNArchiveConfig));
-    nnArchiveConfig.def(py::init<const std::vector<uint8_t>&, NNArchiveEntry::Compression>(),
+                        DOC(dai, NNArchiveVersionedConfig, NNArchiveVersionedConfig));
+    nnArchiveVersionedConfig.def(py::init<const std::vector<uint8_t>&, NNArchiveEntry::Compression>(),
                         py::arg("data"),
                         py::arg("compression") = NNArchiveEntry::Compression::AUTO,
-                        DOC(dai, NNArchiveConfig, NNArchiveConfig));
-    nnArchiveConfig.def(py::init([](const std::function<int()>& openCallback,
+                        DOC(dai, NNArchiveVersionedConfig, NNArchiveVersionedConfig));
+    nnArchiveVersionedConfig.def(py::init([](const std::function<int()>& openCallback,
                                     const std::function<std::vector<uint8_t>()>& readCallback,
                                     const std::function<int64_t(int64_t, NNArchiveEntry::Seek)>& seekCallback,
                                     const std::function<int64_t(int64_t)>& skipCallback,
                                     const std::function<int()>& closeCallback,
                                     NNArchiveEntry::Compression compression) {
         auto readCallbackWrapper = [readCallback]() { return std::make_shared<std::vector<uint8_t>>(readCallback()); };
-        return NNArchiveConfig(openCallback, readCallbackWrapper, seekCallback, skipCallback, closeCallback, compression);
+        return NNArchiveVersionedConfig(openCallback, readCallbackWrapper, seekCallback, skipCallback, closeCallback, compression);
     }));
-    nnArchiveConfig.def("getConfigV1", &NNArchiveConfig::getConfigV1, DOC(dai, NNArchiveConfig, getConfigV1));
+
+    // Bind NNArchiveVersionedConfig.
+    nnArchiveVersionedConfig.def("getConfig", &NNArchiveVersionedConfig::getConfig<NNArchiveConfig>, DOC(dai, NNArchiveVersionedConfig, getConfig));
+    nnArchiveVersionedConfig.def("getConfigV1", &NNArchiveVersionedConfig::getConfig<v1::Config>, DOC(dai, NNArchiveVersionedConfig, getConfig));
+    nnArchiveVersionedConfig.def("getVersion", &NNArchiveVersionedConfig::getVersion, DOC(dai, NNArchiveVersionedConfig, getVersion));
+
+    // Bind NNArchiveConfigVersion
+    nnArchiveConfigVersion.value("V1", NNArchiveConfigVersion::V1);
 
     archiveEntryCompression.value("AUTO", NNArchiveEntry::Compression::AUTO);
     archiveEntryCompression.value("RAW_FS", NNArchiveEntry::Compression::RAW_FS);

bindings/python/src/pipeline/node/DetectionNetworkBindings.cpp

@@ -1,4 +1,6 @@
 #include <pybind11/cast.h>
+#include <pybind11/pybind11.h>
+
 #include <memory>
 
 #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<std::shared_ptr<Camera>, NNModelDescription, float>(&DetectionNetwork::build), py::arg("input"), py::arg("model"), py::arg("fps") = 30.0f)
+        // .def("build", py::overload_cast<Node::Output&, const depthai::model::ModelVariant&>(&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<DetectionNetwork>();
                  self->build(input, nnArchive);

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<const depthai::model::BlobModel&>(&NeuralNetwork::setModel), py::arg("model"))
+        .def("setModel", py::overload_cast<const depthai::model::SuperBlobModel&>(&NeuralNetwork::setModel), py::arg("model"))
+        .def("setModel", py::overload_cast<const depthai::model::DlcModel&>(&NeuralNetwork::setModel), py::arg("model"))
         .def("setNumNCEPerInferenceThread",
              &NeuralNetwork::setNumNCEPerInferenceThread,
              py::arg("numNCEPerThread"),
@@ -53,6 +58,7 @@ void bind_neuralnetwork(pybind11::module& m, void* pCallstack){
         .def("setNNArchive", py::overload_cast<const NNArchive&, int>(&NeuralNetwork::setNNArchive), py::arg("nnArchive"), py::arg("numShaves"), DOC(dai, node, NeuralNetwork, setNNArchive, 2))
         .def("setFromModelZoo", py::overload_cast<NNModelDescription, bool>(&NeuralNetwork::setFromModelZoo), py::arg("description"), py::arg("useCached"), DOC(dai, node, NeuralNetwork, setFromModelZoo))
         .def("build", py::overload_cast<dai::Node::Output&, const NNArchive&>(&NeuralNetwork::build), py::arg("input"), py::arg("nnArchive"), DOC(dai, node, NeuralNetwork, build))
+        .def("build", py::overload_cast<const NNArchiveVersionedConfig&, std::shared_ptr<Camera>, dai::NNModelDescription, float>(&NeuralNetwork::build), py::arg("nnArchiveVersionedConfig"), py::arg("input"), py::arg("modelDesc"), py::arg("fps"), DOC(dai, node, NeuralNetwork, build,2))
         .def("build", py::overload_cast<const NNArchiveConfig&, std::shared_ptr<Camera>, dai::NNModelDescription, float>(&NeuralNetwork::build), py::arg("nnArchiveConfig"), py::arg("input"), py::arg("modelDesc"), py::arg("fps"), DOC(dai, node, NeuralNetwork, build,2))
         .def("setBlob", py::overload_cast<dai::OpenVINO::Blob>(&NeuralNetwork::setBlob), py::arg("blob"), DOC(dai, node, NeuralNetwork, setBlob))
         .def("setBlob", py::overload_cast<const dai::Path&>(&NeuralNetwork::setBlob), py::arg("path"), DOC(dai, node, NeuralNetwork, setBlob, 2))

examples/cpp/v2_examples/NNArchive/nn_archive.cpp

@@ -35,7 +35,7 @@ void initializeNNArchiveFromServer(const std::string& urlBase,        // NOLINT(
     httplib::Client cli(urlBase);
     int64_t curPos = 0;
     uint64_t wholeBytesRead = 0;
-    const dai::NNArchiveConfig config([]() { return 0; },
+    const dai::NNArchiveVersionedConfig config([]() { return 0; },
                                       [&nnArchivePath, &wholeBytesRead, &curPos, &cli]() {
                                           auto bufferSize = 1024;
                                           const auto bytes = std::make_shared<std::vector<uint8_t>>();
@@ -116,7 +116,7 @@ bool initializeNNArchive(const std::string& urlBase,        // NOLINT(bugprone-e
         std::vector<uint8_t> fileContents((std::istreambuf_iterator<char>(archiveFile)), std::istreambuf_iterator<char>());
         detectionNetwork->setNNArchive(dai::NNArchive(fileContents));
     } else if(exampleType == "advanced") {
-        const dai::NNArchiveConfig config(nnArchivePath);
+        const dai::NNArchiveVersionedConfig config(nnArchivePath);
         const auto& configV1 = config.getConfigV1();
         if(!configV1) {
             throw std::runtime_error("Wrong config version");

examples/python/RVC2/ModelZoo/model_zoo.py

@@ -24,7 +24,7 @@
 archive = dai.NNArchive(modelPath)
 blob = archive.getSuperBlob().getBlobWithNumShaves(6)
 print("Number of shaves: ", blob.numShaves)
-print(archive.getConfig().getConfigV1().model.heads[0].metadata.classes)
+print(archive.getConfig().model.heads[0].metadata.classes)
 
 # Construct pipeline and start using downloaded NN model :-)
 with dai.Pipeline() as pipeline:

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
@@ -19,8 +18,20 @@
     camRgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)
     camRgb.setFps(15)
     nnArchive = dai.NNArchive(archivePath)
-    h, w = nnArchive.getConfig().getConfigV1().model.inputs[0].shape[-2:]
+    h, w = nnArchive.getConfig().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
     )

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