--- /dev/null
+// 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 "layers_common.hpp"
+
+
+namespace cv { namespace dnn {
+
+class ArgLayerImpl CV_FINAL : public ArgLayer
+{
+public:
+ enum class ArgOp
+ {
+ MIN = 0,
+ MAX = 1,
+ };
+
+ ArgLayerImpl(const LayerParams& params)
+ {
+ setParamsFrom(params);
+
+ axis = params.get<int>("axis", 0);
+ keepdims = (params.get<int>("keepdims", 1) == 1);
+ select_last_index = (params.get<int>("select_last_index", 0) == 1);
+
+ const std::string& argOp = params.get<std::string>("op");
+
+ if (argOp == "max")
+ {
+ op = ArgOp::MAX;
+ }
+ else if (argOp == "min")
+ {
+ op = ArgOp::MIN;
+ }
+ else
+ {
+ CV_Error(Error::StsBadArg, "Unsupported operation");
+ }
+ }
+
+ virtual bool supportBackend(int backendId) CV_OVERRIDE
+ {
+ return backendId == DNN_BACKEND_OPENCV && preferableTarget == DNN_TARGET_CPU;
+ }
+
+ void handleKeepDims(MatShape& shape, const int axis_) const
+ {
+ if (keepdims)
+ {
+ shape[axis_] = 1;
+ }
+ else
+ {
+ shape.erase(shape.begin() + axis_);
+ }
+ }
+
+ virtual bool getMemoryShapes(const std::vector<MatShape> &inputs,
+ const int requiredOutputs,
+ std::vector<MatShape> &outputs,
+ std::vector<MatShape> &internals) const CV_OVERRIDE
+ {
+ MatShape inpShape = inputs[0];
+
+ const int axis_ = normalize_axis(axis, inpShape);
+ handleKeepDims(inpShape, axis_);
+ outputs.assign(1, inpShape);
+
+ return false;
+ }
+
+ void forward(InputArrayOfArrays inputs_arr, OutputArrayOfArrays outputs_arr, OutputArrayOfArrays internals_arr) CV_OVERRIDE
+ {
+ CV_TRACE_FUNCTION();
+ CV_TRACE_ARG_VALUE(name, "name", name.c_str());
+
+ std::vector<Mat> inputs, outputs;
+ inputs_arr.getMatVector(inputs);
+ outputs_arr.getMatVector(outputs);
+
+ CV_Assert_N(inputs.size() == 1, outputs.size() == 1);
+ std::vector<int> outShape = shape(outputs[0]);
+ Mat output(outShape, CV_32SC1);
+
+ switch (op)
+ {
+ case ArgOp::MIN:
+ cv::reduceArgMin(inputs[0], output, axis, select_last_index);
+ break;
+ case ArgOp::MAX:
+ cv::reduceArgMax(inputs[0], output, axis, select_last_index);
+ break;
+ default:
+ CV_Error(Error::StsBadArg, "Unsupported operation.");
+ }
+
+ output = output.reshape(1, outShape);
+ output.convertTo(outputs[0], CV_32FC1);
+ }
+
+private:
+ // The axis in which to compute the arg indices. Accepted range is [-r, r-1] where r = rank(data).
+ int axis;
+ // Keep the reduced dimension or not
+ bool keepdims;
+ // Whether to select the first or the last index or Max/Min.
+ bool select_last_index;
+ // Operation to be performed
+ ArgOp op;
+};
+
+Ptr<ArgLayer> ArgLayer::create(const LayerParams& params)
+{
+ return Ptr<ArgLayer>(new ArgLayerImpl(params));
+}
+
+}} // namespace cv::dnn
const DispatchMap dispatch;
static const DispatchMap buildDispatchMap();
+ void parseArg (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
void parseMaxPool (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
void parseAveragePool (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
void parseReduce (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
}
}
+void ONNXImporter::parseArg(LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto)
+{
+ const std::string& layer_type = node_proto.op_type();
+ layerParams.type = "Arg";
+ layerParams.set("op", layer_type == "ArgMax" ? "max" : "min");
+ addLayer(layerParams, node_proto);
+}
+
void setCeilMode(LayerParams& layerParams)
{
// auto_pad attribute is deprecated and uses ceil
{
DispatchMap dispatch;
+ dispatch["ArgMax"] = dispatch["ArgMin"] = &ONNXImporter::parseArg;
dispatch["MaxPool"] = &ONNXImporter::parseMaxPool;
dispatch["AveragePool"] = &ONNXImporter::parseAveragePool;
dispatch["GlobalAveragePool"] = dispatch["GlobalMaxPool"] = dispatch["ReduceMean"] = dispatch["ReduceSum"] =