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,13 @@ 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.setFile({}, 0);
             })
         .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+91a462a582fcb01dbd7e95b79cdea101beed5780")

include/depthai/pipeline/AssetManager.hpp

@@ -18,7 +18,19 @@ struct Asset {
     const std::string key;
     std::vector<std::uint8_t> data;
     std::uint32_t alignment = 1;
+    std::vector<std::uint8_t>& getData();
+    std::size_t getSize() const;
     std::string getRelativeUri();
+
+    /// Set the backing file and its expected size.
+    void setFile(std::filesystem::path path, std::size_t size);
+
+   private:
+    friend class AssetManager;
+
+    std::filesystem::path path;
+    std::size_t size = 0;
+    bool dataLoaded = false;
 };
 
 class AssetsMutable : public Assets {
@@ -127,6 +139,8 @@ class AssetManager /*: public Assets*/ {
 
     /// Serializes
     void serialize(AssetsMutable& assets, std::vector<std::uint8_t>& assetStorage, std::string prefix = "") const;
+    /// Calculates the size of the serialized data
+    std::size_t getSerializedSize(std::size_t offset = 0) const;
 };
 
 }  // namespace dai

src/pipeline/AssetManager.cpp

@@ -6,14 +6,81 @@
 #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(!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;
+}
+
+void Asset::setFile(std::filesystem::path path, std::size_t size) {
+    this->path = std::move(path);
+    this->size = size;
+    dataLoaded = false;
+}
+
 AssetManager::AssetManager() {}
 AssetManager::AssetManager(const std::string& rootPath) : rootPath{rootPath} {}
 
@@ -56,6 +123,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);
+    const auto assetSize = asset.path.empty() ? 0 : asset.size;
+    a.setFile(std::move(asset.path), assetSize);
+    a.dataLoaded = asset.dataLoaded;
     a.alignment = asset.alignment;
     return set(std::move(a));
 }
@@ -72,7 +142,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::vector<std::uint8_t>(std::istreambuf_iterator<char>(stream), {});
+    binaryAsset.setFile(path, static_cast<std::size_t>(std::filesystem::file_size(path)));
     // Store asset
     return set(std::move(binaryAsset));
 }
@@ -81,7 +151,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));
 }
@@ -151,27 +221,75 @@ void AssetManager::serialize(AssetsMutable& mutableAssets, std::vector<std::uint
         prefix = rootPath;
     }
 
-    for(auto& kv : assetMap) {
-        auto& a = *kv.second;
-
-        // calculate additional bytes needed to offset to alignment
-        int toAdd = 0;
-        if(a.alignment > 1 && storage.size() % a.alignment != 0) {
-            toAdd = a.alignment - (storage.size() % a.alignment);
+    const auto storageStart = storage.size();
+    const auto mutableAssetsStart = mutableAssets;
+    try {
+        storage.reserve(getSerializedSize(storageStart));
+        for(auto& kv : assetMap) {
+            auto& a = *kv.second;
+
+            const auto assetSize = a.getSize();
+            const auto assetStorageStart = 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);
+
+            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(assetStorageStart);
+                    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>(assetSize), a.alignment);
         }
+    } catch(...) {
+        storage.resize(storageStart);
+        mutableAssets = mutableAssetsStart;
+        throw;
+    }
+}
 
-        // 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());
-
-        // Add to map the currently added asset
-        mutableAssets.set(prefix + a.key, offset, static_cast<uint32_t>(a.data.size()), 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;
 }