Merge pull request #12242 from alalek:fix_12236

modules/ml/include/opencv2/ml.hpp

@@ -239,7 +239,18 @@ public:
     /** @brief Returns vector of symbolic names captured in loadFromCSV() */
     CV_WRAP void getNames(std::vector<String>& names) const;
 
-    CV_WRAP static Mat getSubVector(const Mat& vec, const Mat& idx);
+    /** @brief Extract from 1D vector elements specified by passed indexes.
+    @param vec input vector (supported types: CV_32S, CV_32F, CV_64F)
+    @param idx 1D index vector
+     */
+    static CV_WRAP Mat getSubVector(const Mat& vec, const Mat& idx);
+
+    /** @brief Extract from matrix rows/cols specified by passed indexes.
+    @param matrix input matrix (supported types: CV_32S, CV_32F, CV_64F)
+    @param idx 1D index vector
+    @param layout specifies to extract rows (cv::ml::ROW_SAMPLES) or to extract columns (cv::ml::COL_SAMPLES)
+     */
+    static CV_WRAP Mat getSubMatrix(const Mat& matrix, const Mat& idx, int layout);
 
     /** @brief Reads the dataset from a .csv file and returns the ready-to-use training data.
 

modules/ml/src/data.cpp

@@ -43,6 +43,8 @@
 #include <algorithm>
 #include <iterator>
 
+#include <opencv2/core/utils/logger.hpp>
+
 namespace cv { namespace ml {
 
 static const float MISSED_VAL = TrainData::missingValue();
@@ -54,69 +56,65 @@ Mat TrainData::getTestSamples() const
 {
     Mat idx = getTestSampleIdx();
     Mat samples = getSamples();
-    return idx.empty() ? Mat() : getSubVector(samples, idx);
+    return idx.empty() ? Mat() : getSubMatrix(samples, idx, getLayout());
 }
 
 Mat TrainData::getSubVector(const Mat& vec, const Mat& idx)
 {
-    if( idx.empty() )
-        return vec;
-    int i, j, n = idx.checkVector(1, CV_32S);
-    int type = vec.type();
-    CV_Assert( type == CV_32S || type == CV_32F || type == CV_64F );
-    int dims = 1, m;
-
-    if( vec.cols == 1 || vec.rows == 1 )
+    if (!(vec.cols == 1 || vec.rows == 1))
+        CV_LOG_WARNING(NULL, "'getSubVector(const Mat& vec, const Mat& idx)' call with non-1D input is deprecated. It is not designed to work with 2D matrixes (especially with 'cv::ml::COL_SAMPLE' layout).");
+    return getSubMatrix(vec, idx, vec.rows == 1 ? cv::ml::COL_SAMPLE : cv::ml::ROW_SAMPLE);
+}
+
+template<typename T>
+Mat getSubMatrixImpl(const Mat& m, const Mat& idx, int layout)
+{
+    int nidx = idx.checkVector(1, CV_32S);
+    int dims = m.cols, nsamples = m.rows;
+
+    Mat subm;
+    if (layout == COL_SAMPLE)
     {
-        dims = 1;
-        m = vec.cols + vec.rows - 1;
+        std::swap(dims, nsamples);
+        subm.create(dims, nidx, m.type());
     }
     else
     {
-        dims = vec.cols;
-        m = vec.rows;
+        subm.create(nidx, dims, m.type());
     }
 
-    Mat subvec;
-
-    if( vec.cols == m )
-        subvec.create(dims, n, type);
-    else
-        subvec.create(n, dims, type);
-    if( type == CV_32S )
-        for( i = 0; i < n; i++ )
+    for (int i = 0; i < nidx; i++)
+    {
+        int k = idx.at<int>(i); CV_CheckGE(k, 0, "Bad idx"); CV_CheckLT(k, nsamples, "Bad idx or layout");
+        if (dims == 1)
         {
-            int k = idx.at<int>(i);
-            CV_Assert( 0 <= k && k < m );
-            if( dims == 1 )
-                subvec.at<int>(i) = vec.at<int>(k);
-            else
-                for( j = 0; j < dims; j++ )
-                    subvec.at<int>(i, j) = vec.at<int>(k, j);
+            subm.at<T>(i) = m.at<T>(k);  // at() has "transparent" access for 1D col-based / row-based vectors.
         }
-    else if( type == CV_32F )
-        for( i = 0; i < n; i++ )
+        else if (layout == COL_SAMPLE)
         {
-            int k = idx.at<int>(i);
-            CV_Assert( 0 <= k && k < m );
-            if( dims == 1 )
-                subvec.at<float>(i) = vec.at<float>(k);
-            else
-                for( j = 0; j < dims; j++ )
-                    subvec.at<float>(i, j) = vec.at<float>(k, j);
+            for (int j = 0; j < dims; j++)
+                subm.at<T>(j, i) = m.at<T>(j, k);
         }
-    else
-        for( i = 0; i < n; i++ )
+        else
         {
-            int k = idx.at<int>(i);
-            CV_Assert( 0 <= k && k < m );
-            if( dims == 1 )
-                subvec.at<double>(i) = vec.at<double>(k);
-            else
-                for( j = 0; j < dims; j++ )
-                    subvec.at<double>(i, j) = vec.at<double>(k, j);
+            for (int j = 0; j < dims; j++)
+                subm.at<T>(i, j) = m.at<T>(k, j);
         }
-    return subvec;
+    }
+    return subm;
+}
+
+Mat TrainData::getSubMatrix(const Mat& m, const Mat& idx, int layout)
+{
+    if (idx.empty())
+        return m;
+    int type = m.type();
+    CV_CheckType(type, type == CV_32S || type == CV_32F || type == CV_64F, "");
+    if (type == CV_32S || type == CV_32F)  // 32-bit
+        return getSubMatrixImpl<int>(m, idx, layout);
+    if (type == CV_64F)  // 64-bit
+        return getSubMatrixImpl<double>(m, idx, layout);
+    CV_Error(Error::StsInternal, "");
 }
 
 class TrainDataImpl CV_FINAL : public TrainData
@@ -172,30 +170,30 @@ public:
     }
     Mat getTrainSampleWeights() const CV_OVERRIDE
     {
-        return getSubVector(sampleWeights, getTrainSampleIdx());
+        return getSubVector(sampleWeights, getTrainSampleIdx());  // 1D-vector
     }
     Mat getTestSampleWeights() const CV_OVERRIDE
     {
         Mat idx = getTestSampleIdx();
-        return idx.empty() ? Mat() : getSubVector(sampleWeights, idx);
+        return idx.empty() ? Mat() : getSubVector(sampleWeights, idx);  // 1D-vector
     }
     Mat getTrainResponses() const CV_OVERRIDE
     {
-        return getSubVector(responses, getTrainSampleIdx());
+        return getSubMatrix(responses, getTrainSampleIdx(), cv::ml::ROW_SAMPLE);  // col-based responses are transposed in setData()
     }
     Mat getTrainNormCatResponses() const CV_OVERRIDE
     {
-        return getSubVector(normCatResponses, getTrainSampleIdx());
+        return getSubMatrix(normCatResponses, getTrainSampleIdx(), cv::ml::ROW_SAMPLE);  // like 'responses'
     }
     Mat getTestResponses() const CV_OVERRIDE
     {
         Mat idx = getTestSampleIdx();
-        return idx.empty() ? Mat() : getSubVector(responses, idx);
+        return idx.empty() ? Mat() : getSubMatrix(responses, idx, cv::ml::ROW_SAMPLE);  // col-based responses are transposed in setData()
     }
     Mat getTestNormCatResponses() const CV_OVERRIDE
     {
         Mat idx = getTestSampleIdx();
-        return idx.empty() ? Mat() : getSubVector(normCatResponses, idx);
+        return idx.empty() ? Mat() : getSubMatrix(normCatResponses, idx, cv::ml::ROW_SAMPLE);  // like 'responses'
     }
     Mat getNormCatResponses() const CV_OVERRIDE { return normCatResponses; }
     Mat getClassLabels() const CV_OVERRIDE { return classLabels; }

modules/ml/test/test_mltests2.cpp

@@ -721,5 +721,68 @@ void CV_MLBaseTest::load( const char* filename )
         CV_Error( CV_StsNotImplemented, "invalid stat model name");
 }
 
+
+
+TEST(TrainDataGet, layout_ROW_SAMPLE)  // Details: #12236
+{
+    cv::Mat test = cv::Mat::ones(150, 30, CV_32FC1) * 2;
+    test.col(3) += Scalar::all(3);
+    cv::Mat labels = cv::Mat::ones(150, 3, CV_32SC1) * 5;
+    labels.col(1) += 1;
+    cv::Ptr<cv::ml::TrainData> train_data = cv::ml::TrainData::create(test, cv::ml::ROW_SAMPLE, labels);
+    train_data->setTrainTestSplitRatio(0.9);
+
+    Mat tidx = train_data->getTestSampleIdx();
+    EXPECT_EQ((size_t)15, tidx.total());
+
+    Mat tresp = train_data->getTestResponses();
+    EXPECT_EQ(15, tresp.rows);
+    EXPECT_EQ(labels.cols, tresp.cols);
+    EXPECT_EQ(5, tresp.at<int>(0, 0)) << tresp;
+    EXPECT_EQ(6, tresp.at<int>(0, 1)) << tresp;
+    EXPECT_EQ(6, tresp.at<int>(14, 1)) << tresp;
+    EXPECT_EQ(5, tresp.at<int>(14, 2)) << tresp;
+
+    Mat tsamples = train_data->getTestSamples();
+    EXPECT_EQ(15, tsamples.rows);
+    EXPECT_EQ(test.cols, tsamples.cols);
+    EXPECT_EQ(2, tsamples.at<float>(0, 0)) << tsamples;
+    EXPECT_EQ(5, tsamples.at<float>(0, 3)) << tsamples;
+    EXPECT_EQ(2, tsamples.at<float>(14, test.cols - 1)) << tsamples;
+    EXPECT_EQ(5, tsamples.at<float>(14, 3)) << tsamples;
+}
+
+TEST(TrainDataGet, layout_COL_SAMPLE)  // Details: #12236
+{
+    cv::Mat test = cv::Mat::ones(30, 150, CV_32FC1) * 3;
+    test.row(3) += Scalar::all(3);
+    cv::Mat labels = cv::Mat::ones(3, 150, CV_32SC1) * 5;
+    labels.row(1) += 1;
+    cv::Ptr<cv::ml::TrainData> train_data = cv::ml::TrainData::create(test, cv::ml::COL_SAMPLE, labels);
+    train_data->setTrainTestSplitRatio(0.9);
+
+    Mat tidx = train_data->getTestSampleIdx();
+    EXPECT_EQ((size_t)15, tidx.total());
+
+    Mat tresp = train_data->getTestResponses();  // always row-based, transposed
+    EXPECT_EQ(15, tresp.rows);
+    EXPECT_EQ(labels.rows, tresp.cols);
+    EXPECT_EQ(5, tresp.at<int>(0, 0)) << tresp;
+    EXPECT_EQ(6, tresp.at<int>(0, 1)) << tresp;
+    EXPECT_EQ(6, tresp.at<int>(14, 1)) << tresp;
+    EXPECT_EQ(5, tresp.at<int>(14, 2)) << tresp;
+
+
+    Mat tsamples = train_data->getTestSamples();
+    EXPECT_EQ(15, tsamples.cols);
+    EXPECT_EQ(test.rows, tsamples.rows);
+    EXPECT_EQ(3, tsamples.at<float>(0, 0)) << tsamples;
+    EXPECT_EQ(6, tsamples.at<float>(3, 0)) << tsamples;
+    EXPECT_EQ(6, tsamples.at<float>(3, 14)) << tsamples;
+    EXPECT_EQ(3, tsamples.at<float>(test.rows - 1, 14)) << tsamples;
+}
+
+
+
 } // namespace
 /* End of file. */