提交 7b4aefed 编写于 作者: K Kohei Yoshida

Add support for loading ONNX model from in-memory buffer.

上级 47c45e5b
......@@ -868,6 +868,23 @@ CV__DNN_EXPERIMENTAL_NS_BEGIN
*/
CV_EXPORTS_W Net readNetFromONNX(const String &onnxFile);
/** @brief Reads a network model from <a href="https://onnx.ai/">ONNX</a>
* in-memory buffer.
* @param buffer memory address of the first byte of the buffer.
* @param sizeBuffer size of the buffer.
* @returns Network object that ready to do forward, throw an exception
* in failure cases.
*/
CV_EXPORTS Net readNetFromONNX(const char* buffer, size_t sizeBuffer);
/** @brief Reads a network model from <a href="https://onnx.ai/">ONNX</a>
* in-memory buffer.
* @param buffer in-memory buffer that stores the ONNX model bytes.
* @returns Network object that ready to do forward, throw an exception
* in failure cases.
*/
CV_EXPORTS_W Net readNetFromONNX(const std::vector<uchar>& buffer);
/** @brief Creates blob from .pb file.
* @param path to the .pb file with input tensor.
* @returns Mat.
......
......@@ -57,6 +57,24 @@ public:
CV_Error(Error::StsUnsupportedFormat, "Failed to parse onnx model");
}
ONNXImporter(const char* buffer, size_t sizeBuffer)
{
struct _Buf : public std::streambuf
{
_Buf(const char* buffer, size_t sizeBuffer)
{
char* p = const_cast<char*>(buffer);
setg(p, p, p + sizeBuffer);
}
};
_Buf buf(buffer, sizeBuffer);
std::istream input(&buf);
if (!model_proto.ParseFromIstream(&input))
CV_Error(Error::StsUnsupportedFormat, "Failed to parse onnx model from in-memory byte array.");
}
void populateNet(Net dstNet);
};
......@@ -808,6 +826,19 @@ Net readNetFromONNX(const String& onnxFile)
return net;
}
Net readNetFromONNX(const char* buffer, size_t sizeBuffer)
{
ONNXImporter onnxImporter(buffer, sizeBuffer);
Net net;
onnxImporter.populateNet(net);
return net;
}
Net readNetFromONNX(const std::vector<uchar>& buffer)
{
return readNetFromONNX(reinterpret_cast<const char*>(buffer.data()), buffer.size());
}
Mat readTensorFromONNX(const String& path)
{
opencv_onnx::TensorProto tensor_proto = opencv_onnx::TensorProto();
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册