Detect square symbols in a diagram image in python using OpenCV

Question

I am trying to detect the square shaped symbols in a P&ID (a diagram) image file using OpenCV. I tried following tutorials that use contours, but that method doesn't seem to work with such diagram images. Using Hough Lines I am able to mark the vertical edges of these squares, but I am not sure how to use these edges detect the squares. All squares in an image have the same dimensions, but the dimensions might not be same across different images, hence template matching didn't work for me.

My code using Hough Lines:

import cv2 as cv
import numpy as np
import math

img = cv.imread('test_img.jpg')
img_gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
img_display = img.copy()
ret,thresh = cv.threshold(img_gray,250,255,cv.THRESH_BINARY)
image_inverted = cv.bitwise_not(thresh)

linesP = cv.HoughLinesP(image_inverted, 1, np.pi / 1, 50, None, 50, 2)

if linesP is not None:
    for i in range(0, len(linesP)):
        l = linesP[i][0]
        length = math.sqrt((l[2] - l[0])**2 + (l[3] - l[1])**2)
        if length < 100:
            cv.line(img_display, (l[0], l[1]), (l[2], l[3]), (0,0,255), 1, cv.LINE_AA)
            
cv.imwrite('img_display.png', img_display)

Input image:

Output Image:

In the above code, I've set it to detect only vertical lines because it wasn't detecting horizontal lines reliably.

Answer 1

If you know that the lines are horizontal or vertical you can filter them out by combining erode and dilate (the docs describe how it works).

After seperating horizontal and vertical lines, you can filter them by size. At the end you can fill all remaining closed contours and again use erode/delete to exract the larger shapes.

This is more reliable than using Hough Line Transform and gives you more control over what exactly is extracted.

Here is a demo:

import numpy as np
import cv2 as cv

min_length = 29
max_length = 150


# erode and dilate with rectangular kernel of given dimensions
def erode_dilate(image, dim):
    kernel = cv.getStructuringElement(cv.MORPH_RECT, dim)
    result = cv.erode(image, kernel)
    result = cv.dilate(result, kernel)
    return result


# get contours and filter by max_size
def filter_contours(image, max_size):
    contours, _ = cv.findContours(image, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
    dims = [(cnt, cv.boundingRect(cnt)) for cnt in contours]
    contours = [cnt for cnt, (x, y, w, h) in dims if w <= max_size and h <= max_size]
    return contours


# read image and get inverted threshold mask
img = cv.imread('test_img.jpg')
img_gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
ret, thres = cv.threshold(img_gray, 250, 255, cv.THRESH_BINARY_INV)

# extract horizontal lines
thres_h = erode_dilate(thres, (min_length, 1))
cv.imshow("horizontal", thres_h)

# extract vertical lines
thres_v = erode_dilate(thres, (1, min_length))
cv.imshow("vertical", thres_v)

# filter lines by max_length and draw them back to res
res = np.zeros_like(thres)
cntrs_h = filter_contours(thres_h, max_length)
cv.drawContours(res, cntrs_h, -1, 255, cv.FILLED)
cntrs_v = filter_contours(thres_v, max_length)
cv.drawContours(res, cntrs_v, -1, 255, cv.FILLED)
cv.imshow("filtered horizontal + vertical", res)

# fill remaining shapes
cntrs = filter_contours(res, max_length)
for c in cntrs:
    cv.drawContours(res, [c], -1, 255, cv.FILLED)
cv.imshow("filled", res)

# extract larger shapes
res = erode_dilate(res, (min_length, min_length))
cv.imshow("squares", res)

# draw contours of detected shapes on original image
cntrs = filter_contours(res, max_length)
cv.drawContours(img, cntrs, -1, (0, 0, 255), 2)

cv.imshow("output", img)
cv.waitKey(-1)

cv.destroyAllWindows()

Output: