Skip to content

Commit f723ced

Browse files
r2d3r2d3
committed
add loading TensorFlow/Caffe net from memory buffer
add a corresponding test
1 parent 6e4f943 commit f723ced

File tree

10 files changed

+217
-43
lines changed

10 files changed

+217
-43
lines changed

‎modules/dnn/include/opencv2/dnn/dnn.hpp

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -634,11 +634,33 @@ CV__DNN_EXPERIMENTAL_NS_BEGIN
634634
*/
635635
CV_EXPORTS_W Net readNetFromCaffe(const String &prototxt, const String &caffeModel = String());
636636

637+
/** @brief Reads a network model stored in Caffe model in memory.
638+
* @details This is an overloaded member function, provided for convenience.
639+
* It differs from the above function only in what argument(s) it accepts.
640+
* @param bufferProto buffer containing the content of the .prototxt file
641+
* @param lenProto length of bufferProto
642+
* @param bufferModel buffer containing the content of the .caffemodel file
643+
* @param lenModel length of bufferModel
644+
*/
645+
CV_EXPORTS Net readNetFromCaffe(const char *bufferProto, size_t lenProto,
646+
const char *bufferModel = NULL, size_t lenModel = 0);
647+
637648
/** @brief Reads a network model stored in Tensorflow model file.
638649
* @details This is shortcut consisting from createTensorflowImporter and Net::populateNet calls.
639650
*/
640651
CV_EXPORTS_W Net readNetFromTensorflow(const String &model, const String &config = String());
641652

653+
/** @brief Reads a network model stored in Tensorflow model in memory.
654+
* @details This is an overloaded member function, provided for convenience.
655+
* It differs from the above function only in what argument(s) it accepts.
656+
* @param bufferModel buffer containing the content of the pb file
657+
* @param lenModel length of bufferModel
658+
* @param bufferConfig buffer containing the content of the pbtxt file
659+
* @param lenConfig length of bufferConfig
660+
*/
661+
CV_EXPORTS Net readNetFromTensorflow(const char *bufferModel, size_t lenModel,
662+
const char *bufferConfig = NULL, size_t lenConfig = 0);
663+
642664
/** @brief Reads a network model stored in Torch model file.
643665
* @details This is shortcut consisting from createTorchImporter and Net::populateNet calls.
644666
*/

‎modules/dnn/src/caffe/caffe_importer.cpp

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,17 @@ class CaffeImporter : public Importer
9494
ReadNetParamsFromBinaryFileOrDie(caffeModel, &netBinary);
9595
}
9696

97+
CaffeImporter(const char *dataProto, size_t lenProto,
98+
const char *dataModel, size_t lenModel)
99+
{
100+
CV_TRACE_FUNCTION();
101+
102+
ReadNetParamsFromTextBufferOrDie(dataProto, lenProto, &net);
103+
104+
if (dataModel != NULL && lenModel > 0)
105+
ReadNetParamsFromBinaryBufferOrDie(dataModel, lenModel, &netBinary);
106+
}
107+
97108
void addParam(const Message &msg, const FieldDescriptor *field, cv::dnn::LayerParams &params)
98109
{
99110
const Reflection *refl = msg.GetReflection();
@@ -400,6 +411,15 @@ Net readNetFromCaffe(const String &prototxt, const String &caffeModel /*= String
400411
return net;
401412
}
402413

414+
Net readNetFromCaffe(const char *bufferProto, size_t lenProto,
415+
const char *bufferModel, size_t lenModel)
416+
{
417+
CaffeImporter caffeImporter(bufferProto, lenProto, bufferModel, lenModel);
418+
Net net;
419+
caffeImporter.populateNet(net);
420+
return net;
421+
}
422+
403423
#endif //HAVE_PROTOBUF
404424

405425
CV__DNN_EXPERIMENTAL_NS_END

‎modules/dnn/src/caffe/caffe_io.cpp

Lines changed: 34 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1108,28 +1108,37 @@ const char* UpgradeV1LayerType(const V1LayerParameter_LayerType type) {
11081108

11091109
const int kProtoReadBytesLimit = INT_MAX; // Max size of 2 GB minus 1 byte.
11101110

1111+
bool ReadProtoFromBinary(ZeroCopyInputStream* input, Message *proto) {
1112+
CodedInputStream coded_input(input);
1113+
coded_input.SetTotalBytesLimit(kProtoReadBytesLimit, 536870912);
1114+
1115+
return proto->ParseFromCodedStream(&coded_input);
1116+
}
1117+
11111118
bool ReadProtoFromTextFile(const char* filename, Message* proto) {
11121119
std::ifstream fs(filename, std::ifstream::in);
11131120
CHECK(fs.is_open()) << "Can't open \"" << filename << "\"";
11141121
IstreamInputStream input(&fs);
1115-
bool success = google::protobuf::TextFormat::Parse(&input, proto);
1116-
fs.close();
1117-
return success;
1122+
return google::protobuf::TextFormat::Parse(&input, proto);
11181123
}
11191124

11201125
bool ReadProtoFromBinaryFile(const char* filename, Message* proto) {
11211126
std::ifstream fs(filename, std::ifstream::in | std::ifstream::binary);
11221127
CHECK(fs.is_open()) << "Can't open \"" << filename << "\"";
1123-
ZeroCopyInputStream* raw_input = new IstreamInputStream(&fs);
1124-
CodedInputStream* coded_input = new CodedInputStream(raw_input);
1125-
coded_input->SetTotalBytesLimit(kProtoReadBytesLimit, 536870912);
1128+
IstreamInputStream raw_input(&fs);
1129+
1130+
return ReadProtoFromBinary(&raw_input, proto);
1131+
}
1132+
1133+
bool ReadProtoFromTextBuffer(const char* data, size_t len, Message* proto) {
1134+
ArrayInputStream input(data, len);
1135+
return google::protobuf::TextFormat::Parse(&input, proto);
1136+
}
11261137

1127-
bool success = proto->ParseFromCodedStream(coded_input);
11281138

1129-
delete coded_input;
1130-
delete raw_input;
1131-
fs.close();
1132-
return success;
1139+
bool ReadProtoFromBinaryBuffer(const char* data, size_t len, Message* proto) {
1140+
ArrayInputStream raw_input(data, len);
1141+
return ReadProtoFromBinary(&raw_input, proto);
11331142
}
11341143

11351144
void ReadNetParamsFromTextFileOrDie(const char* param_file,
@@ -1139,13 +1148,27 @@ void ReadNetParamsFromTextFileOrDie(const char* param_file,
11391148
UpgradeNetAsNeeded(param_file, param);
11401149
}
11411150

1151+
void ReadNetParamsFromTextBufferOrDie(const char* data, size_t len,
1152+
NetParameter* param) {
1153+
CHECK(ReadProtoFromTextBuffer(data, len, param))
1154+
<< "Failed to parse NetParameter buffer";
1155+
UpgradeNetAsNeeded("memory buffer", param);
1156+
}
1157+
11421158
void ReadNetParamsFromBinaryFileOrDie(const char* param_file,
11431159
NetParameter* param) {
11441160
CHECK(ReadProtoFromBinaryFile(param_file, param))
11451161
<< "Failed to parse NetParameter file: " << param_file;
11461162
UpgradeNetAsNeeded(param_file, param);
11471163
}
11481164

1165+
void ReadNetParamsFromBinaryBufferOrDie(const char* data, size_t len,
1166+
NetParameter* param) {
1167+
CHECK(ReadProtoFromBinaryBuffer(data, len, param))
1168+
<< "Failed to parse NetParameter buffer";
1169+
UpgradeNetAsNeeded("memory buffer", param);
1170+
}
1171+
11491172
}
11501173
}
11511174
#endif

‎modules/dnn/src/caffe/caffe_io.hpp

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,18 @@ void ReadNetParamsFromTextFileOrDie(const char* param_file,
102102
void ReadNetParamsFromBinaryFileOrDie(const char* param_file,
103103
caffe::NetParameter* param);
104104

105+
// Read parameters from a memory buffer into a NetParammeter proto message.
106+
void ReadNetParamsFromBinaryBufferOrDie(const char* data, size_t len,
107+
caffe::NetParameter* param);
108+
void ReadNetParamsFromTextBufferOrDie(const char* data, size_t len,
109+
caffe::NetParameter* param);
110+
111+
// Utility functions used internally by Caffe and TensorFlow loaders
112+
bool ReadProtoFromTextFile(const char* filename, ::google::protobuf::Message* proto);
113+
bool ReadProtoFromBinaryFile(const char* filename, ::google::protobuf::Message* proto);
114+
bool ReadProtoFromTextBuffer(const char* data, size_t len, ::google::protobuf::Message* proto);
115+
bool ReadProtoFromBinaryBuffer(const char* data, size_t len, ::google::protobuf::Message* proto);
116+
105117
}
106118
}
107119
#endif

‎modules/dnn/src/tensorflow/tf_importer.cpp

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -449,6 +449,9 @@ void ExcludeLayer(tensorflow::GraphDef& net, const int layer_index, const int in
449449
class TFImporter : public Importer {
450450
public:
451451
TFImporter(const char *model, const char *config = NULL);
452+
TFImporter(const char *dataModel, size_t lenModel,
453+
const char *dataConfig = NULL, size_t lenConfig = 0);
454+
452455
void populateNet(Net dstNet);
453456
~TFImporter() {}
454457

@@ -479,6 +482,15 @@ TFImporter::TFImporter(const char *model, const char *config)
479482
ReadTFNetParamsFromTextFileOrDie(config, &netTxt);
480483
}
481484

485+
TFImporter::TFImporter(const char *dataModel, size_t lenModel,
486+
const char *dataConfig, size_t lenConfig)
487+
{
488+
if (dataModel != NULL && lenModel > 0)
489+
ReadTFNetParamsFromBinaryBufferOrDie(dataModel, lenModel, &netBin);
490+
if (dataConfig != NULL && lenConfig > 0)
491+
ReadTFNetParamsFromTextBufferOrDie(dataConfig, lenConfig, &netTxt);
492+
}
493+
482494
void TFImporter::kernelFromTensor(const tensorflow::TensorProto &tensor, Mat &dstBlob)
483495
{
484496
MatShape shape;
@@ -1326,5 +1338,14 @@ Net readNetFromTensorflow(const String &model, const String &config)
13261338
return net;
13271339
}
13281340

1341+
Net readNetFromTensorflow(const char* bufferModel, size_t lenModel,
1342+
const char* bufferConfig, size_t lenConfig)
1343+
{
1344+
TFImporter importer(bufferModel, lenModel, bufferConfig, lenConfig);
1345+
Net net;
1346+
importer.populateNet(net);
1347+
return net;
1348+
}
1349+
13291350
CV__DNN_EXPERIMENTAL_NS_END
13301351
}} // namespace

‎modules/dnn/src/tensorflow/tf_io.cpp

Lines changed: 16 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ Implementation of various functions which are related to Tensorflow models readi
2323

2424
#include "graph.pb.h"
2525
#include "tf_io.hpp"
26+
#include "../caffe/caffe_io.hpp"
2627
#include "../caffe/glog_emulator.hpp"
2728

2829
namespace cv {
@@ -36,41 +37,28 @@ using namespace ::google::protobuf::io;
3637

3738
const int kProtoReadBytesLimit = INT_MAX; // Max size of 2 GB minus 1 byte.
3839

39-
// TODO: remove Caffe duplicate
40-
bool ReadProtoFromBinaryFileTF(const char* filename, Message* proto) {
41-
std::ifstream fs(filename, std::ifstream::in | std::ifstream::binary);
42-
CHECK(fs.is_open()) << "Can't open \"" << filename << "\"";
43-
ZeroCopyInputStream* raw_input = new IstreamInputStream(&fs);
44-
CodedInputStream* coded_input = new CodedInputStream(raw_input);
45-
coded_input->SetTotalBytesLimit(kProtoReadBytesLimit, 536870912);
46-
47-
bool success = proto->ParseFromCodedStream(coded_input);
48-
49-
delete coded_input;
50-
delete raw_input;
51-
fs.close();
52-
return success;
40+
void ReadTFNetParamsFromBinaryFileOrDie(const char* param_file,
41+
tensorflow::GraphDef* param) {
42+
CHECK(ReadProtoFromBinaryFile(param_file, param))
43+
<< "Failed to parse GraphDef file: " << param_file;
5344
}
5445

55-
bool ReadProtoFromTextFileTF(const char* filename, Message* proto) {
56-
std::ifstream fs(filename, std::ifstream::in);
57-
CHECK(fs.is_open()) << "Can't open \"" << filename << "\"";
58-
IstreamInputStream input(&fs);
59-
bool success = google::protobuf::TextFormat::Parse(&input, proto);
60-
fs.close();
61-
return success;
46+
void ReadTFNetParamsFromBinaryBufferOrDie(const char* data, size_t len,
47+
tensorflow::GraphDef* param) {
48+
CHECK(ReadProtoFromBinaryBuffer(data, len, param))
49+
<< "Failed to parse GraphDef buffer";
6250
}
6351

64-
void ReadTFNetParamsFromBinaryFileOrDie(const char* param_file,
52+
void ReadTFNetParamsFromTextFileOrDie(const char* param_file,
6553
tensorflow::GraphDef* param) {
66-
CHECK(ReadProtoFromBinaryFileTF(param_file, param))
67-
<< "Failed to parse GraphDef file: " << param_file;
54+
CHECK(ReadProtoFromTextFile(param_file, param))
55+
<< "Failed to parse GraphDef file: " << param_file;
6856
}
6957

70-
void ReadTFNetParamsFromTextFileOrDie(const char* param_file,
71-
tensorflow::GraphDef* param) {
72-
CHECK(ReadProtoFromTextFileTF(param_file, param))
73-
<< "Failed to parse GraphDef file: " << param_file;
58+
void ReadTFNetParamsFromTextBufferOrDie(const char* data, size_t len,
59+
tensorflow::GraphDef* param) {
60+
CHECK(ReadProtoFromTextBuffer(data, len, param))
61+
<< "Failed to parse GraphDef buffer";
7462
}
7563

7664
}

‎modules/dnn/src/tensorflow/tf_io.hpp

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,13 @@ void ReadTFNetParamsFromBinaryFileOrDie(const char* param_file,
2525
void ReadTFNetParamsFromTextFileOrDie(const char* param_file,
2626
tensorflow::GraphDef* param);
2727

28+
// Read parameters from a memory buffer into a GraphDef proto message.
29+
void ReadTFNetParamsFromBinaryBufferOrDie(const char* data, size_t len,
30+
tensorflow::GraphDef* param);
31+
32+
void ReadTFNetParamsFromTextBufferOrDie(const char* data, size_t len,
33+
tensorflow::GraphDef* param);
34+
2835
}
2936
}
3037

‎modules/dnn/test/test_caffe_importer.cpp

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,24 @@ static std::string _tf(TString filename)
5555
return (getOpenCVExtraDir() + "/dnn/") + filename;
5656
}
5757

58+
TEST(Test_Caffe, memory_read)
59+
{
60+
const string proto = findDataFile("dnn/bvlc_googlenet.prototxt", false);
61+
const string model = findDataFile("dnn/bvlc_googlenet.caffemodel", false);
62+
63+
string dataProto;
64+
ASSERT_TRUE(readFileInMemory(proto, dataProto));
65+
string dataModel;
66+
ASSERT_TRUE(readFileInMemory(model, dataModel));
67+
68+
Net net = readNetFromCaffe(dataProto.c_str(), dataProto.size());
69+
ASSERT_FALSE(net.empty());
70+
71+
Net net2 = readNetFromCaffe(dataProto.c_str(), dataProto.size(),
72+
dataModel.c_str(), dataModel.size());
73+
ASSERT_FALSE(net2.empty());
74+
}
75+
5876
TEST(Test_Caffe, read_gtsrb)
5977
{
6078
Net net = readNetFromCaffe(_tf("gtsrb.prototxt"));
@@ -67,13 +85,26 @@ TEST(Test_Caffe, read_googlenet)
6785
ASSERT_FALSE(net.empty());
6886
}
6987

70-
TEST(Reproducibility_AlexNet, Accuracy)
88+
typedef testing::TestWithParam<tuple<bool> > Reproducibility_AlexNet;
89+
TEST_P(Reproducibility_AlexNet, Accuracy)
7190
{
91+
bool readFromMemory = get<0>(GetParam());
7292
Net net;
7393
{
7494
const string proto = findDataFile("dnn/bvlc_alexnet.prototxt", false);
7595
const string model = findDataFile("dnn/bvlc_alexnet.caffemodel", false);
76-
net = readNetFromCaffe(proto, model);
96+
if (readFromMemory)
97+
{
98+
string dataProto;
99+
ASSERT_TRUE(readFileInMemory(proto, dataProto));
100+
string dataModel;
101+
ASSERT_TRUE(readFileInMemory(model, dataModel));
102+
103+
net = readNetFromCaffe(dataProto.c_str(), dataProto.size(),
104+
dataModel.c_str(), dataModel.size());
105+
}
106+
else
107+
net = readNetFromCaffe(proto, model);
77108
ASSERT_FALSE(net.empty());
78109
}
79110

@@ -86,6 +117,8 @@ TEST(Reproducibility_AlexNet, Accuracy)
86117
normAssert(ref, out);
87118
}
88119

120+
INSTANTIATE_TEST_CASE_P(Test_Caffe, Reproducibility_AlexNet, testing::Values(true, false));
121+
89122
#if !defined(_WIN32) || defined(_WIN64)
90123
TEST(Reproducibility_FCN, Accuracy)
91124
{

‎modules/dnn/test/test_common.hpp

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,4 +57,23 @@ inline void normAssert(cv::InputArray ref, cv::InputArray test, const char *comm
5757
EXPECT_LE(normInf, lInf) << comment;
5858
}
5959

60+
inline bool readFileInMemory(const std::string& filename, std::string& content)
61+
{
62+
std::ios::openmode mode = std::ios::in | std::ios::binary;
63+
std::ifstream ifs(filename.c_str(), mode);
64+
if (!ifs.is_open())
65+
return false;
66+
67+
content.clear();
68+
69+
ifs.seekg(0, std::ios::end);
70+
content.reserve(ifs.tellg());
71+
ifs.seekg(0, std::ios::beg);
72+
73+
content.assign((std::istreambuf_iterator<char>(ifs)),
74+
std::istreambuf_iterator<char>());
75+
76+
return true;
77+
}
78+
6079
#endif

0 commit comments

Comments
 (0)