opencv::KMeans图像分割

#include <opencv2/opencv.hpp>
#include <iostream>

using namespace cv;
using namespace std;

int main(int argc, char** argv) {
    Mat src = imread("D:/vcprojects/images/toux.jpg");
    if (src.empty()) {
        printf("could not load image...
");
        return -1;
    }
    namedWindow("input image", CV_WINDOW_AUTOSIZE);
    imshow("input image", src);

    Scalar colorTab[] = {
        Scalar(0, 0, 255),
        Scalar(0, 255, 0),
        Scalar(255, 0, 0),
        Scalar(0, 255, 255),
        Scalar(255, 0, 255)
    };

    int width = src.cols;
    int height = src.rows;
    int dims = src.channels();

    // 像素点个数
    int sampleCount = width*height;
    int clusterCount = 4;
    //将数据装载到一行
    Mat points(sampleCount, dims, CV_32F, Scalar(10));
    Mat labels;
    Mat centers(clusterCount, 1, points.type());

    // RGB 数据转换到样本数据
    int index = 0;
    for (int row = 0; row < height; row++) {
        for (int col = 0; col < width; col++) {
            index = row*width + col;
            Vec3b bgr = src.at<Vec3b>(row, col);
            points.at<float>(index, 0) = static_cast<int>(bgr[0]);
            points.at<float>(index, 1) = static_cast<int>(bgr[1]);
            points.at<float>(index, 2) = static_cast<int>(bgr[2]);
        }
    }

    // 运行K-Means
    TermCriteria criteria = TermCriteria(TermCriteria::EPS + TermCriteria::COUNT, 10, 0.1);
    kmeans(points, clusterCount, labels, criteria, 3, KMEANS_PP_CENTERS, centers);

    // 显示图像分割结果
    Mat result = Mat::zeros(src.size(), src.type());
    for (int row = 0; row < height; row++) {
        for (int col = 0; col < width; col++) {
            index = row*width + col;
            int label = labels.at<int>(index, 0);
            result.at<Vec3b>(row, col)[0] = colorTab[label][0];
            result.at<Vec3b>(row, col)[1] = colorTab[label][1];
            result.at<Vec3b>(row, col)[2] = colorTab[label][2];
        }
    }

    for (int i = 0; i < centers.rows; i++) {
        int x = centers.at<float>(i, 0);
        int y = centers.at<float>(i, 1);
        printf("center %d = c.x : %d, c.y : %d
", i, x, y);
    }
    
    imshow("KMeans Image Segmentation Demo", result);
    waitKey(0);
    return 0;
}