// 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.
// Copyright (C) 2026, BigVision LLC, all rights reserved.
// Third party copyrights are property of their respective owners.
//
// Companion sample for tutorial:
//   doc/tutorials/dnn/dnn_custom_layers/dnn_custom_layers.md (ONNX section)
//
// Models used by this sample are generated by:
//   opencv_extra/testdata/dnn/onnx/generate_custom_layer_models.py
// and live at:
//   opencv_extra/testdata/dnn/onnx/models/custom_layer_default_domain.onnx
//   opencv_extra/testdata/dnn/onnx/models/custom_layer_custom_domain.onnx

#include <opencv2/dnn.hpp>
#include <opencv2/core.hpp>
#include <iostream>
#include <vector>
#include <string>

using namespace cv;
using namespace cv::dnn;
using namespace std;

//! [CustomScaleBiasLayer]
// y = scale * x + bias, with scale/bias read from ONNX node attributes.
class CustomScaleBiasLayer CV_FINAL : public Layer
{
public:
    CustomScaleBiasLayer(const LayerParams& params) : Layer(params)
    {
        scale = params.get<float>("scale", 1.f);
        bias  = params.get<float>("bias",  0.f);
    }

    static Ptr<Layer> create(LayerParams& params)
    {
        return makePtr<CustomScaleBiasLayer>(params);
    }

    bool getMemoryShapes(const vector<MatShape>& inpts,
                         const int /*requiredOutputs*/,
                         vector<MatShape>& outShapes,
                         vector<MatShape>& /*internals*/) const CV_OVERRIDE
    {
        outShapes.assign(1, inpts[0]);
        return false;
    }

    void forward(InputArrayOfArrays inputs_arr, OutputArrayOfArrays outputs_arr,
                 OutputArrayOfArrays) CV_OVERRIDE
    {
        vector<Mat> inps, outs;
        inputs_arr.getMatVector(inps);
        outputs_arr.getMatVector(outs);
        inps[0].convertTo(outs[0], outs[0].type(), scale, bias);
    }

private:
    float scale, bias;
};
//! [CustomScaleBiasLayer]

static void printShape(const string& tag, const Mat& m)
{
    cout << tag;
    for (int d = 0; d < m.dims; ++d)
        cout << (d ? "x" : "") << m.size[d];
    cout << "\n";
}

int main(int argc, char** argv)
{
    const string keys =
        "{ help  h |                                  | Print help message }"
        "{ model m |                                  | Path to ONNX model containing the custom op }"
        "{ op    o | MyCustomOp                       | Op key for registration. Default-domain: just the op_type "
                                                       "(e.g. MyCustomOp). Custom-domain: <domain>.<op_type> "
                                                       "(e.g. my.namespace.MyDomainOp) }";
    CommandLineParser parser(argc, argv, keys);
    parser.about("Demonstrates importing an ONNX model that contains a custom (non-standard) op "
                 "by registering a user-defined layer with cv::dnn::LayerFactory.");
    if (parser.has("help") || !parser.has("model"))
    {
        parser.printMessage();
        return 0;
    }

    const string modelPath = parser.get<string>("model");
    const string opKey     = parser.get<string>("op");

    //! [Register CustomScaleBiasLayer]
    // ONNX op-type lookup: layers in the default `ai.onnx` domain are registered
    // under their op_type; layers in a non-default domain are registered under
    // "<domain>.<op_type>" (e.g. "my.namespace.MyDomainOp").
    LayerFactory::registerLayer(opKey, CustomScaleBiasLayer::create);
    //! [Register CustomScaleBiasLayer]

    Net net = readNetFromONNX(modelPath);
    if (net.empty())
    {
        cerr << "Failed to load model: " << modelPath << "\n";
        LayerFactory::unregisterLayer(opKey);
        return 1;
    }

    // The companion ONNX models accept a 1x3x4x4 float input.
    Mat input(vector<int>{1, 3, 4, 4}, CV_32F, Scalar(1.0f));
    net.setInput(input);
    Mat out = net.forward();

    cout << "Loaded " << modelPath << " using custom op key '" << opKey << "'.\n";
    printShape("Input shape:  ", input);
    printShape("Output shape: ", out);
    cout << "Output[0,0,0,0] = " << out.ptr<float>()[0]
         << "  (= scale * 1.0 + bias for the registered op)\n";

    LayerFactory::unregisterLayer(opKey);
    return 0;
}