高斯混合模型(GMM)分割数据
#include<opencv2/opencv.hpp>
#include<iostream>
using namespace std;
using namespace cv;
using namespace cv::ml;
int main()
{
Mat img = Mat::zeros(500, 500, CV_8UC3);
RNG rng(1234);
Scalar colorTab[] = {
Scalar(0, 0, 255),
Scalar(0, 255, 0),
Scalar(255, 0, 0),
Scalar(0, 255, 255),
Scalar(255, 0, 255)
};
int numCluster = 4;
int numSample = rng.uniform(5, 1000);
Mat points(numSample, 2, CV_32FC1);
for (int k = 0; k < numCluster; k++)
{
Point center;
center.x = rng.uniform(0, img.cols);
center.y = rng.uniform(0, img.rows);
Mat pointChunk = points.rowRange(k * numSample / numCluster,
k == numCluster - 1 ? numSample:(k + 1) * numSample / numCluster );
rng.fill(pointChunk, RNG::NORMAL, Scalar(center.x, center.y), Scalar(img.cols * 0.05, img.rows * 0.05));
}
randShuffle(points, 1, &rng);
Mat labels;
Ptr<ml::EM> em_model = ml::EM::create();
em_model->setClustersNumber(numCluster);
em_model->setCovarianceMatrixType(ml::EM::COV_MAT_SPHERICAL);
em_model->setTermCriteria(TermCriteria(TermCriteria::EPS +
TermCriteria::COUNT, 100, 0.1));
em_model->trainEM(points, noArray(), labels, noArray());
Mat sample(1, 2, CV_32FC1);
for (int row = 0; row < img.rows; row++)
{
for (int col = 0; col < img.cols; col++)
{
sample.at<float>(0) = (float)col;
sample.at<float>(1) = (float)row;
Vec2d predict = em_model->predict2(sample, noArray());
int response = cvRound(predict[1]);
circle(img, Point(col, row), 1, colorTab[response] * 0.5, -1);
}
}
for (int i = 0; i < numSample; i++)
{
Point p(cvRound(points.at<float>(i, 0)), cvRound(points.at<float>(i, 1)));
circle(img, p, 1, colorTab[labels.at<int>(i)], -1);
}
imshow("GMM-EM", img);
waitKey(0);
return 0;
}