OpenCV 3: Using k-Nearest Neighbors to analyse RGB image

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:

• 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.

• Is there a way to feed the whole image into the findNearest method?

This would make it possible to avoid the slow python loop.

• 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)
else:
color = np.asarray([0, 0, 255]).astype(np.float32)

colorLayers[int_result][row, col] = color

cv2.imshow('Original', img)
cv2.imshow('Blue', colorLayers[0])
c

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)))
cv2.waitKey(0)
cv2.destroyAllWindows()

StackExchange Code Review Q#143529, answer score: 2