I'm new to computer vision and numpy.

I wrote a simple script to seperate red, green and blue colors from the original image by using the kNN algorithm. After reading through some numpy tutorials, I'm under the impression that the performance of my script can be improved heavily. So here are my questions:

  1. Is there a way to avoid the call to np_pixel = np.asarray(...) for each of the pixels? I can't find a format for the trainData that would allow me that.
  2. Is there a way to feed the whole image into the findNearest method? This would make it possible to avoid the slow python loop.
  3. Do you see any other ways to improve the performance or length of this script?
import cv2
import numpy as np

trainData = np.array([[255, 0, 0], [0, 255, 0], [0, 0, 255]], dtype=np.float32)
# Labels each one of the BGR channels as numbers 0, 1, 2 respectively
responses = np.array([[0], [1], [2]], dtype=np.float32)

imgPath = "../Images/RGB.png"
img = cv2.imread(imgPath)
if img is None:
    raise FileNotFoundError("'{0}' could not be opened!".format(imgPath))

knn = cv2.ml.KNearest_create()
knn.train(trainData, cv2.ml.ROW_SAMPLE, responses)

height, width, depth = img.shape
colorLayers = dict()
colorLayers[0] = np.zeros((height, width, 3), dtype=np.float32)
colorLayers[1] = np.zeros((height, width, 3), dtype=np.float32)
colorLayers[2] = np.zeros((height, width, 3), dtype=np.float32)

for row in range(height):
    for col in range(width):
        np_pixel = np.asarray([img[row, col]], dtype=np.float32)
        ret, results, neighbours, dist = knn.findNearest(np_pixel, 1)
        int_result = int(results)

        if int_result == 0:
            color = np.asarray([255, 0, 0]).astype(np.float32)
        elif int_result == 1:
            color = np.asarray([0, 255, 0]).astype(np.float32)
            color = np.asarray([0, 0, 255]).astype(np.float32)

        colorLayers[int_result][row, col] = color

cv2.imshow('Original', img)
cv2.imshow('Blue', colorLayers[0])
cv2.imshow('Green', colorLayers[1])
cv2.imshow('Red', colorLayers[2])
  • \$\begingroup\$ Could you link to an example image to use? Or can we use any image? \$\endgroup\$
    – Graipher
    Oct 7, 2016 at 15:58
  • \$\begingroup\$ @Graipher You can use any RGB image you like. The script clusters each pixel color to either R, G or B and then fills the pixel with that color. Here's a picture that works well with this script. \$\endgroup\$
    – dev-random
    Oct 8, 2016 at 7:14

1 Answer 1


I found a way to get rid of the python loop. np.vstack() converts the image to a single row which in this case is the format findNearest() expects.

Running below code through IPythons timeit function yields a huge perfomance boost: The version in the original answer takes 2.5s per loop. The new version runs at 15.3ms per loop.

import cv2
import numpy as np

imgPath = r"../Images/landscape2.jpg"
img_BGR = cv2.imread(imgPath)
if img_BGR is None:
    raise FileNotFoundError("'{0}' could not be opened!".format(imgPath))

BGR_COLORS = dict(blue=(255, 0, 0), green=(0, 255, 0), red=(0, 0, 255), black=(0, 0, 0), white=(255, 255, 255))
LABELS = dict(blue=np.array([0]), green=np.array([1]), red=np.array([2]), black=np.array([3]), white=np.array([4]))
trainData = np.array([BGR_COLORS['blue'], BGR_COLORS['green'], BGR_COLORS['red'], BGR_COLORS['black'], BGR_COLORS['white']], dtype=np.float32)
responses = np.array([[LABELS['blue']], LABELS['green'], LABELS['red'], LABELS['black'], LABELS['white']], dtype=np.float32)

knn = cv2.ml.KNearest_create()
knn.train(trainData, cv2.ml.ROW_SAMPLE, responses)

img_vstacked = np.vstack(img_BGR).astype(np.float32)
ret, results, neighbours, dist = knn.findNearest(img_vstacked, 1)

height, width, depth = img_BGR.shape
results_int = results.reshape(height, width).astype(np.uint8)

def colorPixels(image, results, colorName):
    image[results[:, :] == LABELS[colorName]] = BGR_COLORS[colorName]

img_clustered = img_BGR.copy()
for colorName in BGR_COLORS.keys():
    colorPixels(img_clustered, results_int, colorName)

cv2.imshow("Original vs. Clustered", np.hstack((img_BGR, img_clustered)))

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.