Image segmentation using k-means clustering

segmentation.py

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

image = cv2.imread("sample-images/minions.jpg")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
plt.imshow(image)

pixel_values = image.reshape((-1, 3))
pixel_values = np.float32(pixel_values)

criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.2)

k = 5
_, labels, (centers) = cv2.kmeans(pixel_values, k, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)

centers = np.uint8(centers)
segmented_image = centers[labels.flatten()]
segmented_image = segmented_image.reshape(image.shape)

plt.imshow(segmented_image)

gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
gaussian = cv2.GaussianBlur( gray,(5,5),0)
median = cv2.medianBlur(gaussian, 3)
bilateral = cv2.bilateralFilter(median, 7, 50, 50)
edges = cv2.Canny(bilateral, 0, 60, L2gradient=False)

plt.imshow(edges)

coloured_edges = cv2.bitwise_and(segmented_image, segmented_image, mask=edges)
coloured_edges = 255 - coloured_edges
print(np.average(coloured_edges))

plt.imshow(coloured_edges)