From 8dfc2364766a42d6eb5600a1eafe6963c024586b Mon Sep 17 00:00:00 2001 From: Varun Jaiswal <96684656+varun-jaiswal17@users.noreply.github.com> Date: Thu, 16 Apr 2026 16:34:26 +0530 Subject: [PATCH] Merge pull request #28812 from varun-jaiswal17:fix-flatten-onnx-axis Fix flatten onnx axis==rank case #28812 fix ONNX Flatten layer incorrect output when axis equals input rank edge case Resolves : https://github.com/opencv/opencv/pull/28781#discussion_r3060088247 OpenCV Extra: https://github.com/opencv/opencv_extra/pull/1342 ### Pull Request Readiness Checklist See details at https://github.com/opencv/opencv/wiki/How_to_contribute#making-a-good-pull-request - [x] I agree to contribute to the project under Apache 2 License. - [x] To the best of my knowledge, the proposed patch is not based on a code under GPL or another license that is incompatible with OpenCV - [x] The PR is proposed to the proper branch - [x] There is a reference to the original bug report and related work - [x] There is accuracy test, performance test and test data in opencv_extra repository, if applicable Patch to opencv_extra has the same branch name. - [x] The feature is well documented and sample code can be built with the project CMake --- modules/dnn/src/layers/flatten_layer.cpp | 20 ++++++++---- modules/dnn/src/onnx/onnx_importer.cpp | 40 +++++++++++++++++++++--- modules/dnn/test/test_onnx_importer.cpp | 1 + 3 files changed, 51 insertions(+), 10 deletions(-) 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)