Fixed multidimensional count non-zero IPP implementation

mshabunin · mshabunin · commit f23b6ba652ec · 2017-05-23T13:23:59.000+03:00
diff --git a/modules/core/src/stat.cpp b/modules/core/src/stat.cpp
@@ -1333,23 +1333,23 @@ static bool ipp_countNonZero( Mat &src, int &res )
     {
         IppStatus       status;
         const Mat      *arrays[] = {&src, NULL};
-        uchar          *ptrs[1]  = {NULL};
-        NAryMatIterator it(arrays, ptrs);
+        Mat            planes[1];
+        NAryMatIterator it(arrays, planes, 1);
         IppiSize        size  = {(int)it.size*src.channels(), 1};
-
+        res = 0;
         for (size_t i = 0; i < it.nplanes; i++, ++it)
         {
             if(depth == CV_8U)
-                status = CV_INSTRUMENT_FUN_IPP(ippiCountInRange_8u_C1R, (const Ipp8u *)src.ptr(), (int)src.step, size, &count, 0, 0);
+                status = CV_INSTRUMENT_FUN_IPP(ippiCountInRange_8u_C1R, it.planes->ptr<Ipp8u>(), (int)it.planes->step, size, &count, 0, 0);
             else if(depth == CV_32F)
-                status = CV_INSTRUMENT_FUN_IPP(ippiCountInRange_32f_C1R, (const Ipp32f *)src.ptr(), (int)src.step, size, &count, 0, 0);
+                status = CV_INSTRUMENT_FUN_IPP(ippiCountInRange_32f_C1R, it.planes->ptr<Ipp32f>(), (int)it.planes->step, size, &count, 0, 0);
             else
                 return false;
 
-            if(status < 0)
+            if(status < 0 || (int)it.planes->total()*src.channels() < count)
                 return false;
 
-            res += (size.width*size.height - count);
+            res += (int)it.planes->total()*src.channels() - count;
         }
     }
 
diff --git a/modules/core/test/test_countnonzero.cpp b/modules/core/test/test_countnonzero.cpp
@@ -250,3 +250,29 @@ void CV_CountNonZeroTest::run(int)
 }
 
 TEST (Core_CountNonZero, accuracy) { CV_CountNonZeroTest test; test.safe_run(); }
+
+
+typedef testing::TestWithParam<std::tr1::tuple<int, int> > CountNonZeroND;
+
+TEST_P (CountNonZeroND, ndim)
+{
+    const int dims = std::tr1::get<0>(GetParam());
+    const int type = std::tr1::get<1>(GetParam());
+    const int ONE_SIZE = 5;
+
+    vector<int> sizes(dims);
+    fill(sizes.begin(), sizes.end(), ONE_SIZE);
+
+    Mat data(sizes, CV_MAKETYPE(type, 1));
+    data = 0;
+    EXPECT_EQ(0, cv::countNonZero(data));
+    data = Scalar::all(1);
+    EXPECT_EQ(pow(ONE_SIZE, dims), cv::countNonZero(data));
+}
+
+INSTANTIATE_TEST_CASE_P(Core, CountNonZeroND,
+    testing::Combine(
+        testing::Range(2, 9),
+        testing::Values(CV_8U, CV_8S, CV_32F)
+    )
+);
