diff --git a/bindings/python/src/nn_archive/NNArchiveBindings.cpp b/bindings/python/src/nn_archive/NNArchiveBindings.cpp index a3620a782b..e87847cad3 100644 --- a/bindings/python/src/nn_archive/NNArchiveBindings.cpp +++ b/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(model->data()), model->size()); + }, + DOC(dai, NNArchive, getOtherModelFormat)); 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)); diff --git a/bindings/python/src/pipeline/AssetManagerBindings.cpp b/bindings/python/src/pipeline/AssetManagerBindings.cpp index cc2367aae2..a7d2f00b87 100644 --- a/bindings/python/src/pipeline/AssetManagerBindings.cpp +++ b/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(); - return py::array_t(a.data.size(), a.data.data(), obj); + auto& data = a.getData(); + return py::array_t(data.size(), data.data(), obj); }, [](py::object& obj, py::array_t array) { dai::Asset& a = obj.cast(); a.data = {array.data(), array.data() + array.size()}; + a.setFile({}, 0); }) .def_readwrite("alignment", &Asset::alignment); diff --git a/cmake/Depthai/DepthaiDeviceRVC4Config.cmake b/cmake/Depthai/DepthaiDeviceRVC4Config.cmake index 212636025d..96311596e4 100644 --- a/cmake/Depthai/DepthaiDeviceRVC4Config.cmake +++ b/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+48c2e4a32587a266f11e35d008768a95b72712df") +set(DEPTHAI_DEVICE_RVC4_VERSION "0.0.1+842f297705fdfd63273779bef2e2965c312e55bc") diff --git a/include/depthai/pipeline/AssetManager.hpp b/include/depthai/pipeline/AssetManager.hpp index ccacb67894..bb7099b3b5 100644 --- a/include/depthai/pipeline/AssetManager.hpp +++ b/include/depthai/pipeline/AssetManager.hpp @@ -18,7 +18,19 @@ struct Asset { const std::string key; std::vector data; std::uint32_t alignment = 1; + std::vector& 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& assetStorage, std::string prefix = "") const; + /// Calculates the size of the serialized data + std::size_t getSerializedSize(std::size_t offset = 0) const; }; } // namespace dai diff --git a/src/pipeline/AssetManager.cpp b/src/pipeline/AssetManager.cpp index 8197acf336..b33da50601 100644 --- a/src/pipeline/AssetManager.cpp +++ b/src/pipeline/AssetManager.cpp @@ -6,14 +6,81 @@ #include "utility/spdlog-fmt.hpp" // std +#include +#include #include +#include namespace dai { +namespace { + +constexpr std::size_t MAX_ASSET_STORAGE_SIZE = std::numeric_limits::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& 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(size); + std::size_t loadedSize = 0; + while(loadedSize < size) { + const auto bytesToRead = std::min(size - loadedSize, 1024 * 1024); + stream.read(reinterpret_cast(loadedData.data() + loadedSize), bytesToRead); + const auto bytesRead = stream.gcount(); + if(bytesRead != static_cast(bytesToRead)) { + throw std::runtime_error(fmt::format("Cannot load asset, file at path {} has changed size.", path)); + } + loadedSize += static_cast(bytesRead); + } + if(stream.peek() != std::char_traits::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 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 AssetManager::set(const std::string& key, const std: // Create an asset Asset binaryAsset(key); binaryAsset.alignment = alignment; - binaryAsset.data = std::vector(std::istreambuf_iterator(stream), {}); + binaryAsset.setFile(path, static_cast(std::filesystem::file_size(path))); // Store asset return set(std::move(binaryAsset)); } @@ -81,7 +151,7 @@ std::shared_ptr 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 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(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 buffer(1024 * 1024); + std::size_t streamedSize = 0; + while(streamedSize < assetSize) { + const auto bytesToRead = std::min(buffer.size(), assetSize - streamedSize); + stream.read(reinterpret_cast(buffer.data()), bytesToRead); + auto bytesRead = stream.gcount(); + if(bytesRead != static_cast(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(bytesRead); + } + if(stream.peek() != std::char_traits::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(assetSize), a.alignment); } + } catch(...) { + storage.resize(storageStart); + mutableAssets = mutableAssetsStart; + throw; + } +} - // calculate offset - std::uint32_t offset = static_cast(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(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) { diff --git a/src/pipeline/Pipeline.cpp b/src/pipeline/Pipeline.cpp index 3fe26e6309..7be94abf4b 100644 --- a/src/pipeline/Pipeline.cpp +++ b/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(asset->data.size())); + defaultDevice->setCameraTuningBlob(asset->getRelativeUri(), static_cast(asset->getSize())); } else if(defaultDeviceProperties != nullptr) { defaultDeviceProperties->cameraTuningBlobUri = asset->getRelativeUri(); - defaultDeviceProperties->cameraTuningBlobSize = static_cast(asset->data.size()); + defaultDeviceProperties->cameraTuningBlobSize = static_cast(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(asset->data.size())); + defaultDevice->setCameraSocketTuningBlob(socket, asset->getRelativeUri(), static_cast(asset->getSize())); } else if(defaultDeviceProperties != nullptr) { defaultDeviceProperties->cameraSocketTuningBlobUri[socket] = asset->getRelativeUri(); - defaultDeviceProperties->cameraSocketTuningBlobSize[socket] = static_cast(asset->data.size()); + defaultDeviceProperties->cameraSocketTuningBlobSize[socket] = static_cast(asset->getSize()); } } @@ -1389,9 +1394,9 @@ std::vector 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(); @@ -1399,9 +1404,9 @@ std::vector PipelineImpl::loadResourceCwd(fs::path uri, fs::path cwd, b 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 diff --git a/src/pipeline/node/NeuralNetwork.cpp b/src/pipeline/node/NeuralNetwork.cpp index 4403979fcb..cb5a473004 100644 --- a/src/pipeline/node/NeuralNetwork.cpp +++ b/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(asset->data.size()); + properties.blobSize = static_cast(asset->getSize()); properties.modelSource = Properties::ModelSource::BLOB; } diff --git a/tests/src/ondevice_tests/pipeline/node/neural_network_node_test.cpp b/tests/src/ondevice_tests/pipeline/node/neural_network_node_test.cpp index c723f408fe..447f473dac 100644 --- a/tests/src/ondevice_tests/pipeline/node/neural_network_node_test.cpp +++ b/tests/src/ondevice_tests/pipeline/node/neural_network_node_test.cpp @@ -1,7 +1,11 @@ #include #include +#include +#include +#include #include #include +#include #include "depthai/common/CameraBoardSocket.hpp" #include "depthai/depthai.hpp" @@ -50,6 +54,83 @@ TEST_CASE("NNArchive API") { } } +TEST_CASE("RVC4 NeuralNetwork model loading paths", "[rvc4]") { + std::vector 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 auto inputSize = archive.getInputSize(); + REQUIRE(inputSize.has_value()); + auto inputType = dai::ImgFrame::Type::BGR888i; + const auto modelInputType = archive.getConfig().model.inputs[0].preprocessing.daiType; + if(modelInputType.has_value()) { + const auto convertedInputType = magic_enum::enum_cast(*modelInputType); + REQUIRE(convertedInputType.has_value()); + inputType = *convertedInputType; + } + + auto inputFrame = std::make_shared(); + cv::Mat frame(inputSize->second, inputSize->first, CV_8UC3, cv::Scalar(0, 255, 0)); + inputFrame->setCvFrame(frame, inputType); + + 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(modelData->data()), static_cast(modelData->size())); + REQUIRE(modelFile.good()); + } + + const auto startPipeline = [&inputFrame](const std::string& path, const auto& configure, bool runInference = false) { + INFO(path); + dai::Pipeline pipeline; + auto neuralNetwork = pipeline.create(); + configure(neuralNetwork); + + auto inputQueue = neuralNetwork->input.createInputQueue(); + auto outputQueue = neuralNetwork->out.createOutputQueue(); + pipeline.start(); + REQUIRE(pipeline.isRunning()); + if(runInference) { + inputQueue->send(inputFrame); + REQUIRE(outputQueue->get() != nullptr); + } + pipeline.stop(); + pipeline.wait(); + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + }; + + startPipeline("setModelPath(.dlc)", [&directModelPath](const auto& neuralNetwork) { neuralNetwork->setModelPath(directModelPath); }, true); + startPipeline("setOtherModelFormat(.dlc)", [&directModelPath](const auto& neuralNetwork) { neuralNetwork->setOtherModelFormat(directModelPath); }, true); + startPipeline("setOtherModelFormat(vector)", [&modelData](const auto& neuralNetwork) { neuralNetwork->setOtherModelFormat(*modelData); }, true); + startPipeline("setModelPath(NNArchive)", [&archivePath](const auto& neuralNetwork) { neuralNetwork->setModelPath(archivePath); }, true); + startPipeline("setNNArchive(NNArchive)", [&archive](const auto& neuralNetwork) { neuralNetwork->setNNArchive(archive); }, true); + startPipeline("setFromModelZoo", [&description](const auto& neuralNetwork) { neuralNetwork->setFromModelZoo(description, true); }, 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()->build(dai::CameraBoardSocket::CAM_A); diff --git a/tests/src/onhost_tests/serialization_test.cpp b/tests/src/onhost_tests/serialization_test.cpp index fcf3921de0..b19d95a5c2 100644 --- a/tests/src/onhost_tests/serialization_test.cpp +++ b/tests/src/onhost_tests/serialization_test.cpp @@ -1,4 +1,7 @@ #include +#include +#include +#include // Include depthai library #include @@ -47,3 +50,81 @@ TEST_CASE("Roundtrip") { REQUIRE(des.numFramesPool == 42); } } + +TEST_CASE("AssetManager uses the current size of memory-backed assets") { + dai::AssetManager assetManager; + auto asset = assetManager.set("asset", std::vector{1, 2}); + asset->data.push_back(3); + + REQUIRE(asset->getSize() == 3); + REQUIRE(assetManager.getSerializedSize() == 3); +} + +TEST_CASE("AssetManager rejects storage beyond 4 GiB") { + dai::Asset asset("oversized"); + asset.setFile("placeholder", static_cast(std::numeric_limits::max()) + 1); + + dai::AssetManager assetManager; + assetManager.set(std::move(asset)); + + REQUIRE_THROWS_WITH(assetManager.getSerializedSize(), "Asset storage cannot exceed 4 GiB"); +} + +TEST_CASE("AssetManager preserves the size of path-backed assets when renaming them") { + dai::Asset asset("source"); + asset.setFile("placeholder", 42); + + dai::AssetManager assetManager; + auto renamedAsset = assetManager.set("renamed", std::move(asset)); + + REQUIRE(renamedAsset->getSize() == 42); +} + +TEST_CASE("AssetManager restores storage when a path-backed asset changes") { + const auto path = std::filesystem::temp_directory_path() / "depthai_asset_manager_serialization_test.bin"; + { + std::ofstream stream(path, std::ios::binary); + stream.write("ab", 2); + } + + dai::AssetManager assetManager; + assetManager.set("asset", path); + { + std::ofstream stream(path, std::ios::binary | std::ios::app); + stream.write("c", 1); + } + + dai::AssetsMutable assets; + std::vector storage{42}; + REQUIRE_THROWS(assetManager.serialize(assets, storage)); + REQUIRE(storage == std::vector{42}); + + std::filesystem::remove(path); +} + +TEST_CASE("AssetManager rolls back all assets when a later path-backed asset changes") { + const auto path = std::filesystem::temp_directory_path() / "depthai_asset_manager_transaction_test.bin"; + { + std::ofstream stream(path, std::ios::binary); + stream.write("ab", 2); + } + + dai::AssetManager assetManager; + assetManager.set("first", std::vector{1, 2}); + assetManager.set("second", path); + { + std::ofstream stream(path, std::ios::binary | std::ios::app); + stream.write("c", 1); + } + + dai::AssetsMutable assets; + assets.set("existing", 0, 1, 1); + std::vector storage{42}; + REQUIRE_THROWS(assetManager.serialize(assets, storage)); + REQUIRE(storage == std::vector{42}); + REQUIRE(assets.has("existing")); + REQUIRE_FALSE(assets.has("first")); + REQUIRE_FALSE(assets.has("second")); + + std::filesystem::remove(path); +}