diff --git a/modules/dnn/src/layers/flatten_layer.cpp b/modules/dnn/src/layers/flatten_layer.cpp index 5ac79da771..7747bf9aad 100644 --- a/modules/dnn/src/layers/flatten_layer.cpp +++ b/modules/dnn/src/layers/flatten_layer.cpp @@ -117,9 +117,12 @@ public: CV_Assert(startAxis >= 0 && startAxis <= numAxes); if (_onnxMode) { + int onnxAxis = _startAxis; + if (onnxAxis < 0) onnxAxis += numAxes; + onnxAxis = std::max(0, std::min(onnxAxis, numAxes)); size_t outer = 1, inner = 1; int i = 0; - for (; i < startAxis; i++) + for (; i < onnxAxis; i++) outer *= inputs[0][i]; for (; i < numAxes; i++) inner *= inputs[0][i]; @@ -165,8 +168,10 @@ public: inputs_arr.getMatVector(inputs); int numAxes = inputs[0].dims; - _startAxis = normalize_axis(_startAxis, numAxes); - _endAxis = normalize_axis(_endAxis, numAxes); + if (!_onnxMode) { + _startAxis = normalize_axis(_startAxis, numAxes); + _endAxis = normalize_axis(_endAxis, numAxes); + } } #ifdef HAVE_OPENCL @@ -253,11 +258,14 @@ public: std::vector dims = ieInpNode.get_shape(); int numAxes = dims.size(); - int startAxis = normalize_axis(_startAxis, numAxes); - int endAxis = normalize_axis(_endAxis, numAxes); + int startAxis = _startAxis; + if (startAxis < 0) startAxis += numAxes; + startAxis = std::max(0, std::min(startAxis, numAxes)); + int endAxis = _endAxis; + if (endAxis < 0) endAxis += numAxes; + endAxis = std::max(0, std::min(endAxis, numAxes - 1)); CV_Assert(startAxis >= 0); - CV_Assert(endAxis >= startAxis && endAxis < numAxes); int64_t flattenedDimensionSize = std::accumulate(dims.begin() + startAxis, dims.begin() + endAxis + 1, 1, std::multiplies()); diff --git a/modules/dnn/src/onnx/onnx_importer.cpp b/modules/dnn/src/onnx/onnx_importer.cpp index 46a21842dc..592a2cdb4a 100644 --- a/modules/dnn/src/onnx/onnx_importer.cpp +++ b/modules/dnn/src/onnx/onnx_importer.cpp @@ -2274,7 +2274,9 @@ void ONNXImporter::parseFlatten(LayerParams& layerParams, const opencv_onnx::Nod { constBlobsExtraInfo.insert(std::make_pair(node_proto.output(0), getBlobExtraInfo(node_proto, 0))); } - int axis = normalize_axis(axis_, input.dims); + int axis = axis_; + if (axis < 0) axis += input.dims; + axis = std::max(0, std::min(axis, input.dims)); int out_size[2] = {1, 1}; for (int i = 0; i < axis; ++i) @@ -2293,16 +2295,46 @@ void ONNXImporter::parseFlatten(LayerParams& layerParams, const opencv_onnx::Nod IterShape_t shapeIt = outShapes.find(node_proto.input(0)); CV_Assert(shapeIt != outShapes.end()); MatShape inpShape = shapeIt->second; - int axis = normalize_axis(axis_, inpShape.size()); + int axis = axis_; + if (axis < 0) axis += (int)inpShape.size(); + axis = std::max(0, std::min(axis, (int)inpShape.size())); - if (axis == 0 || axis == inpShape.size()) + if (axis == (int)inpShape.size()) + { + LayerParams reshapeLp; + reshapeLp.name = layerParams.name + "/reshape"; + reshapeLp.type = "Reshape"; + CV_Assert(layer_id.find(reshapeLp.name) == layer_id.end()); + inpShape.push_back(1); + reshapeLp.set("dim", DictValue::arrayInt(&inpShape[0], inpShape.size())); + opencv_onnx::NodeProto proto; + proto.add_input(node_proto.input(0)); + proto.add_output(reshapeLp.name); + addLayer(reshapeLp, proto); + LayerParams flatLp; + flatLp.name = layerParams.name + "/flatten"; + flatLp.type = "Flatten"; + CV_Assert(layer_id.find(flatLp.name) == layer_id.end()); + flatLp.set("axis", 0); + flatLp.set("end_axis", (int)inpShape.size() - 2); + opencv_onnx::NodeProto proto2; + proto2.add_input(reshapeLp.name); + proto2.add_output(flatLp.name); + addLayer(flatLp, proto2); + layerParams.type = "Identity"; + node_proto.set_input(0, flatLp.name); + addLayer(layerParams, node_proto); + return; + } + + if (axis == 0) { LayerParams reshapeLp; reshapeLp.name = layerParams.name + "/reshape"; reshapeLp.type = "Reshape"; CV_Assert(layer_id.find(reshapeLp.name) == layer_id.end()); - inpShape.insert(axis == 0 ? inpShape.begin() : inpShape.end(), 1); + inpShape.insert(inpShape.begin(), 1); reshapeLp.set("dim", DictValue::arrayInt(&inpShape[0], inpShape.size())); opencv_onnx::NodeProto proto; diff --git a/modules/dnn/test/test_onnx_importer.cpp b/modules/dnn/test/test_onnx_importer.cpp index 049e011c12..6f2687842d 100644 --- a/modules/dnn/test/test_onnx_importer.cpp +++ b/modules/dnn/test/test_onnx_importer.cpp @@ -1215,6 +1215,7 @@ TEST_P(Test_ONNX_layers, DynamicReshape) testONNXModels("dynamic_reshape_opset_11"); testONNXModels("flatten_by_prod"); testONNXModels("flatten_const"); + testONNXModels("flatten_axis_numaxes"); } TEST_P(Test_ONNX_layers, Reshape)