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