Lazy loading large assets

bindings/python/src/nn_archive/NNArchiveBindings.cpp

@@ -79,7 +79,14 @@ void NNArchiveBindings::bind(pybind11::module& m, void* pCallstack) {
                   DOC(dai, NNArchive, NNArchive));
     nnArchive.def("getBlob", &NNArchive::getBlob, DOC(dai, NNArchive, getBlob));
     nnArchive.def("getSuperBlob", &NNArchive::getSuperBlob, DOC(dai, NNArchive, getBlob));
-    nnArchive.def("getOtherModelFormat", &NNArchive::getOtherModelFormat, DOC(dai, NNArchive, getOtherModelFormat));
+    nnArchive.def(
+        "getOtherModelFormat",
+        [](const NNArchive& archive) -> py::object {
+            const auto model = archive.getOtherModelFormat();
+            if(!model.has_value()) return py::none();
+            return py::bytes(reinterpret_cast<const char*>(model->data()), model->size());
+        },
+        DOC(dai, NNArchive, getOtherModelFormat));
     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));

bindings/python/src/pipeline/AssetManagerBindings.cpp

@@ -32,11 +32,14 @@ void AssetManagerBindings::bind(pybind11::module& m, void* pCallstack) {
             "data",
             [](py::object& obj) {
                 dai::Asset& a = obj.cast<dai::Asset&>();
-                return py::array_t<std::uint8_t>(a.data.size(), a.data.data(), obj);
+                auto& data = a.getData();
+                return py::array_t<std::uint8_t>(data.size(), data.data(), obj);
             },
             [](py::object& obj, py::array_t<std::uint8_t, py::array::c_style> array) {
                 dai::Asset& a = obj.cast<dai::Asset&>();
                 a.data = {array.data(), array.data() + array.size()};
+                a.path.clear();
+                a.size = a.data.size();
             })
         .def_readwrite("alignment", &Asset::alignment);
 

cmake/Depthai/DepthaiDeviceRVC4Config.cmake

@@ -3,4 +3,4 @@
 set(DEPTHAI_DEVICE_RVC4_MATURITY "snapshot")
 
 # "version if applicable"
-set(DEPTHAI_DEVICE_RVC4_VERSION "0.0.1+b52f10a00fee1fd8eb886753c8801a808f7197ba")
+set(DEPTHAI_DEVICE_RVC4_VERSION "0.0.1+914d5707cd71ef83501248b476ebfe01b59db0b1")

include/depthai/pipeline/AssetManager.hpp

@@ -17,7 +17,12 @@ struct Asset {
     explicit Asset(std::string k) : key(std::move(k)) {}
     const std::string key;
     std::vector<std::uint8_t> data;
+    std::filesystem::path path;
+    std::size_t size = 0;
+    bool dataLoaded = false;
     std::uint32_t alignment = 1;
+    std::vector<std::uint8_t>& getData();
+    std::size_t getSize() const;
     std::string getRelativeUri();
 };
 
@@ -127,6 +132,7 @@ class AssetManager /*: public Assets*/ {
 
     /// Serializes
     void serialize(AssetsMutable& assets, std::vector<std::uint8_t>& assetStorage, std::string prefix = "") const;
+    std::size_t getSerializedSize(std::size_t offset = 0) const;
 };
 
 }  // namespace dai

src/pipeline/AssetManager.cpp

@@ -6,14 +6,75 @@
 #include "utility/spdlog-fmt.hpp"
 
 // std
+#include <algorithm>
+#include <array>
 #include <fstream>
+#include <limits>
 
 namespace dai {
 
+namespace {
+
+constexpr std::size_t MAX_ASSET_STORAGE_SIZE = std::numeric_limits<std::uint32_t>::max();
+
+std::size_t getSerializedEndOffset(std::size_t offset, std::uint32_t alignment, std::size_t assetSize) {
+    if(alignment == 0) {
+        throw std::runtime_error("Asset alignment cannot be zero");
+    }
+
+    if(offset > MAX_ASSET_STORAGE_SIZE || assetSize > MAX_ASSET_STORAGE_SIZE) {
+        throw std::runtime_error("Asset storage cannot exceed 4 GiB");
+    }
+
+    std::size_t padding = 0;
+    if(alignment > 1 && offset % alignment != 0) {
+        padding = alignment - (offset % alignment);
+    }
+
+    if(padding > MAX_ASSET_STORAGE_SIZE - offset || assetSize > MAX_ASSET_STORAGE_SIZE - offset - padding) {
+        throw std::runtime_error("Asset storage cannot exceed 4 GiB");
+    }
+    return offset + padding + assetSize;
+}
+
+}  // namespace
+
 std::string Asset::getRelativeUri() {
     return fmt::format("{}:{}", "asset", key);
 }
 
+std::vector<std::uint8_t>& Asset::getData() {
+    if(data.empty() && !dataLoaded && !path.empty()) {
+        std::ifstream stream(path, std::ios::in | std::ios::binary);
+        if(!stream.is_open()) {
+            throw std::runtime_error(fmt::format("Cannot load asset, file at path {} doesn't exist.", path));
+        }
+
+        auto loadedData = std::vector<std::uint8_t>(size);
+        std::size_t loadedSize = 0;
+        while(loadedSize < size) {
+            const auto bytesToRead = std::min<std::size_t>(size - loadedSize, 1024 * 1024);
+            stream.read(reinterpret_cast<char*>(loadedData.data() + loadedSize), bytesToRead);
+            const auto bytesRead = stream.gcount();
+            if(bytesRead != static_cast<std::streamsize>(bytesToRead)) {
+                throw std::runtime_error(fmt::format("Cannot load asset, file at path {} has changed size.", path));
+            }
+            loadedSize += static_cast<std::size_t>(bytesRead);
+        }
+        if(stream.peek() != std::char_traits<char>::eof()) {
+            throw std::runtime_error(fmt::format("Cannot load asset, file at path {} has changed size.", path));
+        }
+
+        data = std::move(loadedData);
+        dataLoaded = true;
+    }
+    return data;
+}
+
+std::size_t Asset::getSize() const {
+    return path.empty() ? data.size() : size;
+}
+
 AssetManager::AssetManager() {}
 AssetManager::AssetManager(const std::string& rootPath) : rootPath{rootPath} {}
 
@@ -56,6 +117,9 @@ std::shared_ptr<dai::Asset> AssetManager::set(const std::string& key, Asset asse
     // Rename the asset with supplied key and store
     Asset a(key);
     a.data = std::move(asset.data);
+    a.path = std::move(asset.path);
+    a.size = a.path.empty() ? 0 : asset.size;
+    a.dataLoaded = asset.dataLoaded;
     a.alignment = asset.alignment;
     return set(std::move(a));
 }
@@ -72,7 +136,8 @@ std::shared_ptr<dai::Asset> AssetManager::set(const std::string& key, const std:
     // Create an asset
     Asset binaryAsset(key);
     binaryAsset.alignment = alignment;
-    binaryAsset.data = std::vector<std::uint8_t>(std::istreambuf_iterator<char>(stream), {});
+    binaryAsset.path = path;
+    binaryAsset.size = static_cast<std::size_t>(std::filesystem::file_size(path));
     // Store asset
     return set(std::move(binaryAsset));
 }
@@ -81,7 +146,7 @@ std::shared_ptr<dai::Asset> AssetManager::set(const std::string& key, const std:
     // Create an asset
     Asset binaryAsset(key);
     binaryAsset.alignment = alignment;
-    binaryAsset.data = std::move(data);
+    binaryAsset.data = data;
     // Store asset
     return set(std::move(binaryAsset));
 }
@@ -154,24 +219,63 @@ void AssetManager::serialize(AssetsMutable& mutableAssets, std::vector<std::uint
     for(auto& kv : assetMap) {
         auto& a = *kv.second;
 
-        // calculate additional bytes needed to offset to alignment
-        int toAdd = 0;
+        const auto assetSize = a.getSize();
+        const auto storageStart = storage.size();
+
+        // Calculate additional bytes needed to offset to alignment.
+        std::size_t toAdd = 0;
         if(a.alignment > 1 && storage.size() % a.alignment != 0) {
             toAdd = a.alignment - (storage.size() % a.alignment);
         }
 
+        getSerializedEndOffset(storage.size(), a.alignment, assetSize);
+
         // calculate offset
         std::uint32_t offset = static_cast<uint32_t>(storage.size()) + toAdd;
 
         // Add alignment bytes
         storage.resize(storage.size() + toAdd);
 
-        // copy data
-        storage.insert(storage.end(), a.data.begin(), a.data.end());
+        if(!a.path.empty()) {
+            try {
+                std::ifstream stream(a.path, std::ios::in | std::ios::binary);
+                if(!stream.is_open()) {
+                    throw std::runtime_error(fmt::format("Cannot load asset, file at path {} doesn't exist.", a.path));
+                }
+                std::vector<std::uint8_t> buffer(1024 * 1024);
+                std::size_t streamedSize = 0;
+                while(streamedSize < assetSize) {
+                    const auto bytesToRead = std::min(buffer.size(), assetSize - streamedSize);
+                    stream.read(reinterpret_cast<char*>(buffer.data()), bytesToRead);
+                    auto bytesRead = stream.gcount();
+                    if(bytesRead != static_cast<std::streamsize>(bytesToRead)) {
+                        throw std::runtime_error(fmt::format("Asset at path {} changed while serializing.", a.path));
+                    }
+                    storage.insert(storage.end(), buffer.data(), buffer.data() + bytesRead);
+                    streamedSize += static_cast<std::size_t>(bytesRead);
+                }
+                if(stream.peek() != std::char_traits<char>::eof()) {
+                    throw std::runtime_error(fmt::format("Asset at path {} changed while serializing.", a.path));
+                }
+            } catch(...) {
+                storage.resize(storageStart);
+                throw;
+            }
+        } else {
+            storage.insert(storage.end(), a.data.begin(), a.data.end());
+        }
 
         // Add to map the currently added asset
-        mutableAssets.set(prefix + a.key, offset, static_cast<uint32_t>(a.data.size()), a.alignment);
+        mutableAssets.set(prefix + a.key, offset, static_cast<uint32_t>(assetSize), a.alignment);
+    }
+}
+
+std::size_t AssetManager::getSerializedSize(std::size_t offset) const {
+    for(const auto& kv : assetMap) {
+        const auto& a = *kv.second;
+        offset = getSerializedEndOffset(offset, a.alignment, a.getSize());
     }
+    return offset;
 }
 
 void AssetsMutable::set(const std::string& key, std::uint32_t offset, std::uint32_t size, std::uint32_t alignment) {

src/pipeline/Pipeline.cpp

@@ -299,6 +299,11 @@ void PipelineImpl::serialize(PipelineSchema& schema, Assets& assets, std::vector
 
     // Serialize all asset managers into asset storage
     assetStorage.clear();
+    std::size_t storageSize = assetManager.getSerializedSize();
+    for(auto& node : getAllNodes()) {
+        storageSize = node->getAssetManager().getSerializedSize(storageSize);
+    }
+    assetStorage.reserve(storageSize);
     AssetsMutable mutableAssets;
     // Pipeline assets
     assetManager.serialize(mutableAssets, assetStorage, "/pipeline/");
@@ -588,10 +593,10 @@ void PipelineImpl::setCameraTuningBlobPath(const fs::path& path) {
     auto asset = assetManager.set(assetKey, path);
 
     if(defaultDevice) {
-        defaultDevice->setCameraTuningBlob(asset->getRelativeUri(), static_cast<uint32_t>(asset->data.size()));
+        defaultDevice->setCameraTuningBlob(asset->getRelativeUri(), static_cast<uint32_t>(asset->getSize()));
     } else if(defaultDeviceProperties != nullptr) {
         defaultDeviceProperties->cameraTuningBlobUri = asset->getRelativeUri();
-        defaultDeviceProperties->cameraTuningBlobSize = static_cast<uint32_t>(asset->data.size());
+        defaultDeviceProperties->cameraTuningBlobSize = static_cast<uint32_t>(asset->getSize());
     }
 }
 
@@ -602,10 +607,10 @@ void PipelineImpl::setCameraTuningBlobPath(CameraBoardSocket socket, const fs::p
     auto asset = assetManager.set(assetKey, path);
 
     if(defaultDevice) {
-        defaultDevice->setCameraSocketTuningBlob(socket, asset->getRelativeUri(), static_cast<uint32_t>(asset->data.size()));
+        defaultDevice->setCameraSocketTuningBlob(socket, asset->getRelativeUri(), static_cast<uint32_t>(asset->getSize()));
     } else if(defaultDeviceProperties != nullptr) {
         defaultDeviceProperties->cameraSocketTuningBlobUri[socket] = asset->getRelativeUri();
-        defaultDeviceProperties->cameraSocketTuningBlobSize[socket] = static_cast<uint32_t>(asset->data.size());
+        defaultDeviceProperties->cameraSocketTuningBlobSize[socket] = static_cast<uint32_t>(asset->getSize());
     }
 }
 
@@ -1389,19 +1394,19 @@ std::vector<uint8_t> PipelineImpl::loadResourceCwd(fs::path uri, fs::path cwd, b
              if(asset != nullptr) {
                  if(moveAsset) {
                      p.assetManager.remove(uriString);
-                     return std::move(asset->data);
+                     return std::move(asset->getData());
                  }
-                 return asset->data;
+                 return asset->getData();
              }
              for(auto& node : p.nodes) {
                  auto& assetManager = node->getAssetManager();
                  auto asset = assetManager.get(uriString);
                  if(asset != nullptr) {
                      if(moveAsset) {
                          assetManager.remove(uriString);
-                         return std::move(asset->data);
+                         return std::move(asset->getData());
                      }
-                     return asset->data;
+                     return asset->getData();
                  }
              }
              // Asset not found anywhere

src/pipeline/node/NeuralNetwork.cpp

@@ -229,7 +229,7 @@ void NeuralNetwork::setBlob(OpenVINO::Blob blob) {
     }
     auto asset = assetManager.set("__blob", std::move(blob.data));
     properties.blobUri = asset->getRelativeUri();
-    properties.blobSize = static_cast<uint32_t>(asset->data.size());
+    properties.blobSize = static_cast<uint32_t>(asset->getSize());
     properties.modelSource = Properties::ModelSource::BLOB;
 }
 

tests/src/ondevice_tests/pipeline/node/neural_network_node_test.cpp

@@ -1,7 +1,11 @@
 #include <catch2/catch_all.hpp>
 #include <catch2/catch_test_macros.hpp>
+#include <chrono>
+#include <filesystem>
+#include <fstream>
 #include <magic_enum/magic_enum.hpp>
 #include <opencv2/videoio.hpp>
+#include <thread>
 
 #include "depthai/common/CameraBoardSocket.hpp"
 #include "depthai/depthai.hpp"
@@ -50,6 +54,67 @@ TEST_CASE("NNArchive API") {
     }
 }
 
+TEST_CASE("RVC4 NeuralNetwork model loading paths", "[rvc4]") {
+    std::vector<dai::DeviceModelZoo> supportedDeviceModels;
+    {
+        dai::Pipeline discoveryPipeline;
+        const auto device = discoveryPipeline.getDefaultDevice();
+        if(device->getPlatform() != dai::Platform::RVC4) {
+            SKIP("RVC4-only test");
+        }
+        supportedDeviceModels = device->getSupportedDeviceModels();
+    }
+    REQUIRE_FALSE(supportedDeviceModels.empty());
+
+    const dai::NNModelDescription description{"yolov6-nano", "RVC4"};
+    const auto archivePath = dai::getModelFromZoo(description);
+    const dai::NNArchive archive{archivePath};
+    const auto modelData = archive.getOtherModelFormat();
+    REQUIRE(modelData.has_value());
+
+    const std::filesystem::path directModelPath =
+        std::filesystem::temp_directory_path()
+        / ("depthai-rvc4-model-loading-test_" + std::to_string(std::chrono::steady_clock::now().time_since_epoch().count()) + ".dlc");
+    {
+        std::ofstream modelFile(directModelPath, std::ios::binary | std::ios::trunc);
+        REQUIRE(modelFile.is_open());
+        modelFile.write(reinterpret_cast<const char*>(modelData->data()), static_cast<std::streamsize>(modelData->size()));
+        REQUIRE(modelFile.good());
+    }
+
+    const auto startPipeline = [](const std::string& path, const auto& configure) {
+        INFO(path);
+        dai::Pipeline pipeline;
+        auto neuralNetwork = pipeline.create<dai::node::NeuralNetwork>();
+        configure(neuralNetwork);
+
+        // The queue helpers provide the required single input connection while
+        // keeping this focused on model initialization rather than inference.
+        auto inputQueue = neuralNetwork->input.createInputQueue();
+        auto outputQueue = neuralNetwork->out.createOutputQueue();
+        pipeline.start();
+        REQUIRE(pipeline.isRunning());
+        pipeline.stop();
+        pipeline.wait();
+        std::this_thread::sleep_for(std::chrono::milliseconds(100));
+    };
+
+    startPipeline("setModelPath(.dlc)", [&directModelPath](const auto& neuralNetwork) { neuralNetwork->setModelPath(directModelPath); });
+    startPipeline("setOtherModelFormat(.dlc)", [&directModelPath](const auto& neuralNetwork) { neuralNetwork->setOtherModelFormat(directModelPath); });
+    startPipeline("setOtherModelFormat(vector)", [&modelData](const auto& neuralNetwork) { neuralNetwork->setOtherModelFormat(*modelData); });
+    startPipeline("setModelPath(NNArchive)", [&archivePath](const auto& neuralNetwork) { neuralNetwork->setModelPath(archivePath); });
+    startPipeline("setNNArchive(NNArchive)", [&archive](const auto& neuralNetwork) { neuralNetwork->setNNArchive(archive); });
+    startPipeline("setFromModelZoo", [&description](const auto& neuralNetwork) { neuralNetwork->setFromModelZoo(description, true); });
+
+    for(const auto model : supportedDeviceModels) {
+        INFO(magic_enum::enum_name(model));
+        startPipeline("setModelFromDeviceZoo", [model](const auto& neuralNetwork) { neuralNetwork->setModelFromDeviceZoo(model); });
+    }
+
+    std::error_code ec;
+    std::filesystem::remove(directModelPath, ec);
+}
+
 TEST_CASE("Multi-Input NeuralNetwork API") {
     dai::Pipeline p;
     auto camera = p.create<dai::node::Camera>()->build(dai::CameraBoardSocket::CAM_A);