Commit 266692e1 authored by Vitaliy Lyudvichenko's avatar Vitaliy Lyudvichenko

Improving of Caffe importer compatibility

parent 8ecae046
......@@ -48,6 +48,7 @@ using namespace cv::dnn;
#include <iostream>
#include <fstream>
#include <sstream>
#include <algorithm>
#include <google/protobuf/message.h>
#include <google/protobuf/text_format.h>
......@@ -63,12 +64,21 @@ using ::google::protobuf::Reflection;
namespace
{
class CaffeImporter : public Importer
{
template<typename T>
static cv::String toString(const T &v)
{
std::ostringstream ss;
ss << v;
return ss.str();
}
class CaffeImporter : public Importer
{
caffe::NetParameter net;
caffe::NetParameter netBinary;
public:
public:
CaffeImporter(const char *pototxt, const char *caffeModel)
{
......@@ -250,10 +260,14 @@ namespace
int layerId, outNum;
};
std::vector<BlobNote> addedBlobs;
std::map<String, int> layerCounter;
void populateNet(Net dstNet)
{
int layersSize = net.layer_size();
std::vector<BlobNote> addedBlobs;
layerCounter.clear();
addedBlobs.clear();
addedBlobs.reserve(layersSize + 1);
//setup input layer names
......@@ -277,17 +291,23 @@ namespace
extractLayerParams(layer, layerParams);
extractBinaryLayerParms(layer, layerParams);
int repetitions = layerCounter[name]++;
if (repetitions)
name += String("_") + toString(repetitions);
int id = dstNet.addLayer(name, type, layerParams);
for (int inNum = 0; inNum < layer.bottom_size(); inNum++)
addInput(layer.bottom(inNum), id, inNum, dstNet, addedBlobs);
addInput(layer.bottom(inNum), id, inNum, dstNet);
for (int outNum = 0; outNum < layer.top_size(); outNum++)
addOutput(layer, id, outNum, addedBlobs);
addOutput(layer, id, outNum);
}
addedBlobs.clear();
}
void addOutput(const caffe::LayerParameter &layer, int layerId, int outNum, std::vector<BlobNote> &addedBlobs)
void addOutput(const caffe::LayerParameter &layer, int layerId, int outNum)
{
const std::string &name = layer.top(outNum);
......@@ -311,7 +331,7 @@ namespace
addedBlobs.push_back(BlobNote(name, layerId, outNum));
}
void addInput(const std::string &name, int layerId, int inNum, Net &dstNet, std::vector<BlobNote> &addedBlobs)
void addInput(const std::string &name, int layerId, int inNum, Net &dstNet)
{
int idx;
for (idx = (int)addedBlobs.size() - 1; idx >= 0; idx--)
......@@ -322,7 +342,7 @@ namespace
if (idx < 0)
{
CV_Error(Error::StsObjectNotFound, "Can't found output blob \"" + name + "\"");
CV_Error(Error::StsObjectNotFound, "Can't find output blob \"" + name + "\"");
return;
}
......@@ -334,8 +354,7 @@ namespace
}
};
};
}
......
......@@ -60,7 +60,7 @@ namespace dnn
{
template<typename T>
String toString(const T &v)
static String toString(const T &v)
{
std::ostringstream ss;
ss << v;
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment