Add shape-detect.cpp

1 files changed, 108 insertions, 0 deletions

diff --git a/shape-detect.cpp b/shape-detect.cpp
new file mode 100644
index 0000000..f373adf
--- /dev/null
+++ b/shape-detect.cpp
@@ -0,0 +1,108 @@
+/**
+ * Simple shape detector program.
+ * It loads an image and tries to find simple shapes (rectangle, triangle, circle, etc) in it.
+ * This program is a modified version of `squares.cpp` found in the OpenCV sample dir.
+ */
+#include <opencv2/highgui/highgui.hpp>
+#include <opencv2/imgproc/imgproc.hpp>
+#include <cmath>
+
+/**
+ * Helper function to find a cosine of angle between vectors
+ * from pt0->pt1 and pt0->pt2
+ */
+static double angle(cv::Point pt1, cv::Point pt2, cv::Point pt0)
+{
+	double dx1 = pt1.x - pt0.x;
+	double dy1 = pt1.y - pt0.y;
+	double dx2 = pt2.x - pt0.x;
+	double dy2 = pt2.y - pt0.y;
+	return (dx1*dx2 + dy1*dy2)/sqrt((dx1*dx1 + dy1*dy1)*(dx2*dx2 + dy2*dy2) + 1e-10);
+}
+
+/**
+ * Helper function to display text in the center of a contour
+ */
+void setLabel(cv::Mat& im, const std::string label, std::vector<cv::Point>& contour)
+{
+	int fontface = cv::FONT_HERSHEY_SIMPLEX;
+	double scale = 0.4;
+	int thickness = 1;
+	int baseline = 0;
+
+	cv::Size text = cv::getTextSize(label, fontface, scale, thickness, &baseline);
+	cv::Rect r = cv::boundingRect(contour);
+
+	cv::Point pt(r.x + ((r.width - text.width) / 2), r.y + ((r.height + text.height) / 2));
+	cv::rectangle(im, pt + cv::Point(0, baseline), pt + cv::Point(text.width, -text.height), CV_RGB(255,255,255), CV_FILLED);
+	cv::putText(im, label, pt, fontface, scale, CV_RGB(0,0,0), thickness, 8);
+}
+
+int main()
+{
+	cv::Mat src = cv::imread("basic-shapes-2.png");
+	if (src.empty())
+		return -1;
+
+	// Convert to grayscale
+	cv::Mat gray;
+	cv::cvtColor(src, gray, CV_BGR2GRAY);
+
+	// Use Canny instead of threshold to catch squares with gradient shading
+	cv::Mat bw;
+	cv::Canny(gray, bw, 0, 50, 5);
+
+	// Find contours
+	std::vector<std::vector<cv::Point> > contours;
+	cv::findContours(bw.clone(), contours, CV_RETR_EXTERNAL, CV_CHAIN_APPROX_SIMPLE);
+
+	std::vector<cv::Point> approx;
+	cv::Mat dst = src.clone();
+
+	for (int i = 0; i < contours.size(); i++)
+	{
+		// Approximate contour with accuracy proportional
+		// to the contour perimeter
+		cv::approxPolyDP(cv::Mat(contours[i]), approx, cv::arcLength(cv::Mat(contours[i]), true)*0.02, true);
+
+		// Skip small or non-convex objects 
+		if (std::fabs(cv::contourArea(contours[i])) < 100 || !cv::isContourConvex(approx))
+			continue;
+
+		if (approx.size() == 3)
+			setLabel(dst, "TRI", contours[i]);    // Triangles
+		else if (approx.size() == 5)
+			setLabel(dst, "PENTA", contours[i]);  // Pentagons
+		else if (approx.size() == 6)
+			setLabel(dst, "HEXA", contours[i]);   // Hexagons
+		else if (approx.size() == 4)
+		{
+			// Detect and label rectangles
+			double maxcos = 0;
+			for (int j = 2; j < 5; j++)
+			{
+				double cos = std::fabs(angle(approx[j%4], approx[j-2], approx[j-1]));
+				maxcos = std::max(maxcos, cos);
+			}
+			if (maxcos < 0.3)
+				setLabel(dst, "RECT", contours[i]);
+		}
+		else
+		{
+			// Detect and label circles
+			double area = cv::contourArea(contours[i]);
+			cv::Rect r = cv::boundingRect(contours[i]);
+			int radius = r.width / 2;
+
+			if (std::abs(1 - ((double)r.width / r.height)) <= 0.2 &&
+			    std::abs(1 - (area / (CV_PI * std::pow(radius, 2)))) <= 0.2)
+				setLabel(dst, "CIR", contours[i]);
+		}
+	}
+
+	cv::imshow("src", src);
+	cv::imshow("dst", dst);
+	cv::waitKey(0);
+	return 0;
+}
+