Segment by floodfill, find contours, find polynomial contours, find quadrilaterals, find convex quadrilaterals, find the centroid of convex quadrilateras

floodfill.py

#!/usr/bin/env python

'''
Floodfill sample improved by Gary Bradski to find the contour, convert to polygon, test for qaudrilateral
make sure the quadrilateral is convex and find the center of the convex quadrilateral. Also, read in a directory of
images, not just one.

Usage:
  floodfill.py <Directory of images>

  Click on the image to set seed point

Keys:
  c     - toggle 4/8 connectivity
  f     - toggle floating range
  g     - toggle gaussian blur
  SPACE - go to next image
  ESC   - exit
'''

# Python 2/3 compatibility
from __future__ import print_function

import numpy as np
import cv2
from os import listdir
from os.path import isfile, join
import sys
from copy import deepcopy

class App():

    def __init__(self):
        cv2.namedWindow('floodfill', cv2.WINDOW_NORMAL)
        cv2.resizeWindow('floodfill', 1000, 1000)
        cv2.setMouseCallback('floodfill', self.onmouse)
        cv2.createTrackbar('lo', 'floodfill', 20, 255, self.update)
        cv2.createTrackbar('hi', 'floodfill', 20, 255, self.update)
        self.gaussian = False

    def update(self, dummy=None):
        if self.seed_pt is None:
            cv2.imshow('floodfill', self.img)
            return
        flooded = self.img.copy()
        self.mask[:] = 0
        lo = cv2.getTrackbarPos('lo', 'floodfill')
        hi = cv2.getTrackbarPos('hi', 'floodfill')
        # SEGMENT THE AREA VIA FLOODFILL
        flags = self.connectivity
        if self.fixed_range:
            flags |= cv2.FLOODFILL_FIXED_RANGE
        cv2.floodFill(flooded, self.mask, self.seed_pt, (255, 255, 255), (lo,) * 3, (hi,) * 3, flags)

        self.mask *= 0 #Reset the mask
        flags = self.connectivity
        flags |= cv2.FLOODFILL_MASK_ONLY | cv2.FLOODFILL_FIXED_RANGE
        flags |= (128<<8)
        cv2.floodFill(self.img, self.mask, self.seed_pt, (255, 255, 255), (lo,) * 3, (hi,) * 3, flags)

        cv2.imshow('mask',self.mask)
        cv2.imshow('i',self.img)

        cnts, _ = cv2.findContours(self.mask[1:-1,1:-1], cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        print("contors (len {})".format(len(cnts)))
        cnts = sorted(cnts, key=cv2.contourArea, reverse=True)[:10] # 3rd one is where we clicked

        # loop over our contours
        cnt = 0
        print("STARTING OVER CONTOURS")
        for c in cnts:
            peri = cv2.arcLength(c, True)

            if peri > 400 or peri < 40:
                print("{:3d} peri len = {} ... SKIPPING!".format(cnt,peri))
                continue
            print("{:3d} peri len = {}".format(cnt, peri))
            cnt += 1
            # APPROXIMATE AS A POLYNOMIAL
            approx = cv2.approxPolyDP(c, 0.02 * peri, True)
            # FIND QUADRILATERALS
            if len(approx) == 4:
                print("   Is a quadrilateral")
                # FIND THE CONVEX ONES
                hull = cv2.convexHull(approx, returnPoints = False)
                print("   Is a quadrilateral with len(hull)= {}".format(len(hull)))
                if len(hull) == 4:
                    print("      So: is also convex!")
                # FIND THE CENTER OF THE CONTOUR
                M = cv2.moments(approx)
                cX = int(M["m10"] / M["m00"])
                cY = int(M["m01"] / M["m00"])
                cv2.drawContours(flooded, [approx], -1, (0, 255, 0), cv2.FILLED) #1)
                cv2.circle(flooded,(cX,cY),3,(0,200,0),1)
        cv2.circle(flooded, self.seed_pt, 2, (0, 0, 255), -1)
        cv2.imshow('floodfill', flooded)

    def onmouse(self, event, x, y, flags, param):
        if event == cv2.EVENT_LBUTTONDOWN:
            self.seed_pt = x, y
            self.update()

    def run(self,fn):
        self.img = cv2.imread(cv2.samples.findFile(fn))
        if self.img is None:
            print('Failed to load image file:', fn)
            sys.exit(1)
        self.img_orig = deepcopy(self.img)
        self.img_gauss = cv2.GaussianBlur(self.img,(5,5),0)
        h, w = self.img.shape[:2]
        self.mask = np.zeros((h+2, w+2), np.uint8)
        self.seed_pt = None
        self.fixed_range = True
        self.connectivity = 4

        self.update()


        while True:
            ch = cv2.waitKey()
            if ch == ord(' '): # Next image
                print("==================================")
                break
            if ch == 27: # Exit program
                print('DONE!')
                break
            if ch == ord('f'): #float range
                self.fixed_range = not self.fixed_range
                print('using %s range' % ('floating', 'fixed')[self.fixed_range])
                self.update()
            if ch == ord('g'): #Gaussian
                self.gaussian = not self.gaussian
                print('gaussian = {}'.format(self.gaussian))
                if self.gaussian:
                    self.img = self.img_gauss
                else:
                    self.img = self.img_orig
                self.update()
            if ch == ord('c'): # 4 or 8 connectivity
                self.connectivity = 12-self.connectivity
                print('connectivity =', self.connectivity)
                self.update()

        return ch


if __name__ == '__main__':
    # Window
    from_path = sys.argv[1]
    img_files = [f for f in listdir(from_path) if isfile(join(from_path, f))]
    cnt = 0
    for f in img_files:
        print("{:3d} For file {}:".format(cnt,f))
        cnt += 1
        print(__doc__)
        k = App().run(join(from_path,f))
        #k = cv.waitKey() & 0xFF
        if k == 27 or k == ord('q'):
            break
    cv2.destroyAllWindows()