Recognising simple shapes with OpenCV in Python

Convolutional neural networks are usually used to recognise objects in images but that's a bit of an overkill if you just want to recognise simple flat shapes. Here's how to use OpenCV to do simple object recognition.

Basically the trick to recognise polygons is to convert your image into an approximate polygon representation using something like edge detection and then count the number of sides in the polygon. OpenCV handles a lot of this stuff for you so its quite easy.

Here's the image we will be working on:


The first step is to binarise the pixels of the image, that is they are made either black or white. First the image must be turned into greyscale so that there is just one number per pixel and then anything that is not white (greyscale value 255) is replaced with 255 (white) whilst pure white is replaced with 0 (black).

import cv2
import numpy as np

img = cv2.imread('polygons.png', cv2.IMREAD_COLOR)

#Make image greyscale
grey = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
cv2.imshow('greyscale', grey)

#Binarise greyscale pixels
(_, thresh) = cv2.threshold(grey, 254, 255, 1)
cv2.imshow('thresholded', thresh)

Next we're going to find contours around these white shapes, that is, reduce the image into a set of points with each point being a corner in the shape. This is done using the findContours function which returns a list consisting of contour points for every shape and their hierarchy. The hierarchy is the way each shape is nested within other shapes. We won't need this here since we don't have any shapes within shapes but it will be useful in other situations.

#Get shape contours
(contours, hierarchy) = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
thresh_copy = cv2.cvtColor(thresh, cv2.COLOR_GRAY2RGB)
cv2.drawContours(thresh_copy, contours, contourIdx=-1, color=(0, 255, 0), thickness=2)
print('number of sides per shape:')
for contour in contours:
    print('', contour.shape[0])
print()
cv2.imshow('contours', thresh_copy)

number of sides per shape:
 119
 122
 4
 124

Unfortunately, the number of sides extracted from each shape is nowhere near what it should be. This is because the corners in the shapes are rounded which results in multiple sides around the corner. We therefore need to simplify these contours so that insignificant sides can be removed. This is done using the Ramer–Douglas–Peucker algorithm which is implemented as the approxPolyDP function. The algorithm requires a parameter called epsilon which controls how roughly to chop up the polygon into a smaller number of sides. This number is dependent on the size of the polygon so we make it a fraction of the shape's perimeter.

#Simplify contours
thresh_copy = cv2.cvtColor(thresh, cv2.COLOR_GRAY2RGB)
print('number of sides per shape:')
for contour in contours:
    perimeter = cv2.arcLength(contour, True)
    e = 0.05*perimeter #The bigger the fraction, the more sides are chopped off the original polygon
    contour = cv2.approxPolyDP(contour, epsilon=e, closed=True)
    cv2.drawContours(thresh_copy, [ contour ], contourIdx=-1, color=(0, 255, 0), thickness=2)
    print('', contour.shape[0])
cv2.imshow('simplified contours', thresh_copy)

number of sides per shape:
 6
 5
 4
 3

And with this we have the number of sides in each polygon.

If you want to check for circles as well as polygons then you will not be able to do so by counting sides. Instead you can get the minimum enclosing circle around a contour and check if its area is close to the area of the contour (before it is simplified):

((centre_x, centre_y), radius) = cv2.minEnclosingCircle(contour)
if cv2.contourArea(contour)/(np.pi*radius**2) > 0.95:
    print('circle')

The minimum enclosing circle is the smallest circle that completely contains the contour. Therefore it's area must necessarily be larger or equal to the shape of the contour, which can only be equal in the case that the contour is a circle.

This is also one way how you can get the position of each shape, by getting the centre point of the enclosing circle. The proper way to get a single coordinate representing the position of the contour is to get the centre of gravity using the contour's moments:

m = cv2.moments(contour)
x = m['m10']/m['m00']
y = m['m01']/m['m00']

---

Here's the full code:

import cv2
import numpy as np

img = cv2.imread('polygons.png', cv2.IMREAD_COLOR)

#Binarise pixels
grey = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
(_, thresh) = cv2.threshold(grey, 254, 255, 1)

#Get shape contours
(contours, hierarchy) = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

#Recognise shapes
for contour in contours:
    m = cv2.moments(contour)
    x = round(m['m10']/m['m00'])
    y = round(m['m01']/m['m00'])

    shape_name = None

    (_, radius) = cv2.minEnclosingCircle(contour)
    if cv2.contourArea(contour)/(np.pi*radius**2) > 0.95:
        shape_name = 'circle'
    else:
        e = 0.05*cv2.arcLength(contour, True)
        simple_contour = cv2.approxPolyDP(contour, epsilon=e, closed=True)
        num_sides = simple_contour.shape[0]
        shape_name = { 3: 'triangle', 4: 'quad', 5: 'pentagon', 6: 'hexagon' }.get(num_sides, 'unknown')

    cv2.putText(img, shape_name, (x, y), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.6, color=(0, 255, 0), thickness=2)

cv2.imshow('named shapes', img)