From f978c995237124b17cc5c54ec46f64c6159c893f Mon Sep 17 00:00:00 2001
From: Yuantao Feng <yuantao.feng@opencv.org.cn>
Date: Mon, 25 Dec 2023 19:42:05 +0800
Subject: [PATCH] Merge pull request #24753 from fengyuentau:einsum_importer

dnn onnx: support constaint inputs in einsum importer #24753

Merge with https://github.com/opencv/opencv_extra/pull/1132.

Resolves https://github.com/opencv/opencv/issues/24697

Credits to @LaurentBerger.

---

This is a workaround. I suggest to get input shapes and calculate the output shapes in `getMemoryShapes` so as to keep the best compatibility. It is not always robust getting shapes during the importer stage and we should avoid that as much as possible.

### 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/onnx/onnx_importer.cpp  | 39 ++++++++++++++++++++-----
 modules/dnn/test/test_onnx_importer.cpp |  4 +++
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/modules/dnn/src/onnx/onnx_importer.cpp b/modules/dnn/src/onnx/onnx_importer.cpp
index d65f155a55..115738999a 100644
--- a/modules/dnn/src/onnx/onnx_importer.cpp
+++ b/modules/dnn/src/onnx/onnx_importer.cpp
@@ -3220,19 +3220,44 @@ void ONNXImporter::parseSimpleLayers(LayerParams& layerParams, const opencv_onnx
     addLayer(layerParams, node_proto);
 }
 
-
 void ONNXImporter::parseEinsum(LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto)
 {
     std::vector<MatShape> einsumInpShapes;
     for (int j = 0; j < node_proto.input_size(); j++)
     {
-        const auto& inputLayerName = node_proto.input(j);
-        auto it = outShapes.find(inputLayerName);
-        if (it != outShapes.end())
-        {
-            einsumInpShapes.emplace_back(it->second);
+        // create Const layer for constants and mark its shape
+        std::vector<int> input_shape;
+        if (layer_id.find(node_proto.input(j)) == layer_id.end()) {
+            Mat blob = getBlob(node_proto, j);
+
+            LayerParams const_params;
+            const_params.name = node_proto.input(j);
+            const_params.type = "Const";
+            const_params.blobs.push_back(blob);
+
+            opencv_onnx::NodeProto proto;
+            proto.add_output(const_params.name);
+            addLayer(const_params, proto);
+
+            input_shape.resize(blob.dims);
+            for (size_t i = 0; i < input_shape.size(); i++) {
+                input_shape[i] = blob.size[i];
+            }
+        }
+
+        // also try getting shape from inferred shapes
+        if (input_shape.empty()) {
+            const auto& inputLayerName = node_proto.input(j);
+            auto it = outShapes.find(inputLayerName);
+            if (it != outShapes.end()) {
+                input_shape = it->second;
+            }
+        }
+
+        if (input_shape.empty()) {
+            CV_Error(Error::StsAssert, format("ERROR input shape of %s not found", node_proto.input(j).c_str()));
         } else {
-            CV_Error(Error::StsAssert, "ERROR input shape not found");
+            einsumInpShapes.emplace_back(input_shape);
         }
     }
 
diff --git a/modules/dnn/test/test_onnx_importer.cpp b/modules/dnn/test/test_onnx_importer.cpp
index ad6efbe77a..080b07243c 100644
--- a/modules/dnn/test/test_onnx_importer.cpp
+++ b/modules/dnn/test/test_onnx_importer.cpp
@@ -1514,6 +1514,10 @@ TEST_P(Test_ONNX_layers, Einsum_transpose)
     testONNXModels("einsum_transpose", npy, 0, 0, false, false, 1);
 }
 
+TEST_P(Test_ONNX_layers, Einsum_const_inputs) {
+    testONNXModels("einsum_const_inputs", npy, 0, 0, false, false, 1);
+}
+
 TEST_P(Test_ONNX_layers, Pad2d_Unfused)
 {
     testONNXModels("ReflectionPad2d");