diff --git a/modules/dnn/src/onnx/onnx_graph_simplifier.cpp b/modules/dnn/src/onnx/onnx_graph_simplifier.cpp index 4b857ebc16..b4284af7b1 100644 --- a/modules/dnn/src/onnx/onnx_graph_simplifier.cpp +++ b/modules/dnn/src/onnx/onnx_graph_simplifier.cpp @@ -1278,14 +1278,33 @@ public: } } } + // extract axis from original Gather node + axis = 0; + opencv_onnx::NodeProto* origGatherNode = + inpNode.dynamicCast()->node; + for (int i = 0; i < origGatherNode->attribute_size(); i++) { + opencv_onnx::AttributeProto attr = origGatherNode->attribute(i); + if (attr.name() == "axis") + axis = attr.i(); + } } } } return retVal; } + virtual void finalize(const Ptr& net, + const Ptr& fusedNode, + std::vector >& /*inputs*/) CV_OVERRIDE + { + opencv_onnx::NodeProto* node = fusedNode.dynamicCast()->node; + opencv_onnx::AttributeProto* new_attr = node->add_attribute(); + new_attr->set_name("axis"); + new_attr->set_i(axis); + } + private: - int cast, gather; + int cast, gather, axis; }; /* Constant folding shape for Expand.