Traffic Sign Segmentation using K-means Clustering OpenCV

Question

I used K-Means Clustering to perform segmentation on this traffic sign as shown below.

These are my code

Read image and blur

img = cv.imread('000_0001.png')
img_rgb = cv.cvtColor(img, cv.COLOR_BGR2RGB)
kernel_size = 5
img_rgb = cv.blur(img_rgb, (kernel_size, kernel_size))

# reshape 
img_reshape = img_rgb.reshape((-1, 3))
img_reshape = np.float32(img_reshape)

Perform k-means clustering

criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 10, 1.0)
K = 4
attempts = 10
ret, label, center = cv.kmeans(img_reshape, K, None, criteria, attempts, cv.KMEANS_PP_CENTERS)

# reshape into original dimensions
center = np.uint8(center)
res = center[label.flatten()]
result = res.reshape(img_rgb.shape)

plt.figure(figsize = (15, 15))
plt.subplot(1, 2, 1)
plt.imshow(img_rgb)
plt.subplot(1, 2, 2)
plt.imshow(result)
plt.show()

create masks

from numpy import linalg as LN
red = (255, 0, 0)
Idx_min_red = np.argmin(LN.norm(center-red, axis = 1))

white = (255, 255, 255)
Idx_min_white = np.argmin(LN.norm(center-white, axis = 1))

black = (0, 0, 0)
Idx_min_black = np.argmin(LN.norm(center-black, axis = 1))

mask_red = result == center[Idx_min_red]

mask_white = result == center[Idx_min_white]

mask_black = result == center[Idx_min_black]

pre_mask = cv.bitwise_or(np.float32(mask_red), np.float32(mask_white))
mask = cv.bitwise_or(np.float32(pre_mask), np.float32(mask_black))

Segment the image

seg_img = img*(mask.astype("uint8"))

Morphological Transformation

kernel = np.ones((5, 5), dtype = np.uint8)

img_dilate = cv.morphologyEx(seg_img, cv.MORPH_OPEN, kernel)

res = np.hstack((img, img_dilate))

cv.imshow("res", res)
cv.waitKey(0)
cv.destroyAllWindows()

Question here These codes can only segment red traffic signs, is it possible to tweak a little bit on different colours so that it can segment red, blue and yellow traffic signs? (like the one below for example)

Update, this is what I have tried: I used a pipeline to do OR operation on all the masks mask = mask_red | mask_white | mask_black | mask_blue

but then the new mask will fail to segment the image

user_3104151 · Answer

Rather than K means, I'd suggest a connected components based approach to find this. These signs and symbols have relatively uniform color areas that take a fairly significant part of the image, that would result in large connected components. You can later write downstream logic to select the relevant connected components to define your sign and create a segmented image from them.

import cv2
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

colors = (np.array(plt.cm.tab10.colors) * 255).astype('int')

im1 = cv2.imread('37Gu6.png')
im2 = cv2.imread('dU47J.png')

for i in [im1, im2]:
    gray    = cv2.cvtColor(i, cv2.COLOR_BGR2GRAY)
    
    # Calculate image value histogram to get an adaptive threshold value based
    hist    = np.histogram(gray.ravel(), bins=25)
    modeIdx = np.where(hist[0] == hist[0].max())[0][0]
    modeVal = hist[1][modeIdx]
    
    tVal    = modeVal * 1.15 # Our threshold is 115% of the most common hist value
    _, thresh = cv2.threshold(gray, tVal, 255, cv2.THRESH_BINARY)

    # Find connected componenents
    output = cv2.connectedComponentsWithStats(thresh, 8, cv2.CV_32S)
    (numLabels, labels, stats, centroids) = output

    # The 5th column in the CC stats is the area
    # Find the indices of CCs with above average area
    CCArea = stats[:,4]
    aboveAvgLbls = np.where(CCArea > CCArea.mean())[0]
    
    # Generate labeled output image
    outputImg = np.zeros(i.shape, dtype="uint8")

    for l, c in zip(aboveAvgLbls, colors):
        w = np.where(labels == l)
        outputImg[ w[0], w[1] ] = c
        
    f, ax = plt.subplots(1,4,figsize=(20,5))
    for a, img, t in zip(ax, [i, gray, thresh, outputImg], ['RGB', 'Gray', f'Thresh [t={tVal}]', 'Labels']):
        a.imshow(img[:,:,::-1] if img.shape[-1] == 3 else img, 'gray')
        a.set_title(t)

    plt.show()

Traffic Sign Segmentation using K-means Clustering OpenCV

1 Answers