1#!/usr/bin/env python
2
3'''
4Simple "Square Detector" program.
5
6Loads several images sequentially and tries to find squares in each image.
7'''
8
9# Python 2/3 compatibility
10import sys
11PY3 = sys.version_info[0] == 3
12
13if PY3:
14    xrange = range
15
16import numpy as np
17import cv2 as cv
18
19
20def angle_cos(p0, p1, p2):
21    d1, d2 = (p0-p1).astype('float'), (p2-p1).astype('float')
22    return abs( np.dot(d1, d2) / np.sqrt( np.dot(d1, d1)*np.dot(d2, d2) ) )
23
24def find_squares(img):
25    img = cv.GaussianBlur(img, (5, 5), 0)
26    squares = []
27    for gray in cv.split(img):
28        for thrs in xrange(0, 255, 26):
29            if thrs == 0:
30                bin = cv.Canny(gray, 0, 50, apertureSize=5)
31                bin = cv.dilate(bin, None)
32            else:
33                _retval, bin = cv.threshold(gray, thrs, 255, cv.THRESH_BINARY)
34            contours, _hierarchy = cv.findContours(bin, cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
35            for cnt in contours:
36                cnt_len = cv.arcLength(cnt, True)
37                cnt = cv.approxPolyDP(cnt, 0.02*cnt_len, True)
38                if len(cnt) == 4 and cv.contourArea(cnt) > 1000 and cv.isContourConvex(cnt):
39                    cnt = cnt.reshape(-1, 2)
40                    max_cos = np.max([angle_cos( cnt[i], cnt[(i+1) % 4], cnt[(i+2) % 4] ) for i in xrange(4)])
41                    if max_cos < 0.1 and filterSquares(squares, cnt):
42                        squares.append(cnt)
43
44    return squares
45
46def intersectionRate(s1, s2):
47    area, _intersection = cv.intersectConvexConvex(np.array(s1), np.array(s2))
48    return 2 * area / (cv.contourArea(np.array(s1)) + cv.contourArea(np.array(s2)))
49
50def filterSquares(squares, square):
51
52    for i in range(len(squares)):
53        if intersectionRate(squares[i], square) > 0.95:
54            return False
55
56    return True
57
58from tests_common import NewOpenCVTests
59
60class squares_test(NewOpenCVTests):
61
62    def test_squares(self):
63
64        img = self.get_sample('samples/data/pic1.png')
65        squares = find_squares(img)
66
67        testSquares = [
68        [[43, 25],
69        [43, 129],
70        [232, 129],
71        [232, 25]],
72
73        [[252, 87],
74        [324, 40],
75        [387, 137],
76        [315, 184]],
77
78        [[154, 178],
79        [196, 180],
80        [198, 278],
81        [154, 278]],
82
83        [[0, 0],
84        [400, 0],
85        [400, 300],
86        [0, 300]]
87        ]
88
89        matches_counter = 0
90        for i in range(len(squares)):
91            for j in range(len(testSquares)):
92                if intersectionRate(squares[i], testSquares[j]) > 0.9:
93                    matches_counter += 1
94
95        self.assertGreater(matches_counter / len(testSquares), 0.9)
96        self.assertLess( (len(squares) - matches_counter) / len(squares), 0.2)
97
98if __name__ == '__main__':
99    NewOpenCVTests.bootstrap()
100