vcs.maemo.org Git - opencv/blob - samples/swig_python/squares.py

   1 #!/usr/bin/python
   2 #
   3 # The full "Square Detector" program.
   4 # It loads several images subsequentally and tries to find squares in
   5 # each image
   6 #
   7
   8 from opencv.cv import *
   9 from opencv.highgui import *
  10 from math import sqrt
  11
  12 thresh = 50;
  13 img = None;
  14 img0 = None;
  15 storage = None;
  16 wndname = "Square Detection Demo";
  17
  18 def angle( pt1, pt2, pt0 ):
  19     dx1 = pt1.x - pt0.x;
  20     dy1 = pt1.y - pt0.y;
  21     dx2 = pt2.x - pt0.x;
  22     dy2 = pt2.y - pt0.y;
  23     return (dx1*dx2 + dy1*dy2)/sqrt((dx1*dx1 + dy1*dy1)*(dx2*dx2 + dy2*dy2) + 1e-10);
  24
  25 def findSquares4( img, storage ):
  26     N = 11;
  27     sz = cvSize( img.width & -2, img.height & -2 );
  28     timg = cvCloneImage( img ); # make a copy of input image
  29     gray = cvCreateImage( sz, 8, 1 );
  30     pyr = cvCreateImage( cvSize(sz.width/2, sz.height/2), 8, 3 );
  31     # create empty sequence that will contain points -
  32     # 4 points per square (the square's vertices)
  33     squares = cvCreateSeq( 0, sizeof_CvSeq, sizeof_CvPoint, storage );
  34     squares = CvSeq_CvPoint.cast( squares )
  35
  36     # select the maximum ROI in the image
  37     # with the width and height divisible by 2
  38     subimage = cvGetSubRect( timg, cvRect( 0, 0, sz.width, sz.height ))
  39
  40     # down-scale and upscale the image to filter out the noise
  41     cvPyrDown( subimage, pyr, 7 );
  42     cvPyrUp( pyr, subimage, 7 );
  43     tgray = cvCreateImage( sz, 8, 1 );
  44     # find squares in every color plane of the image
  45     for c in range(3):
  46         # extract the c-th color plane
  47         channels = [None, None, None]
  48         channels[c] = tgray
  49         cvSplit( subimage, channels[0], channels[1], channels[2], None )
  50         for l in range(N):
  51             # hack: use Canny instead of zero threshold level.
  52             # Canny helps to catch squares with gradient shading
  53             if( l == 0 ):
  54                 # apply Canny. Take the upper threshold from slider
  55                 # and set the lower to 0 (which forces edges merging)
  56                 cvCanny( tgray, gray, 0, thresh, 5 );
  57                 # dilate canny output to remove potential
  58                 # holes between edge segments
  59                 cvDilate( gray, gray, None, 1 );
  60             else:
  61                 # apply threshold if l!=0:
  62                 #     tgray(x,y) = gray(x,y) < (l+1)*255/N ? 255 : 0
  63                 cvThreshold( tgray, gray, (l+1)*255/N, 255, CV_THRESH_BINARY );
  64
  65             # find contours and store them all as a list
  66             count, contours = cvFindContours( gray, storage, sizeof_CvContour,
  67                 CV_RETR_LIST, CV_CHAIN_APPROX_SIMPLE, cvPoint(0,0) );
  68
  69             if not contours:
  70                 continue
  71
  72             # test each contour
  73             for contour in contours.hrange():
  74                 # approximate contour with accuracy proportional
  75                 # to the contour perimeter
  76                 result = cvApproxPoly( contour, sizeof_CvContour, storage,
  77                     CV_POLY_APPROX_DP, cvContourPerimeter(contours)*0.02, 0 );
  78                 # square contours should have 4 vertices after approximation
  79                 # relatively large area (to filter out noisy contours)
  80                 # and be convex.
  81                 # Note: absolute value of an area is used because
  82                 # area may be positive or negative - in accordance with the
  83                 # contour orientation
  84                 if( result.total == 4 and
  85                     abs(cvContourArea(result)) > 1000 and
  86                     cvCheckContourConvexity(result) ):
  87                     s = 0;
  88                     for i in range(5):
  89                         # find minimum angle between joint
  90                         # edges (maximum of cosine)
  91                         if( i >= 2 ):
  92                             t = abs(angle( result[i], result[i-2], result[i-1]))
  93                             if s<t:
  94                                 s=t
  95                     # if cosines of all angles are small
  96                     # (all angles are ~90 degree) then write quandrange
  97                     # vertices to resultant sequence
  98                     if( s < 0.3 ):
  99                         for i in range(4):
 100                             squares.append( result[i] )
 101
 102     return squares;
 103
 104 # the function draws all the squares in the image
 105 def drawSquares( img, squares ):
 106     cpy = cvCloneImage( img );
 107     # read 4 sequence elements at a time (all vertices of a square)
 108     i=0
 109     while i<squares.total:
 110         pt = []
 111         # read 4 vertices
 112         pt.append( squares[i] )
 113         pt.append( squares[i+1] )
 114         pt.append( squares[i+2] )
 115         pt.append( squares[i+3] )
 116
 117         # draw the square as a closed polyline
 118         cvPolyLine( cpy, [pt], 1, CV_RGB(0,255,0), 3, CV_AA, 0 );
 119         i+=4
 120
 121     # show the resultant image
 122     cvShowImage( wndname, cpy );
 123
 124 def on_trackbar( a ):
 125     if( img ):
 126         drawSquares( img, findSquares4( img, storage ) );
 127
 128 names =  ["../c/pic1.png", "../c/pic2.png", "../c/pic3.png",
 129           "../c/pic4.png", "../c/pic5.png", "../c/pic6.png" ];
 130
 131 if __name__ == "__main__":
 132     # create memory storage that will contain all the dynamic data
 133     storage = cvCreateMemStorage(0);
 134     for name in names:
 135         img0 = cvLoadImage( name, 1 );
 136         if not img0:
 137             print "Couldn't load %s" % name
 138             continue;
 139         img = cvCloneImage( img0 );
 140         # create window and a trackbar (slider) with parent "image" and set callback
 141         # (the slider regulates upper threshold, passed to Canny edge detector)
 142         cvNamedWindow( wndname, 1 );
 143         cvCreateTrackbar( "canny thresh", wndname, thresh, 1000, on_trackbar );
 144         # force the image processing
 145         on_trackbar(0);
 146         # wait for key.
 147         # Also the function cvWaitKey takes care of event processing
 148         c = cvWaitKey(0);
 149         # clear memory storage - reset free space position
 150         cvClearMemStorage( storage );
 151         if( c == '\x1b' ):
 152             break;
 153     cvDestroyWindow( wndname );