From e71758cfdf217590c613c371aa03c5b600e926ce Mon Sep 17 00:00:00 2001 From: Dmitry Kurtaev Date: Wed, 12 Dec 2018 17:36:17 +0300 Subject: [PATCH] Operate with shapes in ONNX models --- modules/dnn/src/onnx/onnx_importer.cpp | 126 +++++++++++++++++++++++- modules/dnn/test/test_onnx_importer.cpp | 4 + 2 files changed, 128 insertions(+), 2 deletions(-) diff --git a/modules/dnn/src/onnx/onnx_importer.cpp b/modules/dnn/src/onnx/onnx_importer.cpp index 22eda5046c..18e26f1b86 100644 --- a/modules/dnn/src/onnx/onnx_importer.cpp +++ b/modules/dnn/src/onnx/onnx_importer.cpp @@ -6,6 +6,7 @@ // Third party copyrights are property of their respective owners. #include "../precomp.hpp" +#include #ifdef HAVE_PROTOBUF @@ -134,9 +135,38 @@ Mat getMatFromTensor(opencv_onnx::TensorProto& tensor_proto) else CV_Error(Error::StsUnsupportedFormat, "Unsupported data type: " + opencv_onnx::TensorProto_DataType_Name(datatype)); + if (tensor_proto.dims_size() == 0) + blob.dims = 1; // To force 1-dimensional cv::Mat for scalars. return blob; } +void runLayer(Ptr layer, const std::vector& inputs, + std::vector& outputs) +{ + std::vector inpShapes(inputs.size()); + int ddepth = CV_32F; + for (size_t i = 0; i < inputs.size(); ++i) + { + inpShapes[i] = shape(inputs[i]); + if (i > 0 && ddepth != inputs[i].depth()) + CV_Error(Error::StsNotImplemented, "Mixed input data types."); + ddepth = inputs[i].depth(); + } + + std::vector outShapes, internalShapes; + layer->getMemoryShapes(inpShapes, 0, outShapes, internalShapes); + + std::vector internals(internalShapes.size()); + outputs.resize(outShapes.size()); + for (size_t i = 0; i < outShapes.size(); ++i) + outputs[i].create(outShapes[i], ddepth); + for (size_t i = 0; i < internalShapes.size(); ++i) + internals[i].create(internalShapes[i], ddepth); + + layer->finalize(inputs, outputs); + layer->forward(inputs, outputs, internals); +} + std::map ONNXImporter::getGraphTensors( const opencv_onnx::GraphProto& graph_proto) { @@ -292,6 +322,26 @@ void ONNXImporter::populateNet(Net dstNet) CV_Assert(model_proto.has_graph()); opencv_onnx::GraphProto graph_proto = model_proto.graph(); std::map constBlobs = getGraphTensors(graph_proto); + // List of internal blobs shapes. + std::map outShapes; + // Add all the inputs shapes. It includes as constant blobs as network's inputs shapes. + for (int i = 0; i < graph_proto.input_size(); ++i) + { + opencv_onnx::ValueInfoProto valueInfoProto = graph_proto.input(i); + CV_Assert(valueInfoProto.has_type()); + opencv_onnx::TypeProto typeProto = valueInfoProto.type(); + CV_Assert(typeProto.has_tensor_type()); + opencv_onnx::TypeProto::Tensor tensor = typeProto.tensor_type(); + CV_Assert(tensor.has_shape()); + opencv_onnx::TensorShapeProto tensorShape = tensor.shape(); + + MatShape inpShape(tensorShape.dim_size()); + for (int j = 0; j < inpShape.size(); ++j) + { + inpShape[j] = tensorShape.dim(j).dim_value(); + } + outShapes[valueInfoProto.name()] = inpShape; + } std::string framework_name; if (model_proto.has_producer_name()) { @@ -301,6 +351,7 @@ void ONNXImporter::populateNet(Net dstNet) // create map with network inputs (without const blobs) std::map layer_id; std::map::iterator layerId; + std::map::iterator shapeIt; // fill map: push layer name, layer id and output id std::vector netInputs; for (int j = 0; j < graph_proto.input_size(); j++) @@ -317,9 +368,9 @@ void ONNXImporter::populateNet(Net dstNet) LayerParams layerParams; opencv_onnx::NodeProto node_proto; - for(int i = 0; i < layersSize; i++) + for(int li = 0; li < layersSize; li++) { - node_proto = graph_proto.node(i); + node_proto = graph_proto.node(li); layerParams = getLayerParams(node_proto); CV_Assert(node_proto.output_size() >= 1); layerParams.name = node_proto.output(0); @@ -598,6 +649,65 @@ void ONNXImporter::populateNet(Net dstNet) { layerParams.type = "Padding"; } + else if (layer_type == "Shape") + { + CV_Assert(node_proto.input_size() == 1); + shapeIt = outShapes.find(node_proto.input(0)); + CV_Assert(shapeIt != outShapes.end()); + MatShape inpShape = shapeIt->second; + + Mat shapeMat(inpShape.size(), 1, CV_32S); + for (int j = 0; j < inpShape.size(); ++j) + shapeMat.at(j) = inpShape[j]; + shapeMat.dims = 1; + + constBlobs.insert(std::make_pair(layerParams.name, shapeMat)); + continue; + } + else if (layer_type == "Gather") + { + CV_Assert(node_proto.input_size() == 2); + CV_Assert(layerParams.has("axis")); + Mat input = getBlob(node_proto, constBlobs, 0); + Mat indexMat = getBlob(node_proto, constBlobs, 1); + CV_Assert_N(indexMat.type() == CV_32S, indexMat.total() == 1); + int index = indexMat.at(0); + int axis = layerParams.get("axis"); + + std::vector ranges(input.dims, Range::all()); + ranges[axis] = Range(index, index + 1); + + Mat out = input(ranges); + constBlobs.insert(std::make_pair(layerParams.name, out)); + continue; + } + else if (layer_type == "Concat") + { + bool hasVariableInps = false; + for (int i = 0; i < node_proto.input_size(); ++i) + { + if (layer_id.find(node_proto.input(i)) != layer_id.end()) + { + hasVariableInps = true; + break; + } + } + + if (!hasVariableInps) + { + std::vector inputs(node_proto.input_size()), concatenated; + for (size_t i = 0; i < inputs.size(); ++i) + { + inputs[i] = getBlob(node_proto, constBlobs, i); + } + Ptr concat = ConcatLayer::create(layerParams); + runLayer(concat, inputs, concatenated); + + CV_Assert(concatenated.size() == 1); + constBlobs.insert(std::make_pair(layerParams.name, concatenated[0])); + continue; + } + } else { for (int j = 0; j < node_proto.input_size(); j++) { @@ -609,12 +719,24 @@ void ONNXImporter::populateNet(Net dstNet) int id = dstNet.addLayer(layerParams.name, layerParams.type, layerParams); layer_id.insert(std::make_pair(layerParams.name, LayerInfo(id, 0))); + + std::vector layerInpShapes, layerOutShapes, layerInternalShapes; for (int j = 0; j < node_proto.input_size(); j++) { layerId = layer_id.find(node_proto.input(j)); if (layerId != layer_id.end()) { dstNet.connect(layerId->second.layerId, layerId->second.outputId, id, j); + // Collect input shapes. + shapeIt = outShapes.find(node_proto.input(j)); + CV_Assert(shapeIt != outShapes.end()); + layerInpShapes.push_back(shapeIt->second); } } + + // Compute shape of output blob for this layer. + Ptr layer = dstNet.getLayer(id); + layer->getMemoryShapes(layerInpShapes, 0, layerOutShapes, layerInternalShapes); + CV_Assert(!layerOutShapes.empty()); + outShapes[layerParams.name] = layerOutShapes[0]; } } diff --git a/modules/dnn/test/test_onnx_importer.cpp b/modules/dnn/test/test_onnx_importer.cpp index 61b06cc7cf..36e7450892 100644 --- a/modules/dnn/test/test_onnx_importer.cpp +++ b/modules/dnn/test/test_onnx_importer.cpp @@ -162,6 +162,10 @@ TEST_P(Test_ONNX_layers, MultyInputs) normAssert(ref, out, "", default_l1, default_lInf); } +TEST_P(Test_ONNX_layers, DynamicReshape) +{ + testONNXModels("dynamic_reshape"); +} INSTANTIATE_TEST_CASE_P(/*nothing*/, Test_ONNX_layers, dnnBackendsAndTargets());