COMP9517 lab3 image segementation

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

from scipy import ndimage as ndi
from skimage.morphology import watershed
from skimage.feature import peak_local_max
from sklearn.cluster import MeanShift

from PIL import Image

size = 100, 100

img_names = ["shapes.png", "strawberry.png"]
ext_names = ["coins.png", "kiwi.png"]

images = [r'C:\Users\harri\Desktop\2020T2\9517\lab3\COMP9517_20T2_Lab3_Images\\'+i for i in img_names]
ext_images = [r'C:\Users\harri\Desktop\2020T2\9517\lab3\COMP9517_20T2_Lab3_Images\\'+i for i in ext_names]


def plot_three_images(figure_title, image1, label1,
                      image2, label2, image3, label3):
    fig = plt.figure()
    fig.suptitle(figure_title)

    # Display the first image
    fig.add_subplot(1, 3, 1)
    plt.imshow(image1)
    plt.axis('off')
    plt.title(label1)

    # Display the second image
    fig.add_subplot(1, 3, 2)
    plt.imshow(image2)
    plt.axis('off')
    plt.title(label2)

    # Display the third image
    fig.add_subplot(1, 3, 3)
    plt.imshow(image3)
    plt.axis('off')
    plt.title(label3)

    plt.show()

for img_path in ext_images:
    img = Image.open(img_path)
    img.thumbnail(size)


    # TODO: perform meanshift on image
    meanShift_img = np.array(img)[:,:,:3]
    R = meanShift_img[:,:,0:1].flatten()[:,np.newaxis]
    G = meanShift_img[:,:,1:2].flatten()[:,np.newaxis]
    B = meanShift_img[:,:,2:3].flatten()[:,np.newaxis]
    colour_samples = np.concatenate((R,G,B),axis=1)
    MS = MeanShift(bin_seeding=True)
    ms_labels = MS.fit_predict(colour_samples)  # CHANGE THIS
    ms_labels = ms_labels.reshape(meanShift_img.shape[0],meanShift_img.shape[1])
    # TODO: perform an optimisation and then watershed on image

    ws_img = img.convert('L')
    gray_ws_image = np.array(ws_img)
    distance = ndi.distance_transform_edt(gray_ws_image)
    local_max = peak_local_max(distance, indices=False, footprint=np.ones((3, 3)),
                            labels=gray_ws_image)
    markers = ndi.label(local_max)[0]

    ws_labels = watershed(-distance,markers,mask = gray_ws_image)  # CHANGE THIS


    # Display the results
    plot_three_images(img_path, img, "Original Image", ms_labels, "MeanShift Labels",
                      ws_labels, "Watershed Labels")