mirror of
https://github.com/opencv/opencv.git
synced 2024-11-29 05:29:54 +08:00
Add support for loading ONNX model from in-memory buffer.
This commit is contained in:
parent
47c45e5bd3
commit
7b4aefedea
@ -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();
|
||||
|
Loading…
Reference in New Issue
Block a user