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 9import numpy as np 10import cv2 11 12 13def angle_cos(p0, p1, p2): 14 d1, d2 = (p0-p1).astype('float'), (p2-p1).astype('float') 15 return abs( np.dot(d1, d2) / np.sqrt( np.dot(d1, d1)*np.dot(d2, d2) ) ) 16 17def find_squares(img): 18 img = cv2.GaussianBlur(img, (5, 5), 0) 19 squares = [] 20 for gray in cv2.split(img): 21 for thrs in xrange(0, 255, 26): 22 if thrs == 0: 23 bin = cv2.Canny(gray, 0, 50, apertureSize=5) 24 bin = cv2.dilate(bin, None) 25 else: 26 retval, bin = cv2.threshold(gray, thrs, 255, cv2.THRESH_BINARY) 27 bin, contours, hierarchy = cv2.findContours(bin, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) 28 for cnt in contours: 29 cnt_len = cv2.arcLength(cnt, True) 30 cnt = cv2.approxPolyDP(cnt, 0.02*cnt_len, True) 31 if len(cnt) == 4 and cv2.contourArea(cnt) > 1000 and cv2.isContourConvex(cnt): 32 cnt = cnt.reshape(-1, 2) 33 max_cos = np.max([angle_cos( cnt[i], cnt[(i+1) % 4], cnt[(i+2) % 4] ) for i in xrange(4)]) 34 if max_cos < 0.1: 35 squares.append(cnt) 36 return squares 37 38if __name__ == '__main__': 39 from glob import glob 40 for fn in glob('../data/pic*.png'): 41 img = cv2.imread(fn) 42 squares = find_squares(img) 43 cv2.drawContours( img, squares, -1, (0, 255, 0), 3 ) 44 cv2.imshow('squares', img) 45 ch = 0xFF & cv2.waitKey() 46 if ch == 27: 47 break 48 cv2.destroyAllWindows() 49