OpenCV: Sudoku Recognise

sudoku-recgonise.cpp

#include <iostream>
#include <opencv2/opencv.hpp>
using namespace std;
using namespace cv;

// #define SHOW_PROCESS

const int size = 303;
const int border = 3;
const int cut = 3;
const int areaThresh = 30;
const int sampleSize = 8;

bool polyToSquare (const Mat &src, Mat &dst, const vector<Point> &poly)
{
	int sum, minSum = INT_MAX, ptTopLeft, dif, maxDif = INT_MIN, ptTopRight;
	for(int i = 0; i < 4; i++) {
		sum = poly[i].x + poly[i].y;
		dif = poly[i].x - poly[i].y;
		if(sum < minSum) {
			minSum = sum;
			ptTopLeft = i;
		}
		if(dif > maxDif) {
			maxDif = dif;
			ptTopRight = i;
		}
	}
	if(ptTopLeft != (ptTopRight+1)%4 && ptTopRight != (ptTopLeft+1)%4)
		return false;

	Point2f srcPoints[4] = {
		Point2f(poly[ptTopLeft].x, poly[ptTopLeft].y),
		Point2f(poly[ptTopRight].x, poly[ptTopRight].y),
		Point2f(poly[(ptTopRight+2)%4].x, poly[(ptTopRight+2)%4].y),
		Point2f(poly[(ptTopLeft+2)%4].x, poly[(ptTopLeft+2)%4].y)
	},
	dstPoints[4] = {
		Point2f(0, 0),
		Point2f(dst.cols-1, 0),
		Point2f(0, dst.rows-1),
		Point2f(dst.cols-1, dst.rows-1)
	};

	Mat wrapMatrix = getPerspectiveTransform(srcPoints, dstPoints);
	warpPerspective(src, dst, wrapMatrix, dst.size());
	dst = dst(Rect(border, border, dst.cols-border*2, dst.rows-border*2));
	return true;
}

void getROI (const Mat& src, Mat& dst)
{
	int left, right, top, bottom;
	left = src.cols;
	right = 0;
	top = src.rows;
	bottom = 0;

	for(int i = 0; i < src.rows; i++)
	{
		for(int j = 0; j < src.cols; j++)
		{
			if(src.at<uchar>(i, j) > 0)
			{
				if(j < left) left = j;
				if(j > right) right = j;
				if(i < top) top = i;
				if(i > bottom) bottom = i;
			}
		}
	}

	int width = right - left;
	int height = bottom - top;
	int len = (width < height) ? height : width;
	dst = Mat::zeros(len, len, CV_8UC1);

	Rect dstRect((len - width)/2, (len - height)/2, width, height);
	Rect srcRect(left, top, width, height);
	Mat dstROI = dst(dstRect);
	Mat srcROI = src(srcRect);
	srcROI.copyTo(dstROI);
}

void removeSmallArea (Mat& src)
{
	Mat img;
	src.copyTo(img);

	vector<vector<Point>> contours;
	findContours(img, contours, CV_RETR_LIST, CV_CHAIN_APPROX_SIMPLE);

	Mat mask = Mat::zeros(src.size(), src.type());

	for(int i = 0; i != contours.size(); ++i) {
		if(fabs(contourArea(contours[i])) > areaThresh)
			drawContours(mask, contours, i, Scalar(255), CV_FILLED);
	}

	for(int i = 0; i < src.rows; i++)
	{
		for(int j = 0; j < src.cols; j++)
		{
			if(mask.at<uchar>(i, j) == 0)
			{
				src.at<uchar>(i, j) = 0;
			}
		}
	}
}

bool process(const string& file)
{
	Mat img, src;
	src = imread(file, CV_LOAD_IMAGE_GRAYSCALE);
	src.copyTo(img);

	adaptiveThreshold(img, img, 255, CV_ADAPTIVE_THRESH_MEAN_C, CV_THRESH_BINARY_INV, 25, 10);
	medianBlur(img, img, 5);

	vector<vector<Point>> contours;
	findContours(img, contours, CV_RETR_EXTERNAL, CV_CHAIN_APPROX_SIMPLE);
	double area, maxArea = 100;
	int maxIdx;
	for(int i = 0; i != contours.size(); ++i) {
		area = fabs(contourArea(contours[i]));
		if(area > maxArea)
		{
			maxIdx = i;
			maxArea = area;
		}
	}

	vector<Point> poly;
	approxPolyDP(contours[maxIdx], poly, arcLength(contours[maxIdx], true) * 0.02, true);
	if(poly.size() != 4 || fabs(contourArea(poly)) < 5000)
		return false;

#ifdef SHOW_PROCESS
	Mat tmp;
	src.copyTo(tmp);
	drawContours(tmp, contours, maxIdx, Scalar::all(0), 3);
	imshow("max contour", tmp);
	waitKey();
#endif

	Mat dst(size, size, src.type());
	if(!polyToSquare(src, dst, poly))
		return false;

	adaptiveThreshold(dst, dst, 255, CV_ADAPTIVE_THRESH_MEAN_C, CV_THRESH_BINARY_INV, 25, 10);
	
	int step = size/9;
	Mat cell, num, temp, data;
	temp = Mat::zeros(sampleSize, sampleSize, CV_8UC1);
	data = Mat::zeros(1, sampleSize*sampleSize, CV_32FC1);
	CvSVM svm = CvSVM(); 
	svm.load( "../SVM_DATA.xml" );
	int result[81];
	memset(result, 0, sizeof(result));

	for(int i = 0; i < 9; i++) {
		for(int j = 0; j < 9; j++) {
			cell = dst(Rect(j*step+cut, i*step+cut, step-cut*2, step-cut*2));
			removeSmallArea(cell);
			if(sum(cell)[0] == 0)
				continue;

			getROI(cell, num);
			resize(num, temp, temp.size());

			for(int i = 0; i < sampleSize; i++) {
				for(int j = 0; j < sampleSize; j++) {
					data.at<float>(0, i*sampleSize+j) = temp.at<uchar>(i, j);
				}
			}

			normalize(data, data);
			result[i*9+j] =  char(svm.predict(data)) - '0';
		}
		for(int k = 0; k < 9; k++) {
			cout << result[i*9+k] << "\t";
		}
		cout << endl;
	}

	return true;
}

int main( int argc, char** argv )
{
	process("../res/1.jpg");
	return 0;
}