From 3215d7e6ea9681d8e735eb77aa807f91d628102c Mon Sep 17 00:00:00 2001 From: vrooomy Date: Thu, 16 Apr 2026 18:09:36 +0530 Subject: [PATCH] fix Flatten axis=rank bug --- modules/dnn/src/onnx/onnx_importer.cpp | 40 +++++++++++++++++++++---- modules/dnn/test/test_onnx_importer.cpp | 1 + 2 files changed, 36 insertions(+), 5 deletions(-) diff --git a/modules/dnn/src/onnx/onnx_importer.cpp b/modules/dnn/src/onnx/onnx_importer.cpp index 11d24eadb8..11fde4278a 100644 --- a/modules/dnn/src/onnx/onnx_importer.cpp +++ b/modules/dnn/src/onnx/onnx_importer.cpp @@ -2225,7 +2225,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) @@ -2244,18 +2246,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.insert(axis == 0 ? inpShape.begin() : inpShape.end(), 1); + 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(inpShape.begin(), 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); diff --git a/modules/dnn/test/test_onnx_importer.cpp b/modules/dnn/test/test_onnx_importer.cpp index a69e94990b..cd1ae9d2fd 100644 --- a/modules/dnn/test/test_onnx_importer.cpp +++ b/modules/dnn/test/test_onnx_importer.cpp @@ -1170,6 +1170,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)