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

@@ -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<size_t> 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<size_t>());
 

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;

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)