Merge pull request #20733 from rogday:argmaxnd

Implement ArgMax and ArgMin

* add reduceArgMax and reduceArgMin

* fix review comments

* address review concerns

modules/core/src/minmax.cpp

@@ -7,6 +7,9 @@
 #include "opencl_kernels_core.hpp"
 #include "opencv2/core/openvx/ovx_defs.hpp"
 #include "stat.hpp"
+#include "opencv2/core/detail/dispatch_helper.impl.hpp"
+
+#include <algorithm>
 
 #undef HAVE_IPP
 #undef CV_IPP_RUN_FAST
@@ -1570,3 +1573,118 @@ void cv::minMaxLoc( InputArray _img, double* minVal, double* maxVal,
     if( maxLoc )
         std::swap(maxLoc->x, maxLoc->y);
 }
+
+enum class ReduceMode
+{
+    FIRST_MIN = 0, //!< get index of first min occurrence
+    LAST_MIN  = 1, //!< get index of last min occurrence
+    FIRST_MAX = 2, //!< get index of first max occurrence
+    LAST_MAX  = 3, //!< get index of last max occurrence
+};
+
+template <typename T>
+struct reduceMinMaxImpl
+{
+    void operator()(const cv::Mat& src, cv::Mat& dst, ReduceMode mode, const int axis) const
+    {
+        switch(mode)
+        {
+        case ReduceMode::FIRST_MIN:
+            reduceMinMaxApply<std::less>(src, dst, axis);
+            break;
+        case ReduceMode::LAST_MIN:
+            reduceMinMaxApply<std::less_equal>(src, dst, axis);
+            break;
+        case ReduceMode::FIRST_MAX:
+            reduceMinMaxApply<std::greater>(src, dst, axis);
+            break;
+        case ReduceMode::LAST_MAX:
+            reduceMinMaxApply<std::greater_equal>(src, dst, axis);
+            break;
+        }
+    }
+
+    template <template<class> class Cmp>
+    static void reduceMinMaxApply(const cv::Mat& src, cv::Mat& dst, const int axis)
+    {
+        Cmp<T> cmp;
+
+        const auto *src_ptr = src.ptr<T>();
+        auto *dst_ptr = dst.ptr<int32_t>();
+
+        const size_t outer_size = src.total(0, axis);
+        const auto mid_size = static_cast<size_t>(src.size[axis]);
+
+        const size_t outer_step = src.total(axis);
+        const size_t dst_step = dst.total(axis);
+
+        const size_t mid_step = src.total(axis + 1);
+
+        for (size_t outer = 0; outer < outer_size; ++outer)
+        {
+            const size_t outer_offset = outer * outer_step;
+            const size_t dst_offset = outer * dst_step;
+            for (size_t mid = 0; mid != mid_size; ++mid)
+            {
+                const size_t src_offset = outer_offset + mid * mid_step;
+                for (size_t inner = 0; inner < mid_step; inner++)
+                {
+                    int32_t& index = dst_ptr[dst_offset + inner];
+
+                    const size_t prev = outer_offset + index * mid_step + inner;
+                    const size_t curr = src_offset + inner;
+
+                    if (cmp(src_ptr[curr], src_ptr[prev]))
+                    {
+                        index = static_cast<int32_t>(mid);
+                    }
+                }
+            }
+        }
+    }
+};
+
+static void reduceMinMax(cv::InputArray src, cv::OutputArray dst, ReduceMode mode, int axis)
+{
+    CV_INSTRUMENT_REGION();
+
+    cv::Mat srcMat = src.getMat();
+    axis = (axis + srcMat.dims) % srcMat.dims;
+    CV_Assert(srcMat.channels() == 1 && axis >= 0 && axis < srcMat.dims);
+
+    std::vector<int> sizes(srcMat.dims);
+    std::copy(srcMat.size.p, srcMat.size.p + srcMat.dims, sizes.begin());
+    sizes[axis] = 1;
+
+    dst.create(srcMat.dims, sizes.data(), CV_32SC1); // indices
+    cv::Mat dstMat = dst.getMat();
+    dstMat.setTo(cv::Scalar::all(0));
+
+    if (!srcMat.isContinuous())
+    {
+        srcMat = srcMat.clone();
+    }
+
+    bool needs_copy = !dstMat.isContinuous();
+    if (needs_copy)
+    {
+        dstMat = dstMat.clone();
+    }
+
+    cv::detail::depthDispatch<reduceMinMaxImpl>(srcMat.depth(), srcMat, dstMat, mode, axis);
+
+    if (needs_copy)
+    {
+        dstMat.copyTo(dst);
+    }
+}
+
+void cv::reduceArgMin(InputArray src, OutputArray dst, int axis, bool lastIndex)
+{
+    reduceMinMax(src, dst, lastIndex ? ReduceMode::LAST_MIN : ReduceMode::FIRST_MIN, axis);
+}
+
+void cv::reduceArgMax(InputArray src, OutputArray dst, int axis, bool lastIndex)
+{
+    reduceMinMax(src, dst, lastIndex ? ReduceMode::LAST_MAX : ReduceMode::FIRST_MAX, axis);
+}