opencv/samples/python/cv20squares.py

167 lines
5.2 KiB
Python
Raw Normal View History

"""
Find Squares in image by finding countours and filtering
"""
#Results slightly different from C version on same images, but is
#otherwise ok
import math
2011-07-12 14:56:03 +02:00
import cv2.cv as cv
def angle(pt1, pt2, pt0):
"calculate angle contained by 3 points(x, y)"
dx1 = pt1[0] - pt0[0]
dy1 = pt1[1] - pt0[1]
dx2 = pt2[0] - pt0[0]
dy2 = pt2[1] - pt0[1]
nom = dx1*dx2 + dy1*dy2
denom = math.sqrt( (dx1*dx1 + dy1*dy1) * (dx2*dx2 + dy2*dy2) + 1e-10 )
ang = nom / denom
return ang
def is_square(contour):
"""
Squareness checker
Square contours should:
-have 4 vertices after approximation,
-have relatively large area (to filter out noisy contours)
-be convex.
-have angles between sides close to 90deg (cos(ang) ~0 )
Note: absolute value of an area is used because area may be
positive or negative - in accordance with the contour orientation
"""
area = math.fabs( cv.ContourArea(contour) )
isconvex = cv.CheckContourConvexity(contour)
s = 0
if len(contour) == 4 and area > 1000 and isconvex:
for i in range(1, 4):
# find minimum angle between joint edges (maximum of cosine)
pt1 = contour[i]
pt2 = contour[i-1]
pt0 = contour[i-2]
t = math.fabs(angle(pt0, pt1, pt2))
if s <= t:s = t
# if cosines of all angles are small (all angles are ~90 degree)
# then its a square
if s < 0.3:return True
return False
def find_squares_from_binary( gray ):
"""
use contour search to find squares in binary image
returns list of numpy arrays containing 4 points
"""
squares = []
storage = cv.CreateMemStorage(0)
contours = cv.FindContours(gray, storage, cv.CV_RETR_TREE, cv.CV_CHAIN_APPROX_SIMPLE, (0,0))
storage = cv.CreateMemStorage(0)
while contours:
#approximate contour with accuracy proportional to the contour perimeter
arclength = cv.ArcLength(contours)
polygon = cv.ApproxPoly( contours, storage, cv.CV_POLY_APPROX_DP, arclength * 0.02, 0)
if is_square(polygon):
squares.append(polygon[0:4])
contours = contours.h_next()
return squares
def find_squares4(color_img):
"""
Finds multiple squares in image
Steps:
-Use Canny edge to highlight contours, and dilation to connect
the edge segments.
-Threshold the result to binary edge tokens
-Use cv.FindContours: returns a cv.CvSequence of cv.CvContours
-Filter each candidate: use Approx poly, keep only contours with 4 vertices,
enough area, and ~90deg angles.
Return all squares contours in one flat list of arrays, 4 x,y points each.
"""
#select even sizes only
width, height = (color_img.width & -2, color_img.height & -2 )
timg = cv.CloneImage( color_img ) # make a copy of input image
gray = cv.CreateImage( (width,height), 8, 1 )
# select the maximum ROI in the image
cv.SetImageROI( timg, (0, 0, width, height) )
# down-scale and upscale the image to filter out the noise
pyr = cv.CreateImage( (width/2, height/2), 8, 3 )
cv.PyrDown( timg, pyr, 7 )
cv.PyrUp( pyr, timg, 7 )
tgray = cv.CreateImage( (width,height), 8, 1 )
squares = []
# Find squares in every color plane of the image
# Two methods, we use both:
# 1. Canny to catch squares with gradient shading. Use upper threshold
# from slider, set the lower to 0 (which forces edges merging). Then
# dilate canny output to remove potential holes between edge segments.
# 2. Binary thresholding at multiple levels
N = 11
for c in [0, 1, 2]:
#extract the c-th color plane
cv.SetImageCOI( timg, c+1 );
cv.Copy( timg, tgray, None );
cv.Canny( tgray, gray, 0, 50, 5 )
cv.Dilate( gray, gray)
squares = squares + find_squares_from_binary( gray )
# Look for more squares at several threshold levels
for l in range(1, N):
cv.Threshold( tgray, gray, (l+1)*255/N, 255, cv.CV_THRESH_BINARY )
squares = squares + find_squares_from_binary( gray )
return squares
RED = (0,0,255)
GREEN = (0,255,0)
def draw_squares( color_img, squares ):
"""
Squares is py list containing 4-pt numpy arrays. Step through the list
and draw a polygon for each 4-group
"""
color, othercolor = RED, GREEN
for square in squares:
cv.PolyLine(color_img, [square], True, color, 3, cv.CV_AA, 0)
color, othercolor = othercolor, color
cv.ShowImage(WNDNAME, color_img)
WNDNAME = "Squares Demo"
def main():
"""Open test color images, create display window, start the search"""
cv.NamedWindow(WNDNAME, 1)
for name in [ "../c/pic%d.png" % i for i in [1, 2, 3, 4, 5, 6] ]:
img0 = cv.LoadImage(name, 1)
try:
img0
except ValueError:
print "Couldn't load %s\n" % name
continue
# slider deleted from C version, same here and use fixed Canny param=50
img = cv.CloneImage(img0)
cv.ShowImage(WNDNAME, img)
# force the image processing
draw_squares( img, find_squares4( img ) )
# wait for key.
if cv.WaitKey(-1) % 0x100 == 27:
break
if __name__ == "__main__":
main()