Merge pull request #27508 from abhishek-gola:if_layer_add

IfLayer add to new DNN engine #27508

### 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/include/opencv2/dnn/all_layers.hpp

@@ -585,6 +585,18 @@ CV__DNN_INLINE_NS_BEGIN
         static Ptr<RequantizeLayer> create(const LayerParams &params);
     };
 
+    // Forward declaration for computational Graph used by IfLayer
+    class Graph;
+
+    class CV_EXPORTS IfLayer : public Layer
+    {
+    public:
+        virtual int branch(InputArray arr) const = 0;
+
+        /** Factory: creates an IfLayer implementation. */
+        static Ptr<IfLayer> create(const LayerParams& params);
+    };
+
     class CV_EXPORTS ConcatLayer : public Layer
     {
     public:

modules/dnn/src/graph_buffer_allocator.cpp

@@ -135,6 +135,19 @@ struct BufferAllocator
             releaseBuffer(toBuf);
     }
 
+    template<typename _Tp> std::ostream&
+    dumpArgVec(std::ostream& strm, const std::string& name, const vector<_Tp>& vec) const
+    {
+        CV_Assert(vec.size() == netimpl->args.size());
+        strm << name << ": [";
+        size_t i, sz = vec.size();
+        for (i = 0; i < sz; i++) {
+            strm << "\n\t" << netimpl->args[i].name << ": " << vec[i];
+        }
+        strm << "]";
+        return strm;
+    }
+
     void assign()
     {
         netimpl->useCounts(usecounts);
@@ -152,6 +165,22 @@ struct BufferAllocator
     {
         if (!graph)
             return;
+        // Pre-assign buffers for *sub-graph* TEMP inputs/outputs only.
+        // (The main graph has already been handled by regular allocation logic.)
+        bool isSubGraph = graph.get() != netimpl->mainGraph.get();
+        if (isSubGraph)
+        {
+            const std::vector<Arg>& gr_inputs = graph->inputs();
+            for (const Arg& inarg : gr_inputs)
+            {
+                if (netimpl->argKind(inarg) == DNN_ARG_TEMP &&
+                    !netimpl->isConstArg(inarg) &&
+                    bufidxs.at(inarg.idx) < 0)
+                {
+                    bufidxs.at(inarg.idx) = getFreeBuffer();
+                }
+            }
+        }
         const std::vector<Ptr<Layer> >& prog = graph->prog();
         for (const auto& layer: prog) {
             bool inplace = false;
@@ -164,6 +193,13 @@ struct BufferAllocator
             size_t ninputs = inputs.size();
             size_t noutputs = outputs.size();
 
+            //std::cout << "graph '" << graph->name() << "', op '" << layer->name << "' (" << layer->type << ")\n";
+            //std::cout << "usecounts: " << usecounts << "\n";
+            //dumpArgVec(std::cout, "usecounts", usecounts) << "\n";
+            //std::cout << "freebufs: " << freebufs << "\n";
+            //std::cout << "buf_usecounts: " << buf_usecounts << "\n";
+            //dumpArgVec(std::cout, "bufidxs", bufidxs) << "\n";
+
             /*
              Determine if we can possibly re-use some of the input buffers for the output as well,
              in other words, whether we can run the operation in-place.
@@ -242,20 +278,26 @@ struct BufferAllocator
                     Arg thenOutarg = thenOutargs[i];
                     Arg elseOutarg = elseOutargs[i];
 
-                    if (!netimpl->isConstArg(thenOutarg) && usecounts[thenOutarg.idx] == 1)
+                    if (!netimpl->isConstArg(thenOutarg) &&
+                        usecounts[thenOutarg.idx] == 1 &&
+                        bufidxs[thenOutarg.idx] >= 0)
                         shareBuffer(outarg, thenOutarg);
-                    if (!netimpl->isConstArg(elseOutarg) && usecounts[elseOutarg.idx] == 1)
+                    if (!netimpl->isConstArg(elseOutarg) &&
+                        usecounts[elseOutarg.idx] == 1 &&
+                        bufidxs[thenOutarg.idx] >= 0)
                         shareBuffer(outarg, elseOutarg);
                 }
 
                 assign(thenBranch);
                 assign(elseBranch);
-
                 for (size_t i = 0; i < noutputs; i++) {
                     Arg thenOutarg = thenOutargs[i];
                     Arg elseOutarg = elseOutargs[i];
-                    releaseBuffer(bufidxs[thenOutarg.idx]);
-                    releaseBuffer(bufidxs[elseOutarg.idx]);
+                    if (!netimpl->isConstArg(thenOutarg) &&
+                        bufidxs[thenOutarg.idx] >= 0 &&
+                        !netimpl->isConstArg(elseOutarg) &&
+                        bufidxs[elseOutarg.idx] >= 0)
+                        shareBuffer(thenOutarg, elseOutarg);
                 }
             } else if (opname == "Loop") {
                 /*

modules/dnn/src/graph_const_fold.cpp

@@ -26,7 +26,6 @@ struct ConstFolding
         size_t nargs = netimpl->args.size();
         netimpl->__tensors__.resize(nargs);
         netimpl->useCounts(usecounts);
-        netimpl->scratchBufs.clear();
         processGraph(netimpl->mainGraph);
         netimpl->scratchBufs.clear();
     }
@@ -46,6 +45,7 @@ struct ConstFolding
 
     bool processGraph(Ptr<Graph>& graph)
     {
+        netimpl->scratchBufs.clear();
         bool modified = false;
         const std::vector<Ptr<Layer> >& prog = graph->prog();
         size_t i, nops = prog.size();
@@ -63,6 +63,7 @@ struct ConstFolding
                     if (processGraph(g))
                         modified = true;
                 }
+                continue;
             }
             const std::vector<Arg>& inputs = layer->inputs;
             const std::vector<Arg>& outputs = layer->outputs;

modules/dnn/src/init.cpp

@@ -84,6 +84,7 @@ void initializeLayerFactory()
     static ProtobufShutdown protobufShutdown; CV_UNUSED(protobufShutdown);
 #endif
 
+    CV_DNN_REGISTER_LAYER_CLASS(If,             IfLayer);
     CV_DNN_REGISTER_LAYER_CLASS(Concat,         ConcatLayer);
     CV_DNN_REGISTER_LAYER_CLASS(Concat2,        Concat2Layer);
     CV_DNN_REGISTER_LAYER_CLASS(ConstantOfShape, ConstantOfShapeLayer);

modules/dnn/src/layers/if_layer.cpp

@@ -0,0 +1,80 @@
+// This file is part of OpenCV project.
+// It is subject to the license terms in the LICENSE file found in the top-level directory
+// of this distribution and at http://opencv.org/license.html.
+#include "../precomp.hpp"
+#include "../net_impl.hpp"
+#include "layers_common.hpp"
+#include <opencv2/dnn.hpp>
+
+namespace cv { namespace dnn {
+
+class IfLayerImpl CV_FINAL : public IfLayer
+{
+public:
+    explicit IfLayerImpl(const LayerParams& params)
+    {
+        setParamsFrom(params);
+    }
+    virtual ~IfLayerImpl() = default;
+
+    std::vector<Ptr<Graph>>* subgraphs() const CV_OVERRIDE { return &thenelse; }
+
+    bool getMemoryShapes(const std::vector<MatShape>& /*inputs*/,
+                         const int requiredOutputs,
+                         std::vector<MatShape>& outputs,
+                         std::vector<MatShape>& internals) const CV_OVERRIDE
+    {
+        outputs.assign(std::max(1, requiredOutputs), MatShape());
+        internals.clear();
+        return false;
+    }
+
+    bool dynamicOutputShapes() const CV_OVERRIDE { return true; }
+
+    int branch(InputArray arr) const CV_OVERRIDE
+    {
+        Mat buf, *inp;
+        if (arr.kind() == _InputArray::MAT) {
+            inp = (Mat*)arr.getObj();
+        } else {
+            buf = arr.getMat();
+            inp = &buf;
+        }
+        CV_Assert(inp->total() == 1u);
+        bool flag;
+        switch (inp->depth())
+        {
+        case CV_8U: case CV_8S: case CV_Bool:
+            flag = *inp->ptr<char>() != 0; break;
+        case CV_16U: case CV_16S:
+            flag = *inp->ptr<short>() != 0; break;
+        case CV_16F:
+            flag = *inp->ptr<hfloat>() != 0; break;
+        case CV_16BF:
+            flag = *inp->ptr<hfloat>() != 0; break;
+        case CV_32U: case CV_32S:
+            flag = *inp->ptr<int>() != 0; break;
+        case CV_32F:
+            flag = *inp->ptr<float>() != 0; break;
+        case CV_64U: case CV_64S:
+            flag = *inp->ptr<long long>() != 0; break;
+        case CV_64F:
+            flag = *inp->ptr<double>() != 0; break;
+        default:
+            CV_Error_(Error::StsBadArg,
+                    ("If-layer condition: unsupported tensor type %s",
+                    typeToString(inp->type()).c_str()));
+        }
+        return (int)!flag;
+    }
+
+private:
+    mutable std::vector<Ptr<Graph>> thenelse;
+};
+
+Ptr<IfLayer> IfLayer::create(const LayerParams& params)
+{
+    return makePtr<IfLayerImpl>(params);
+}
+
+}}  // namespace cv::dnn

modules/dnn/src/net_impl2.cpp

@@ -253,7 +253,6 @@ Arg Net::Impl::newArg(const std::string& name, ArgKind kind, bool allowEmptyName
     return Arg(idx);
 }
 
-
 int Net::Impl::findDim(const std::string& dimname, bool insert)
 {
     if (!dimname.empty()) {
@@ -595,7 +594,6 @@ void Net::Impl::forwardGraph(Ptr<Graph>& graph, InputArrayOfArrays inputs_,
     if (graphofs_it == graphofs.end()) {
         CV_Error_(Error::StsObjectNotFound, ("graph '%s' does not belong to the model", graph->name().c_str()));
     }
-
     std::ostream& strm_ = dump_strm ? *dump_strm : std::cout;
     const std::vector<Ptr<Layer> >& prog = graph->prog();
     size_t i, nops = prog.size();
@@ -611,10 +609,8 @@ void Net::Impl::forwardGraph(Ptr<Graph>& graph, InputArrayOfArrays inputs_,
 
     size_t graph_ofs = (size_t)graphofs_it->second;
     CV_Assert(graph_ofs + nops <= totalLayers);
-
     if (inputs_.empty()) {
         // inputs are already set; it's only possible to do with the main graph
-        CV_Assert(isMainGraph);
         for (i = 0; i < n_gr_inputs; i++)
             CV_CheckFalse(argTensor(gr_inputs[i]).empty(), "Some of the model inputs were not set");
     }
@@ -660,7 +656,6 @@ void Net::Impl::forwardGraph(Ptr<Graph>& graph, InputArrayOfArrays inputs_,
                 traceArg(strm_, "Input", i, inp, false);
             }
         }
-
         bool dynamicOutShapes = layer->dynamicOutputShapes();
         if (!dynamicOutShapes) {
             allocateLayerOutputs(layer, inpTypes, inpShapes, outTypes, outShapes, outOrigData, outMats,
@@ -676,11 +671,27 @@ void Net::Impl::forwardGraph(Ptr<Graph>& graph, InputArrayOfArrays inputs_,
 
         timestamp = getTickCount();
 
-        // [TODO] handle If/Loop/...
-        CV_Assert(!layer->subgraphs());
-        if (finalizeLayers)
-            layer->finalize(inpMats, outMats);
-        layer->forward(inpMats, outMats, tempMats);
+        std::vector<Ptr<Graph> >* subgraphs = layer->subgraphs();
+        if (!subgraphs) {
+            if (finalizeLayers)
+                layer->finalize(inpMats, outMats);
+            layer->forward(inpMats, outMats, tempMats);
+        }
+        else {
+            Ptr<IfLayer> iflayer = layer.dynamicCast<IfLayer>();
+            if (iflayer) {
+                int branch = iflayer->branch(inpMats[0]);
+                Ptr<Graph> subgraph = subgraphs->at(branch);
+                std::vector<Mat> branchInputs;
+                if (inpMats.size() > 1)
+                    branchInputs.assign(inpMats.begin() + 1, inpMats.end());
+                forwardGraph(subgraph, branchInputs, outMats, false);
+            }
+            else {
+                CV_Error_(Error::StsNotImplemented,
+                          ("unknown layer type '%s' with subgraphs", layer->type.c_str()));
+            }
+        }
         CV_Assert(outMats.size() == noutputs);
 
         for (i = 0; i < noutputs; i++) {
@@ -748,6 +759,11 @@ void Net::Impl::updateUseCounts(const Ptr<Graph>& graph, std::vector<int>& useco
 {
     if (!graph)
         return;
+    const std::vector<Arg>& gr_outputs = graph->outputs();
+    for (const Arg& output: gr_outputs) {
+        CV_Assert(output.idx < (int)usecounts.size());
+        usecounts[output.idx]++;
+    }
     const std::vector<Ptr<Layer> >& prog = graph->prog();
     for (const Ptr<Layer>& layer: prog) {
         const std::vector<Arg>& inputs = layer->inputs;

modules/dnn/src/onnx/onnx_importer2.cpp

@@ -172,6 +172,7 @@ protected:
     void parseCast                 (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
     void parseClip                 (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
     void parseConcat               (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
+    void parseIf                   (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
     void parseConstant             (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
     void parseConstantOfShape      (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
     void parseConv                 (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
@@ -468,7 +469,9 @@ LayerParams ONNXImporter2::getLayerParams(const opencv_onnx::NodeProto& node_pro
             }
             else if (attribute_proto.has_g())
             {
-                CV_Error(Error::StsNotImplemented, format("DNN/ONNX/Attribute[%s]: 'Graph' is not supported", attribute_name.c_str()));
+                // CV_Error(Error::StsNotImplemented, format("DNN/ONNX/Attribute[%s]: 'Graph' is not supported", attribute_name.c_str()));
+                continue;
+
             }
             else if (attribute_proto.graphs_size() > 0)
             {
@@ -1488,6 +1491,31 @@ void ONNXImporter2::parseConcat(LayerParams& layerParams, const opencv_onnx::Nod
     addLayer(layerParams, node_proto);
 }
 
+void ONNXImporter2::parseIf(LayerParams& layerParams,
+                            const opencv_onnx::NodeProto& node_proto)
+{
+    CV_Assert(node_proto.input_size() >= 1);
+    layerParams.type = "If";
+
+    addLayer(layerParams, node_proto);
+
+    std::vector<Ptr<Graph> > thenelse(2);
+    for (int i = 0; i < node_proto.attribute_size(); ++i)
+    {
+        const auto& attr = node_proto.attribute(i);
+        if (attr.name() == "then_branch" || attr.name() == "else_branch") {
+            opencv_onnx::GraphProto branch = attr.g();
+            Ptr<Graph> graph = parseGraph(&branch, false);
+            thenelse[(int)(attr.name() == "else_branch")] = graph;
+        }
+    }
+
+    CV_Assert_N(!thenelse[0].empty(), !thenelse[1].empty());
+
+    Ptr<Layer>& ifLayer = curr_prog.back();
+    *ifLayer->subgraphs() = thenelse;
+}
+
 // https://github.com/onnx/onnx/blob/master/docs/Operators.md#Resize
 void ONNXImporter2::parseResize(LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto)
 {
@@ -2363,6 +2391,7 @@ void ONNXImporter2::buildDispatchMap_ONNX_AI(int opset_version)
     dispatch["Gather"] = &ONNXImporter2::parseGather;
     dispatch["GatherElements"] = &ONNXImporter2::parseGatherElements;
     dispatch["Concat"] = &ONNXImporter2::parseConcat;
+    dispatch["If"] = &ONNXImporter2::parseIf;
     dispatch["Resize"] = &ONNXImporter2::parseResize;
     dispatch["Upsample"] = &ONNXImporter2::parseUpsample;
     dispatch["SoftMax"] = dispatch["Softmax"] = dispatch["LogSoftmax"] = &ONNXImporter2::parseSoftMax;

modules/dnn/test/test_layers.cpp

@@ -2816,4 +2816,41 @@ TEST(Layer_LSTM, repeatedInference)
     EXPECT_EQ(diff2, 0.);
 }
 
+TEST(Layer_If, resize)
+{
+    // Skip this test when the classic DNN engine is explicitly requested. The
+    // "if" layer is supported only by the new engine.
+    auto engine_forced = static_cast<cv::dnn::EngineType>(
+            cv::utils::getConfigurationParameterSizeT("OPENCV_FORCE_DNN_ENGINE", cv::dnn::ENGINE_AUTO));
+    if (engine_forced == cv::dnn::ENGINE_CLASSIC)
+    {
+        // Mark the test as skipped and exit early.
+        applyTestTag(CV_TEST_TAG_DNN_SKIP_PARSER);
+        return;
+    }
+
+    const std::string imgname   = findDataFile("cv/shared/lena.png", true);
+    const std::string modelname = findDataFile("dnn/onnx/models/if_layer.onnx", true);
+
+    dnn::Net net = dnn::readNetFromONNX(modelname, ENGINE_NEW);
+    Mat src = imread(imgname), blob;
+    dnn::blobFromImage(src, blob, 1.0, cv::Size(), cv::Scalar(), false, false);
+
+    for (int f = 0; f <= 1; f++) {
+        Mat cond(1, 1, CV_BoolC1, cv::Scalar(f));
+
+        net.setInput(cond, "cond");
+        net.setInput(blob, "image");
+
+        std::vector<Mat> outs;
+        net.forward(outs);
+
+        std::vector<Mat> images;
+        dnn::imagesFromBlob(outs[0], images);
+        EXPECT_EQ(images.size(), 1u);
+        EXPECT_EQ(images[0].rows*(4 >> f), src.rows);
+        EXPECT_EQ(images[0].cols*(4 >> f), src.cols);
+    }
+}
+
 }} // namespace

modules/dnn/test/test_onnx_conformance.cpp

@@ -980,6 +980,7 @@ public:
     static std::set<std::string> opencl_fp16_deny_list;
     static std::set<std::string> opencl_deny_list;
     static std::set<std::string> cpu_deny_list;
+    static std::set<std::string> classic_deny_list;
 #ifdef HAVE_HALIDE
     static std::set<std::string> halide_deny_list;
 #endif
@@ -1058,6 +1059,18 @@ public:
             #include "test_onnx_conformance_layer_filter_opencv_cpu_denylist.inl.hpp"
         };
 
+        EngineType engine_forced =
+            (EngineType)utils::getConfigurationParameterSizeT(
+                "OPENCV_FORCE_DNN_ENGINE", ENGINE_AUTO);
+
+        if (engine_forced == ENGINE_CLASSIC) {
+            classic_deny_list = {
+#include "test_onnx_conformance_layer_filter_opencv_classic_denylist.inl.hpp"
+            };
+        } else {
+            classic_deny_list = {};
+        }
+
 #ifdef HAVE_HALIDE
         halide_deny_list = {
             #include "test_onnx_conformance_layer_filter__halide_denylist.inl.hpp"
@@ -1088,6 +1101,7 @@ std::set<std::string> Test_ONNX_conformance::opencv_deny_list;
 std::set<std::string> Test_ONNX_conformance::opencl_fp16_deny_list;
 std::set<std::string> Test_ONNX_conformance::opencl_deny_list;
 std::set<std::string> Test_ONNX_conformance::cpu_deny_list;
+std::set<std::string> Test_ONNX_conformance::classic_deny_list;
 #ifdef HAVE_HALIDE
 std::set<std::string> Test_ONNX_conformance::halide_deny_list;
 #endif
@@ -1113,6 +1127,12 @@ TEST_P(Test_ONNX_conformance, Layer_Test)
         applyTestTag(CV_TEST_TAG_DNN_SKIP_PARSER, CV_TEST_TAG_DNN_SKIP_ONNX_CONFORMANCE);
     }
 
+    // SKIP some more if we are in the 'classic engine' mode, where we don't support certain layers.
+    if (classic_deny_list.find(name) != classic_deny_list.end())
+    {
+        applyTestTag(CV_TEST_TAG_DNN_SKIP_PARSER, CV_TEST_TAG_DNN_SKIP_ONNX_CONFORMANCE);
+    }
+
     // SKIP when the test case is in the global deny list.
     if (global_deny_list.find(name) != global_deny_list.end())
     {

modules/dnn/test/test_onnx_conformance_layer_filter__openvino.inl.hpp

@@ -781,7 +781,7 @@ CASE(test_identity_opt)
 CASE(test_identity_sequence)
     // no filter
 CASE(test_if)
-    // no filter
+    SKIP;
 CASE(test_if_opt)
     // no filter
 CASE(test_if_seq)

modules/dnn/test/test_onnx_conformance_layer_filter_opencv_classic_denylist.inl.hpp

@@ -0,0 +1 @@
+"test_if",

modules/dnn/test/test_onnx_conformance_layer_parser_denylist.inl.hpp

@@ -127,7 +127,6 @@
 "test_gru_with_initial_bias", // ---- same as above ---
 "test_identity_opt", //  23221 illegal hardware instruction
 "test_identity_sequence",  // Issue:: Unkonwn error
-"test_if", // Issue::'Graph' is not supported in function 'getLayerParams'
 "test_if_opt", // Issue::Failed to allocate 17059022683624350 bytes in function 'OutOfMemoryError'
 "test_if_seq", // Issue::typeProto.has_tensor_type() in function 'dumpValueInfoProto'
 "test_isinf", // Issue::Can't create layer "onnx_node_output_0!y" of type "IsInf" in function 'getLayerInstance'