Skip to content

Instantly share code, notes, and snippets.

@daviddoria
Created December 10, 2016 21:08
Show Gist options
  • Save daviddoria/943b82f1877f4ed3541544c48c22926d to your computer and use it in GitHub Desktop.
Save daviddoria/943b82f1877f4ed3541544c48c22926d to your computer and use it in GitHub Desktop.
#include <iostream>
#include <opencv2/ml/ml.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc.hpp>
int main(int, char**)
{
// Data for visual representation
int width = 512, height = 512;
cv::Mat image = cv::Mat::zeros(height, width, CV_8UC3);
// Set up training data
const size_t numberOfSamples = 4;
int labels[numberOfSamples] = { 1, -1, -1, -1 };
cv::Mat labelsMat(numberOfSamples, 1, CV_32SC1, labels);
float trainingData[numberOfSamples][2] = { { 501, 10 }, { 255, 10 }, { 501, 255 }, { 10, 501 } };
cv::Mat trainingDataMat(numberOfSamples, 2, CV_32FC1, trainingData);
// Set up SVM's parameters
cv::Ptr<cv::ml::SVM> svm = cv::ml::SVM::create();
svm->setType(cv::ml::SVM::C_SVC);
svm->setKernel(cv::ml::SVM::LINEAR);
svm->setTermCriteria(cv::TermCriteria(cv::TermCriteria::MAX_ITER, 100, 1e-6));
// Train the SVM with given parameters
cv::Ptr<cv::ml::TrainData> td = cv::ml::TrainData::create(trainingDataMat, cv::ml::ROW_SAMPLE, labelsMat);
svm->train(td);
// Or train the SVM with optimal parameters
//svm->trainAuto(td);
cv::Vec3b green(0, 255, 0), blue(255, 0, 0);
// Show the decision regions given by the SVM
for (int i = 0; i < image.rows; ++i) {
for (int j = 0; j < image.cols; ++j) {
cv::Mat sampleMat = (cv::Mat_<float>(1, 2) << j, i);
float response = svm->predict(sampleMat);
if (response == 1) {
image.at<cv::Vec3b>(i, j) = green;
}
else if (response == -1) {
image.at<cv::Vec3b>(i, j) = blue;
}
}
}
// Show the training data
int thickness = -1;
int lineType = 8;
for(size_t sampleID = 0; sampleID < numberOfSamples; ++sampleID) {
cv::Scalar color;
if(labels[sampleID] == 1) {
color = cv::Scalar(0,0,0);
}
else {
color = cv::Scalar(255,255,255);
}
cv::circle(image, cv::Point(trainingData[sampleID][0], trainingData[sampleID][1]), 5, color, thickness, lineType);
}
// Show support vectors
thickness = 2;
lineType = 8;
cv::Mat sv = svm->getSupportVectors();
for (int i = 0; i < sv.rows; ++i) {
const float* v = sv.ptr<float>(i);
circle(image, cv::Point((int)v[0], (int)v[1]), 6, cv::Scalar(128, 128, 128), thickness, lineType);
}
cv::imwrite("result.png", image); // save the image
cv::imshow("SVM Simple Example", image); // show it to the user
cv::waitKey(0);
return EXIT_SUCCESS;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment