diff --git a/CMakeLists.txt b/CMakeLists.txt index af3d6a4a2b..54d34a7ac0 100644 --- a/CMakeLists.txt +++ b/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 diff --git a/bindings/python/src/modelzoo/ZooBindings.cpp b/bindings/python/src/modelzoo/ZooBindings.cpp index 0209060748..7a9fec3285 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,50 @@ 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_> model(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_> 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_> 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_> blobmodel(m_model, "BlobModel"); + blobmodel.def("settings", &depthai::model::BlobModel::settings, DOC(depthai::model, BlobModel, settings)); + // Bind SuperBlobModel + py::class_> superblobmodel(m_model, "SuperBlobModel"); + superblobmodel.def("settings", &depthai::model::SuperBlobModel::settings, DOC(depthai::model, SuperBlobModel, settings)); + + // Bind 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")); + m_zoo.def("load", &depthai::model::zoo::load, py::arg("ModelDescription")); + +} diff --git a/bindings/python/src/nn_archive/NNArchiveBindings.cpp b/bindings/python/src/nn_archive/NNArchiveBindings.cpp index d2f7834ba5..8070120eb4 100644 --- a/bindings/python/src/nn_archive/NNArchiveBindings.cpp +++ b/bindings/python/src/nn_archive/NNArchiveBindings.cpp @@ -5,7 +5,7 @@ #include #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(m, "NNArchive", DOC(dai, NNArchive)); py::class_ nnArchiveOptions(m, "NNArchiveOptions", DOC(dai, NNArchiveOptions)); - // NNArchiveConfig - py::class_ nnArchiveConfig(m, "NNArchiveConfig", DOC(dai, NNArchiveConfig)); + // NNArchiveVersionedConfig + py::class_ nnArchiveVersionedConfig(m, "NNArchiveVersionedConfig", DOC(dai, NNArchiveVersionedConfig)); + py::enum_ nnArchiveConfigVersion(m, "NNArchiveConfigVersion", DOC(dai, NNArchiveConfigVersion)); // NNArchiveEntry py::class_ 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, DOC(dai, NNArchive, getConfig)); + nnArchive.def("getConfigV1", &NNArchive::getConfig, 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(), + // Bind NNArchiveVersionedConfig + nnArchiveVersionedConfig.def(py::init(), py::arg("path"), py::arg("compression") = NNArchiveEntry::Compression::AUTO, - DOC(dai, NNArchiveConfig, NNArchiveConfig)); - nnArchiveConfig.def(py::init&, NNArchiveEntry::Compression>(), + DOC(dai, NNArchiveVersionedConfig, NNArchiveVersionedConfig)); + nnArchiveVersionedConfig.def(py::init&, NNArchiveEntry::Compression>(), py::arg("data"), py::arg("compression") = NNArchiveEntry::Compression::AUTO, - DOC(dai, NNArchiveConfig, NNArchiveConfig)); - nnArchiveConfig.def(py::init([](const std::function& openCallback, + DOC(dai, NNArchiveVersionedConfig, NNArchiveVersionedConfig)); + nnArchiveVersionedConfig.def(py::init([](const std::function& openCallback, const std::function()>& readCallback, const std::function& seekCallback, const std::function& skipCallback, const std::function& closeCallback, NNArchiveEntry::Compression compression) { auto readCallbackWrapper = [readCallback]() { return std::make_shared>(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, DOC(dai, NNArchiveVersionedConfig, getConfig)); + nnArchiveVersionedConfig.def("getConfigV1", &NNArchiveVersionedConfig::getConfig, 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); 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/bindings/python/src/pipeline/node/NeuralNetworkBindings.cpp b/bindings/python/src/pipeline/node/NeuralNetworkBindings.cpp index c6288efc70..39a0ebc61b 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"), @@ -53,6 +58,7 @@ void bind_neuralnetwork(pybind11::module& m, void* pCallstack){ .def("setNNArchive", py::overload_cast(&NeuralNetwork::setNNArchive), py::arg("nnArchive"), py::arg("numShaves"), DOC(dai, node, NeuralNetwork, setNNArchive, 2)) .def("setFromModelZoo", py::overload_cast(&NeuralNetwork::setFromModelZoo), py::arg("description"), py::arg("useCached"), DOC(dai, node, NeuralNetwork, setFromModelZoo)) .def("build", py::overload_cast(&NeuralNetwork::build), py::arg("input"), py::arg("nnArchive"), DOC(dai, node, NeuralNetwork, build)) + .def("build", py::overload_cast, 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, 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(&NeuralNetwork::setBlob), py::arg("blob"), DOC(dai, node, NeuralNetwork, setBlob)) .def("setBlob", py::overload_cast(&NeuralNetwork::setBlob), py::arg("path"), DOC(dai, node, NeuralNetwork, setBlob, 2)) diff --git a/examples/cpp/v2_examples/NNArchive/nn_archive.cpp b/examples/cpp/v2_examples/NNArchive/nn_archive.cpp index 8a1fcecfe2..ce6ff67d66 100644 --- a/examples/cpp/v2_examples/NNArchive/nn_archive.cpp +++ b/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>(); @@ -116,7 +116,7 @@ bool initializeNNArchive(const std::string& urlBase, // NOLINT(bugprone-e std::vector fileContents((std::istreambuf_iterator(archiveFile)), std::istreambuf_iterator()); 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"); diff --git a/examples/python/RVC2/ModelZoo/model_zoo.py b/examples/python/RVC2/ModelZoo/model_zoo.py index 44bcfd38dc..54cc2e0420 100644 --- a/examples/python/RVC2/ModelZoo/model_zoo.py +++ b/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: diff --git a/examples/python/RVC2/NNArchive/nn_archive.py b/examples/python/RVC2/NNArchive/nn_archive.py index a1d62aff79..2c939bf2af 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 @@ -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 ) 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/examples/python/RVC2/NNArchive/nn_archive_superblob.py b/examples/python/RVC2/NNArchive/nn_archive_superblob.py index 17ef46b7b2..3f27ead449 100644 --- a/examples/python/RVC2/NNArchive/nn_archive_superblob.py +++ b/examples/python/RVC2/NNArchive/nn_archive_superblob.py @@ -25,11 +25,8 @@ # There is no blob available assert archive.getBlob() is None -# One can access the NNArchive config as follows -config = archive.getConfig() - # You can access any config version -v1config: dai.nn_archive.v1.Config = config.getConfigV1() +v1config: dai.nn_archive.v1.Config = archive.getConfig() # Print some config fields print("-" * 10) 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..743eab5243 --- /dev/null +++ b/include/depthai/models/Models.hpp @@ -0,0 +1,129 @@ +#pragma once + +#include +#include +#include +#include +#include +#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 = ""; + 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 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(std::shared_ptr model, std::shared_ptr settings) : modelPtr_(model), settingsPtr_(settings) {} + + inline const dai::OpenVINO::Blob& model() const { + return *modelPtr_; + } + + inline const BlobSettings& settings() const { + return *settingsPtr_; + } + + private: + std::shared_ptr modelPtr_; + std::shared_ptr settingsPtr_; +}; + +class SuperBlobModel : public Model { + public: + ModelType type() const override { + return ModelType::SUPERBLOB; + } + + SuperBlobModel(std::shared_ptr model, std::shared_ptr settings) : modelPtr_(model), settingsPtr_(settings) {} + + inline const dai::OpenVINO::SuperBlob& model() const { + return *modelPtr_; + } + + inline const SuperBlobSettings& settings() const { + return *settingsPtr_; + } + + private: + std::shared_ptr modelPtr_; + std::shared_ptr settingsPtr_; +}; + +class DlcModel : public Model { + public: + ModelType type() const override { + return ModelType::DLC; + } + + 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_; + std::shared_ptr settingsPtr_; +}; + +using ModelVariant = std::variant; + +// Helper functions +ModelType getModelType(const ModelVariant& model); + +} // namespace model +} // namespace depthai 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/nn_archive/NNArchive.hpp b/include/depthai/nn_archive/NNArchive.hpp index 5836fafe10..7653188f00 100644 --- a/include/depthai/nn_archive/NNArchive.hpp +++ b/include/depthai/nn_archive/NNArchive.hpp @@ -4,7 +4,7 @@ #include #include -#include "depthai/nn_archive/NNArchiveConfig.hpp" +#include "depthai/nn_archive/NNArchiveVersionedConfig.hpp" #include "depthai/nn_archive/NNArchiveEntry.hpp" #include "depthai/openvino/OpenVINO.hpp" #include "depthai/utility/arg.hpp" @@ -57,11 +57,22 @@ class NNArchive { std::optional getModelPath() const; /** - * @brief Get NNArchive config, i.e. contents of `config.json` inside the archive. + * @brief Get NNArchive config wrapper * - * @return NNArchiveConfig: Archive config + * @return NNArchiveVersionedConfig: Archive config */ - const NNArchiveConfig& getConfig() const; + const NNArchiveVersionedConfig& getVersionedConfig() const; + + /** + * @brief Get NNArchive config. + * + * @tparam T: Type of config to get + * @return const T&: Config + */ + template + const T& getConfig() const { + return getVersionedConfig().getConfig(); + } /** * @brief Get type of model contained in NNArchive @@ -81,7 +92,7 @@ class NNArchive { NNArchiveOptions archiveOptions; // Archive config - std::shared_ptr archiveConfigPtr; + std::shared_ptr archiveVersionedConfigPtr; // Blob related stuff std::shared_ptr blobPtr; diff --git a/include/depthai/nn_archive/NNArchiveConfig.hpp b/include/depthai/nn_archive/NNArchiveConfig.hpp deleted file mode 100644 index 4a7077d483..0000000000 --- a/include/depthai/nn_archive/NNArchiveConfig.hpp +++ /dev/null @@ -1,54 +0,0 @@ -#pragma once - -// C std -#include - -// C++ std -#include -#include -#include - -// libraries -#include - -// internal -#include "depthai/nn_archive/NNArchiveEntry.hpp" -#include "depthai/nn_archive/v1/Config.hpp" -#include "depthai/utility/Path.hpp" - -namespace dai { - -class NNArchiveConfig { - public: - /** - * @data Should point to a whole compressed NNArchive read to memory if compression is not set to RAW_FS. - * If compression is set to RAW_FS, then this should point to just the config.json file read to memory. - */ - explicit NNArchiveConfig(const std::vector& data, NNArchiveEntry::Compression compression = NNArchiveEntry::Compression::AUTO); - /** - * @path Should point to: - * 1) if compression is set to RAW_FS: to just the config.json file. - * 2) if compression is set to AUTO: to whole compressed NNArchive or just the config.json file which must end in .json . - * 3) else: to whole compressed NNArchive. - * see NNArchive class for parameter explanation - */ - explicit NNArchiveConfig(const dai::Path& path, NNArchiveEntry::Compression compression = NNArchiveEntry::Compression::AUTO); - /** - * Returned data should be just the config.json if compression == RAW_FS or the whole NNArchive otherwise - * see NNArchive class for parameter explanation - */ - NNArchiveConfig(const std::function& openCallback, - const std::function>()>& readCallback, - const std::function& seekCallback, - const std::function& skipCallback, - const std::function& closeCallback, - NNArchiveEntry::Compression compression = NNArchiveEntry::Compression::AUTO); - - std::optional getConfigV1() const; - - private: - class Impl; - spimpl::impl_ptr pimpl; -}; - -} // namespace dai diff --git a/include/depthai/nn_archive/NNArchiveVersionedConfig.hpp b/include/depthai/nn_archive/NNArchiveVersionedConfig.hpp new file mode 100644 index 0000000000..5289e864b1 --- /dev/null +++ b/include/depthai/nn_archive/NNArchiveVersionedConfig.hpp @@ -0,0 +1,96 @@ +#pragma once + +// C std +#include + +// C++ std +#include +#include +#include +#include +#include + +// libraries +#include + +// internal +#include "depthai/nn_archive/NNArchiveEntry.hpp" +#include "depthai/nn_archive/v1/Config.hpp" +#include "depthai/utility/Path.hpp" + +namespace dai { + +enum class NNArchiveConfigVersion { + V1, +}; + +using NNArchiveConfig = std::variant; + +// This assumes that the config is always stored in the same order as the enum +inline NNArchiveConfigVersion getNNArchiveConfigVersion(const NNArchiveConfig& config) { + return static_cast(config.index()); +} + +class NNArchiveVersionedConfig { + public: + /** + * @data Should point to a whole compressed NNArchive read to memory if compression is not set to RAW_FS. + * If compression is set to RAW_FS, then this should point to just the config.json file read to memory. + */ + explicit NNArchiveVersionedConfig(const std::vector& data, NNArchiveEntry::Compression compression = NNArchiveEntry::Compression::AUTO); + + /** + * @path Should point to: + * 1) if compression is set to RAW_FS: to just the config.json file. + * 2) if compression is set to AUTO: to whole compressed NNArchive or just the config.json file which must end in .json . + * 3) else: to whole compressed NNArchive. + * see NNArchive class for parameter explanation + */ + explicit NNArchiveVersionedConfig(const dai::Path& path, NNArchiveEntry::Compression compression = NNArchiveEntry::Compression::AUTO); + + /** + * Returned data should be just the config.json if compression == RAW_FS or the whole NNArchive otherwise + * see NNArchive class for parameter explanation + */ + NNArchiveVersionedConfig(const std::function& openCallback, + const std::function>()>& readCallback, + const std::function& seekCallback, + const std::function& skipCallback, + const std::function& closeCallback, + NNArchiveEntry::Compression compression = NNArchiveEntry::Compression::AUTO); + + /** + * @brief Construct NNArchiveVersionedConfig from a specific version of a config. + * @param version: Version of the config. + */ + NNArchiveVersionedConfig(const NNArchiveConfig& config) : version(getNNArchiveConfigVersion(config)), config(config) {} + + /** + * @brief Get version of the underlying config. + */ + NNArchiveConfigVersion getVersion() const { + return version; + } + + /** + * @brief Get stored config cast to a specific version. + * @tparam T: Config type to cast to. + */ + template + const T& getConfig() const { + return std::get(config); + } + + private: + void initConfig(const std::optional& maybeJson); + + NNArchiveConfigVersion version; + NNArchiveConfig config; +}; + +template <> +inline const NNArchiveConfig& NNArchiveVersionedConfig::getConfig() const { + return config; +} + +} // namespace dai diff --git a/include/depthai/openvino/OpenVINO.hpp b/include/depthai/openvino/OpenVINO.hpp index 6e42183014..a1cb5460a0 100644 --- a/include/depthai/openvino/OpenVINO.hpp +++ b/include/depthai/openvino/OpenVINO.hpp @@ -76,13 +76,20 @@ 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 * * @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 +106,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..52cb6f2f32 100644 --- a/include/depthai/pipeline/node/DetectionNetwork.hpp +++ b/include/depthai/pipeline/node/DetectionNetwork.hpp @@ -28,7 +28,32 @@ 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); + 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 f38c8a2649..390da547b0 100644 --- a/include/depthai/pipeline/node/DetectionParser.hpp +++ b/include/depthai/pipeline/node/DetectionParser.hpp @@ -3,12 +3,14 @@ #include #include +#include + // standard #include // shared #include -#include +#include #include #include #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. * @@ -149,18 +174,18 @@ class DetectionParser : public DeviceNodeCRTP> mClasses; std::optional mArchive; - std::optional archiveConfig; + std::optional archiveConfig; }; } // namespace node diff --git a/include/depthai/pipeline/node/NeuralNetwork.hpp b/include/depthai/pipeline/node/NeuralNetwork.hpp index 1cdba3298f..718ae059a1 100644 --- a/include/depthai/pipeline/node/NeuralNetwork.hpp +++ b/include/depthai/pipeline/node/NeuralNetwork.hpp @@ -1,10 +1,12 @@ #pragma once +#include #include #include #include #include "depthai/nn_archive/NNArchive.hpp" +#include "depthai/nn_archive/NNArchiveVersionedConfig.hpp" #include "depthai/openvino/OpenVINO.hpp" // standard @@ -37,7 +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 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 @@ -74,6 +83,23 @@ class NeuralNetwork : public DeviceNodeCRTPsetModel(p); }, model); + } + + void setModel(const depthai::model::BlobModel& model) { + setBlob(model.model()); + } + + void setModel(const depthai::model::SuperBlobModel& model) { + int numShaves = model.settings().numShaves; + setBlob(model.model().getBlobWithNumShaves(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..73459bf9c8 --- /dev/null +++ b/src/models/ModelLoader.cpp @@ -0,0 +1,214 @@ +#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; + auto blobPtr = std::make_shared(path); + auto settingsPtr = std::make_shared(); + modelVariant_ = BlobModel(blobPtr, settingsPtr); +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +void ModelLoader::loadSuperblobModel(std::vector data) { + DAI_CHECK_V(false, "Not implemented"); +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +void ModelLoader::loadSuperblobModel(const std::string& path) { + modelType_ = ModelType::SUPERBLOB; + auto superblobPtr = std::make_shared(path); + auto settingsPtr = std::make_shared(); + modelVariant_ = SuperBlobModel(superblobPtr, settingsPtr); +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +void ModelLoader::loadDlcModel(std::vector data) { + DAI_CHECK_V(false, "Not implemented"); +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +void ModelLoader::loadDlcModel(const std::string& path) { + modelType_ = ModelType::DLC; + auto settingsPtr = std::make_shared(); + modelVariant_ = DlcModel(path, settingsPtr); +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +void ModelLoader::loadNNArchive(std::vector data) { + DAI_CHECK_V(false, "Not implemented"); +} + +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +/**********************************************************************************************************************/ +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; + { + 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; + { + 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: + 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..e6709f4380 --- /dev/null +++ b/src/models/main.cpp @@ -0,0 +1,24 @@ +#include +#include +#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; + + return 0; +} \ No newline at end of file diff --git a/src/nn_archive/NNArchive.cpp b/src/nn_archive/NNArchive.cpp index 7447627d02..caae642ec3 100644 --- a/src/nn_archive/NNArchive.cpp +++ b/src/nn_archive/NNArchive.cpp @@ -1,4 +1,5 @@ #include "depthai/nn_archive/NNArchive.hpp" +#include "depthai/nn_archive/NNArchiveVersionedConfig.hpp" // internal private #include "common/ModelType.hpp" @@ -12,10 +13,11 @@ NNArchive::NNArchive(const std::string& archivePath, NNArchiveOptions options) : if(!std::filesystem::exists(archivePath)) DAI_CHECK_V(false, "Archive file does not exist: {}", archivePath); // Read config - archiveConfigPtr.reset(new NNArchiveConfig(archivePath, archiveOptions.compression())); + archiveVersionedConfigPtr.reset(new NNArchiveVersionedConfig(archivePath, archiveOptions.compression())); - // Based on the config, read model path in archive - std::string modelPathInArchive = archiveConfigPtr->getConfigV1()->model.metadata.path; + // Only V1 config is supported at the moment + DAI_CHECK(archiveVersionedConfigPtr->getVersion() == NNArchiveConfigVersion::V1, "Only V1 config is supported at the moment"); + std::string modelPathInArchive = archiveVersionedConfigPtr->getConfig().model.metadata.path; // Read archive type modelType = model::readModelType(modelPathInArchive); @@ -102,8 +104,8 @@ model::ModelType NNArchive::getModelType() const { return modelType; } -const NNArchiveConfig& NNArchive::getConfig() const { - return *archiveConfigPtr; +const NNArchiveVersionedConfig& NNArchive::getVersionedConfig() const { + return *archiveVersionedConfigPtr; } std::vector NNArchive::readModelFromArchive(const std::string& archivePath, const std::string& modelPathInArchive) const { diff --git a/src/nn_archive/NNArchiveBlob.cpp b/src/nn_archive/NNArchiveBlob.cpp index 26ea86c886..8044498881 100644 --- a/src/nn_archive/NNArchiveBlob.cpp +++ b/src/nn_archive/NNArchiveBlob.cpp @@ -7,7 +7,7 @@ #include // internal private -#include "nn_archive/NNArchiveConfigHelper.hpp" +#include "nn_archive/NNArchiveVersionedConfigHelper.hpp" #include "utility/ArchiveUtil.hpp" #include "utility/ErrorMacros.hpp" @@ -18,11 +18,11 @@ class NNArchiveBlob::Impl { std::optional mBlob; template - void init(const NNArchiveConfig& config, const T& dataOrPath, NNArchiveEntry::Compression compression) { + void init(const NNArchiveVersionedConfig& config, const T& dataOrPath, NNArchiveEntry::Compression compression) { if(compression == NNArchiveEntry::Compression::RAW_FS) { mBlob.emplace(OpenVINO::Blob(dataOrPath)); } else { - const auto& blobPath = nn_archive::NNArchiveConfigHelper::getBlobPath(config); + const auto& blobPath = nn_archive::NNArchiveVersionedConfigHelper::getBlobPath(config); utility::ArchiveUtil archive(dataOrPath, compression); std::vector blobBytes; const bool success = archive.readEntry(blobPath, blobBytes); @@ -31,15 +31,15 @@ class NNArchiveBlob::Impl { } } - Impl(const NNArchiveConfig& config, const std::vector& data, NNArchiveEntry::Compression compression) { + Impl(const NNArchiveVersionedConfig& config, const std::vector& data, NNArchiveEntry::Compression compression) { init(config, data, compression); } - Impl(const NNArchiveConfig& config, const Path& path, NNArchiveEntry::Compression compression) { + Impl(const NNArchiveVersionedConfig& config, const Path& path, NNArchiveEntry::Compression compression) { init(config, path, compression); } - Impl(const NNArchiveConfig& config, + Impl(const NNArchiveVersionedConfig& config, const std::function& openCallback, const std::function>()>& readCallback, const std::function& seekCallback, @@ -49,7 +49,7 @@ class NNArchiveBlob::Impl { if(compression == NNArchiveEntry::Compression::RAW_FS) { DAI_CHECK(false, "RAW_FS with callbacks NOT IMPLEMENTED YET for NNArchiveBlob"); } else { - const auto& blobPath = nn_archive::NNArchiveConfigHelper::getBlobPath(config); + const auto& blobPath = nn_archive::NNArchiveVersionedConfigHelper::getBlobPath(config); utility::ArchiveUtil archive(openCallback, readCallback, seekCallback, skipCallback, closeCallback, compression); std::vector blobBytes; const bool success = archive.readEntry(blobPath, blobBytes); @@ -63,13 +63,13 @@ class NNArchiveBlob::Impl { } }; -NNArchiveBlob::NNArchiveBlob(const NNArchiveConfig& config, const std::vector& data, NNArchiveEntry::Compression compression) +NNArchiveBlob::NNArchiveBlob(const NNArchiveVersionedConfig& config, const std::vector& data, NNArchiveEntry::Compression compression) : pimpl(spimpl::make_impl(config, data, compression)) {}; -NNArchiveBlob::NNArchiveBlob(const NNArchiveConfig& config, const Path& path, NNArchiveEntry::Compression compression) +NNArchiveBlob::NNArchiveBlob(const NNArchiveVersionedConfig& config, const Path& path, NNArchiveEntry::Compression compression) : pimpl(spimpl::make_impl(config, path, compression)) {} -NNArchiveBlob::NNArchiveBlob(const NNArchiveConfig& config, +NNArchiveBlob::NNArchiveBlob(const NNArchiveVersionedConfig& config, const std::function& openCallback, const std::function>()>& readCallback, const std::function& seekCallback, diff --git a/src/nn_archive/NNArchiveConfig.cpp b/src/nn_archive/NNArchiveConfig.cpp deleted file mode 100644 index 72ef26aaa9..0000000000 --- a/src/nn_archive/NNArchiveConfig.cpp +++ /dev/null @@ -1,113 +0,0 @@ -#include "depthai/nn_archive/NNArchiveConfig.hpp" - -// C++ std -#include -#include -#include - -// libraries -#include - -#include - -// internal public -#include "depthai/nn_archive/v1/Config.hpp" - -// internal private -#include "nn_archive/v1/Generators.hpp" -#include "utility/ArchiveUtil.hpp" -#include "utility/ErrorMacros.hpp" - -namespace dai { - -class NNArchiveConfig::Impl { - public: - // in the future we will have - std::optional> mConfig; - - void initConfig(const std::optional& maybeJson) { - DAI_CHECK_IN(maybeJson); - dai::nn_archive::v1::Config config; - dai::nn_archive::v1::from_json(*maybeJson, config); - mConfig = config; - } - - Impl(const std::vector& data, NNArchiveEntry::Compression compression) { - bool isJson = compression == NNArchiveEntry::Compression::RAW_FS; - std::optional maybeJson; - if(isJson) { - maybeJson = nlohmann::json::parse(data); - } else { - utility::ArchiveUtil archive(data, compression); - std::vector jsonBytes; - const bool success = archive.readEntry("config.json", jsonBytes); - DAI_CHECK(success, "Didn't find the config.json file inside the NNArchive read from memory."); - maybeJson = nlohmann::json::parse(jsonBytes); - } - initConfig(maybeJson); - } - - Impl(const Path& path, NNArchiveEntry::Compression compression) { - bool isJson = - compression == NNArchiveEntry::Compression::RAW_FS || (compression == NNArchiveEntry::Compression::AUTO && utility::ArchiveUtil::isJsonPath(path)); - std::optional maybeJson; - if(isJson) { - std::ifstream jsonStream(path); - maybeJson = nlohmann::json::parse(jsonStream); - } else { - utility::ArchiveUtil archive(path, compression); - std::vector jsonBytes; - const bool success = archive.readEntry("config.json", jsonBytes); - DAI_CHECK_V(success, "Didn't find the config.json file inside the {} archive.", path); - maybeJson = nlohmann::json::parse(jsonBytes); - } - initConfig(maybeJson); - } - - Impl(const std::function& openCallback, - const std::function>()>& readCallback, - const std::function& seekCallback, - const std::function& skipCallback, - const std::function& closeCallback, - NNArchiveEntry::Compression compression) { - std::optional maybeJson; - if(compression == NNArchiveEntry::Compression::RAW_FS) { - DAI_CHECK(false, "RAW_FS with callbacks NOT IMPLEMENTED YET for NNArchiveConfig"); - } else { - utility::ArchiveUtil archive(openCallback, readCallback, seekCallback, skipCallback, closeCallback, compression); - std::vector jsonBytes; - const bool success = archive.readEntry("config.json", jsonBytes); - DAI_CHECK(success, "Didn't find the config.json file inside the NNArchive."); - maybeJson = nlohmann::json::parse(jsonBytes); - } - initConfig(maybeJson); - } - - template - std::optional getConfig() const { - DAI_CHECK_IN(mConfig); - if(const auto* configPtr(std::get_if(&(*mConfig))); configPtr) { - return *configPtr; - } - return std::nullopt; - } -}; - -NNArchiveConfig::NNArchiveConfig(const std::vector& data, NNArchiveEntry::Compression compression) - : pimpl(spimpl::make_impl(data, compression)){}; - -NNArchiveConfig::NNArchiveConfig(const Path& path, NNArchiveEntry::Compression compression) : pimpl(spimpl::make_impl(path, compression)) {} - -NNArchiveConfig::NNArchiveConfig(const std::function& openCallback, - const std::function>()>& readCallback, - const std::function& seekCallback, - const std::function& skipCallback, - const std::function& closeCallback, - NNArchiveEntry::Compression compression) - : pimpl(spimpl::make_impl(openCallback, readCallback, seekCallback, skipCallback, closeCallback, compression)) {} - -std::optional NNArchiveConfig::getConfigV1() const { - return pimpl->getConfig(); -} - -} // namespace dai diff --git a/src/nn_archive/NNArchiveVersionedConfig.cpp b/src/nn_archive/NNArchiveVersionedConfig.cpp new file mode 100644 index 0000000000..90faf7a437 --- /dev/null +++ b/src/nn_archive/NNArchiveVersionedConfig.cpp @@ -0,0 +1,86 @@ +#include "depthai/nn_archive/NNArchiveVersionedConfig.hpp" + +// C++ std +#include +#include +#include + +// libraries +#include + +#include + +// internal public +#include "depthai/nn_archive/v1/Config.hpp" + +// internal private +#include "nn_archive/v1/Generators.hpp" +#include "utility/ArchiveUtil.hpp" +#include "utility/ErrorMacros.hpp" + +namespace dai { + +NNArchiveVersionedConfig::NNArchiveVersionedConfig(const std::vector& data, NNArchiveEntry::Compression compression) { + bool isJson = compression == NNArchiveEntry::Compression::RAW_FS; + std::optional maybeJson; + if(isJson) { + maybeJson = nlohmann::json::parse(data); + } else { + utility::ArchiveUtil archive(data, compression); + std::vector jsonBytes; + const bool success = archive.readEntry("config.json", jsonBytes); + DAI_CHECK(success, "Didn't find the config.json file inside the NNArchive read from memory."); + maybeJson = nlohmann::json::parse(jsonBytes); + } + initConfig(maybeJson); +} + +NNArchiveVersionedConfig::NNArchiveVersionedConfig(const Path& path, NNArchiveEntry::Compression compression) { + bool isJson = + compression == NNArchiveEntry::Compression::RAW_FS || (compression == NNArchiveEntry::Compression::AUTO && utility::ArchiveUtil::isJsonPath(path)); + std::optional maybeJson; + if(isJson) { + std::ifstream jsonStream(path); + maybeJson = nlohmann::json::parse(jsonStream); + } else { + utility::ArchiveUtil archive(path, compression); + std::vector jsonBytes; + const bool success = archive.readEntry("config.json", jsonBytes); + DAI_CHECK_V(success, "Didn't find the config.json file inside the {} archive.", path); + maybeJson = nlohmann::json::parse(jsonBytes); + } + initConfig(maybeJson); +} + +NNArchiveVersionedConfig::NNArchiveVersionedConfig(const std::function& openCallback, + const std::function>()>& readCallback, + const std::function& seekCallback, + const std::function& skipCallback, + const std::function& closeCallback, + NNArchiveEntry::Compression compression) { + std::optional maybeJson; + if(compression == NNArchiveEntry::Compression::RAW_FS) { + DAI_CHECK(false, "RAW_FS with callbacks NOT IMPLEMENTED YET for NNArchiveVersionedConfig"); + } else { + utility::ArchiveUtil archive(openCallback, readCallback, seekCallback, skipCallback, closeCallback, compression); + std::vector jsonBytes; + const bool success = archive.readEntry("config.json", jsonBytes); + DAI_CHECK(success, "Didn't find the config.json file inside the NNArchive."); + maybeJson = nlohmann::json::parse(jsonBytes); + } + initConfig(maybeJson); +} + +void NNArchiveVersionedConfig::initConfig(const std::optional& maybeJson) { + DAI_CHECK_IN(maybeJson); + + // TODO: Check what version of config this is - for now we will fall back to V1 + dai::nn_archive::v1::Config configV1; + dai::nn_archive::v1::from_json(*maybeJson, configV1); + config = configV1; + + // TODO: This should be loaded from the json + version = NNArchiveConfigVersion::V1; +} + +} // namespace dai 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]; } diff --git a/src/pipeline/node/DetectionNetwork.cpp b/src/pipeline/node/DetectionNetwork.cpp index 85750bcf72..3f26ecb1d3 100644 --- a/src/pipeline/node/DetectionNetwork.cpp +++ b/src/pipeline/node/DetectionNetwork.cpp @@ -12,7 +12,7 @@ #include "depthai/depthai.hpp" #include "depthai/modelzoo/Zoo.hpp" #include "depthai/nn_archive/NNArchive.hpp" -#include "nn_archive/NNArchiveConfig.hpp" +#include "nn_archive/NNArchiveVersionedConfig.hpp" #include "pipeline/DeviceNodeGroup.hpp" #include "utility/ArchiveUtil.hpp" #include "utility/ErrorMacros.hpp" @@ -57,25 +57,19 @@ std::shared_ptr DetectionNetwork::build(Node::Output& input, c std::shared_ptr DetectionNetwork::build(std::shared_ptr camera, dai::NNModelDescription modelDesc, float fps) { setFromModelZoo(modelDesc); - // Get the input size - auto nnArchiveConfig = detectionParser->getNNArchiveConfig().getConfigV1(); - if(!nnArchiveConfig.has_value()) { - DAI_CHECK_V(false, "The DetectionNetwork.build method only supports for NNConfigV1"); - } - if(nnArchiveConfig->model.inputs.size() != 1) { - DAI_CHECK_V(false, "Only single input model is supported"); - } - if(nnArchiveConfig->model.inputs[0].shape.size() != 4) { - DAI_CHECK_V(false, "Only 4D input shape is supported"); - } + // We only support config v1 for now + DAI_CHECK_V(detectionParser->getNNArchiveVersionedConfig().getVersion() == NNArchiveConfigVersion::V1, "Only NNConfigV1 is supported for DetectionNetwork"); + auto nnArchiveConfigV1 = detectionParser->getNNArchiveVersionedConfig().getConfig(); - // Check that the first two dimesions are 1 and 3 - if(nnArchiveConfig->model.inputs[0].shape[0] != 1 || nnArchiveConfig->model.inputs[0].shape[1] != 3) { - DAI_CHECK_V(false, "Only 3 channel input is supported"); - } - auto inputHeight = nnArchiveConfig->model.inputs[0].shape[2]; - auto inputWidth = nnArchiveConfig->model.inputs[0].shape[3]; + // Perform basic checks + DAI_CHECK_V(nnArchiveConfigV1.model.inputs.size() == 1, "Only single input model is supported"); + DAI_CHECK_V(nnArchiveConfigV1.model.inputs[0].shape.size() == 4, "Only 4D input shape is supported"); + DAI_CHECK_V(nnArchiveConfigV1.model.inputs[0].shape[0] == 1 && nnArchiveConfigV1.model.inputs[0].shape[1] == 3, "Only 3 channel input is supported"); + + // Unpack input dimensions + auto inputHeight = nnArchiveConfigV1.model.inputs[0].shape[2]; + auto inputWidth = nnArchiveConfigV1.model.inputs[0].shape[3]; auto type = dai::ImgFrame::Type::BGR888p; auto platform = getDevice()->getPlatform(); @@ -230,12 +224,12 @@ float DetectionNetwork::getConfidenceThreshold() const { } std::vector>> DetectionNetwork::getRequiredInputs() { - const dai::NNArchiveConfig& config = detectionParser->getNNArchiveConfig(); - const auto configV1 = config.getConfigV1(); - DAI_CHECK(configV1.has_value(), "Only NNConfigV1 is supported for DetectionNetwork"); + const dai::NNArchiveVersionedConfig& config = detectionParser->getNNArchiveVersionedConfig(); + DAI_CHECK(config.getVersion() == NNArchiveConfigVersion::V1, "Only NNConfigV1 is supported for DetectionNetwork"); + const auto configV1 = config.getConfig(); - const auto width = configV1->model.inputs[0].shape[2]; - const auto height = configV1->model.inputs[0].shape[3]; + const auto width = configV1.model.inputs[0].shape[2]; + const auto height = configV1.model.inputs[0].shape[3]; auto cap = std::make_shared(); cap->size.value = std::pair(width, height); diff --git a/src/pipeline/node/DetectionParser.cpp b/src/pipeline/node/DetectionParser.cpp index e8cf1d2b7e..769b5461d3 100644 --- a/src/pipeline/node/DetectionParser.cpp +++ b/src/pipeline/node/DetectionParser.cpp @@ -46,17 +46,16 @@ void DetectionParser::setModelPath(const dai::Path& modelPath) { } } -const NNArchiveConfig& DetectionParser::getNNArchiveConfig() const { - DAI_CHECK_V(archiveConfig.has_value(), "NNArchiveConfig is not set. Use setNNArchive(...) first."); +const NNArchiveVersionedConfig& DetectionParser::getNNArchiveVersionedConfig() const { + DAI_CHECK_V(archiveConfig.has_value(), "NNArchiveVersionedConfig is not set. Use setNNArchive(...) first."); return archiveConfig.value(); } -void DetectionParser::setConfig(const dai::NNArchiveConfig& config) { +void DetectionParser::setConfig(const dai::NNArchiveVersionedConfig& config) { archiveConfig = config; - DAI_CHECK_V(config.getConfigV1().has_value(), "Only NNArchive config V1 is supported."); - - auto configV1 = *config.getConfigV1(); + DAI_CHECK_V(config.getVersion() == NNArchiveConfigVersion::V1, "Only NNArchive config V1 is supported."); + auto configV1 = config.getConfig(); const auto model = configV1.model; // TODO(jakgra) is NN Archive valid without this? why is this optional? @@ -108,18 +107,18 @@ void DetectionParser::setConfig(const dai::NNArchiveConfig& config) { void DetectionParser::setNNArchiveBlob(const NNArchive& nnArchive) { DAI_CHECK_V(nnArchive.getModelType() == dai::model::ModelType::BLOB, "NNArchive type is not BLOB"); - setConfig(nnArchive.getConfig()); + setConfig(nnArchive.getVersionedConfig()); setBlob(nnArchive.getBlob().value()); } void DetectionParser::setNNArchiveSuperblob(const NNArchive& nnArchive, int numShaves) { DAI_CHECK_V(nnArchive.getModelType() == dai::model::ModelType::SUPERBLOB, "NNArchive type is not SUPERBLOB"); - setConfig(nnArchive.getConfig()); + setConfig(nnArchive.getVersionedConfig()); setBlob(nnArchive.getSuperBlob()->getBlobWithNumShaves(numShaves)); } void DetectionParser::setNNArchiveOther(const NNArchive& nnArchive) { - setConfig(nnArchive.getConfig()); + setConfig(nnArchive.getVersionedConfig()); } void DetectionParser::setBlob(OpenVINO::Blob blob) { diff --git a/src/pipeline/node/NeuralNetwork.cpp b/src/pipeline/node/NeuralNetwork.cpp index 74b561727a..576e81f827 100644 --- a/src/pipeline/node/NeuralNetwork.cpp +++ b/src/pipeline/node/NeuralNetwork.cpp @@ -25,27 +25,23 @@ std::shared_ptr NeuralNetwork::build(Node::Output& input, const N return std::static_pointer_cast(shared_from_this()); } -std::shared_ptr NeuralNetwork::build(const NNArchiveConfig& nnArchiveCfg, std::shared_ptr camera, dai::NNModelDescription modelDesc, float fps){ +std::shared_ptr NeuralNetwork::build(const NNArchiveConfig& nnArchiveConfig, std::shared_ptr input, dai::NNModelDescription modelDesc, float fps) { + return build(NNArchiveVersionedConfig(nnArchiveConfig), input, modelDesc, fps); +} + +std::shared_ptr NeuralNetwork::build(const NNArchiveVersionedConfig& nnArchiveCfg, std::shared_ptr camera, dai::NNModelDescription modelDesc, float fps){ setFromModelZoo(modelDesc); - // Get the input size - auto nnArchiveConfig = nnArchiveCfg.getConfigV1(); - if(!nnArchiveConfig.has_value()) { - DAI_CHECK_V(false, "The NeuralNetwork.build method only supports for NNConfigV1"); - } - if(nnArchiveConfig->model.inputs.size() != 1) { - DAI_CHECK_V(false, "Only single input model is supported"); - } - if(nnArchiveConfig->model.inputs[0].shape.size() != 4) { - DAI_CHECK_V(false, "Only 4D input shape is supported"); - } + DAI_CHECK_V(nnArchiveCfg.getVersion() == dai::NNArchiveConfigVersion::V1, "Only V1 configs are supported for NeuralNetwork.build method"); - // Check that the first two dimesions are 1 and 3 - if(nnArchiveConfig->model.inputs[0].shape[0] != 1 || nnArchiveConfig->model.inputs[0].shape[1] != 3) { - DAI_CHECK_V(false, "Only 3 channel input is supported"); - } - auto inputHeight = nnArchiveConfig->model.inputs[0].shape[2]; - auto inputWidth = nnArchiveConfig->model.inputs[0].shape[3]; + const auto& configV1 = nnArchiveCfg.getConfig(); + + DAI_CHECK_V(configV1.model.inputs.size() == 1, "Only single input model is supported"); + DAI_CHECK_V(configV1.model.inputs[0].shape.size() == 4, "Only 4D input shape is supported"); + DAI_CHECK_V(configV1.model.inputs[0].shape[0] == 1 && configV1.model.inputs[0].shape[1] == 3, "Only 3 channel input is supported"); + + auto inputHeight = configV1.model.inputs[0].shape[2]; + auto inputWidth = configV1.model.inputs[0].shape[3]; auto type = dai::ImgFrame::Type::BGR888p; auto platform = getDevice()->getPlatform(); @@ -62,9 +58,7 @@ std::shared_ptr NeuralNetwork::build(const NNArchiveConfig& nnArc cap.type = type; cap.fps.value = fps; auto* input = camera->requestOutput(cap, false); - if(!input) { - DAI_CHECK_V(false, "Camera does not have output with requested capabilities"); - } + DAI_CHECK_V(input != nullptr, "Camera does not have output with requested capabilities"); input->link(this->input); return std::static_pointer_cast(shared_from_this()); } diff --git a/src/pipeline/node/SpatialDetectionNetwork.cpp b/src/pipeline/node/SpatialDetectionNetwork.cpp index b8d24f0b76..4485de39fc 100644 --- a/src/pipeline/node/SpatialDetectionNetwork.cpp +++ b/src/pipeline/node/SpatialDetectionNetwork.cpp @@ -38,25 +38,16 @@ std::shared_ptr SpatialDetectionNetwork::build(std::sha dai::NNModelDescription modelDesc, float fps) { setFromModelZoo(modelDesc); - // Get the input size - auto nnArchiveConfig = detectionParser->getNNArchiveConfig().getConfigV1(); - if(!nnArchiveConfig.has_value()) { - DAI_CHECK_V(false, "The DetectionNetwork.build method only supports for NNConfigV1"); - } - if(nnArchiveConfig->model.inputs.size() != 1) { - DAI_CHECK_V(false, "Only single input model is supported"); - } - if(nnArchiveConfig->model.inputs[0].shape.size() != 4) { - DAI_CHECK_V(false, "Only 4D input shape is supported"); - } + DAI_CHECK(detectionParser->getNNArchiveVersionedConfig().getVersion() == NNArchiveConfigVersion::V1, "Only NNArchive config V1 is supported."); + auto configV1 = detectionParser->getNNArchiveVersionedConfig().getConfig(); - // Check that the first two dimesions are 1 and 3 - if(nnArchiveConfig->model.inputs[0].shape[0] != 1 || nnArchiveConfig->model.inputs[0].shape[1] != 3) { - DAI_CHECK_V(false, "Only 3 channel input is supported"); - } - auto inputHeight = nnArchiveConfig->model.inputs[0].shape[2]; - auto inputWidth = nnArchiveConfig->model.inputs[0].shape[3]; + DAI_CHECK(configV1.model.inputs.size() == 1, "Only single input model is supported"); + DAI_CHECK(configV1.model.inputs[0].shape.size() == 4, "Only 4D input shape is supported"); + DAI_CHECK(configV1.model.inputs[0].shape[0] == 1 && configV1.model.inputs[0].shape[1] == 3, "Only 3 channel input is supported"); + + auto inputHeight = configV1.model.inputs[0].shape[2]; + auto inputWidth = configV1.model.inputs[0].shape[3]; auto type = dai::ImgFrame::Type::BGR888p; auto platform = getDevice()->getPlatform(); diff --git a/tests/src/ondevice_tests/pipeline/node/detection_network_test.cpp b/tests/src/ondevice_tests/pipeline/node/detection_network_test.cpp index 77bfed2382..92d1f51c4d 100644 --- a/tests/src/ondevice_tests/pipeline/node/detection_network_test.cpp +++ b/tests/src/ondevice_tests/pipeline/node/detection_network_test.cpp @@ -28,7 +28,7 @@ TEST_CASE("DetectionNetwork can load BLOB properly") { // Classes match auto loadedClasses = nn->getClasses().value(); - auto archiveClasses = nnArchive.getConfig().getConfigV1()->model.heads->at(0).metadata.classes.value(); + auto archiveClasses = nnArchive.getConfig().model.heads->at(0).metadata.classes.value(); REQUIRE(loadedClasses == archiveClasses); // Throws if number of shaves is specified