superpixel_hsv_segmentation.py

import numpy as np
import scipy.sparse as sp
import matplotlib.pyplot as plt

from scipy.spatial.distance import cdist

from sklearn.cluster import AgglomerativeClustering

from skimage.io import imread
from skimage.color import rgb2hsv
from skimage.transform import resize
from skimage.morphology import binary_dilation
from skimage.segmentation import slic, mark_boundaries


def get_neighbors_ids(sp_map, sp_id):
    return set(np.unique(sp_map[binary_dilation(sp_map == sp_id)])) - {sp_id}

def superpixels_pooling(features, sp_map, sp_id):
    return features[sp_map == sp_id].mean(0)

@np.vectorize
def superpixel_id_to_idx(mapping_dict, k):
    return mapping_dict[k]


img = imread("test_image.jpg")[:, :, :3]
img = resize(img, (224, 224))

hsv = rgb2hsv(img)

superpixels = slic(img, n_segments=200)

superpixels_features = {
    sp_id: superpixels_pooling(hsv, superpixels, sp_id)
    for sp_id in np.unique(superpixels)
}

superpixel_id_to_idx_dict = {k: i for i, k in enumerate(sorted(superpixels_features.keys()))}
superpixels_features_array = np.array([t[1] for t in sorted(superpixels_features.items())])

superpixels_distances_matrix = cdist(
    superpixels_features_array,
    superpixels_features_array
)



model = AgglomerativeClustering(n_clusters=6, affinity='precomputed', linkage='complete').fit(superpixels_distances_matrix)


cluster_map = np.zeros((224, 224))
for i, v in enumerate(model.labels_):
    cluster_map[superpixels == list(superpixel_id_to_idx_dict.keys())[i]] = v

plt.imshow(cluster_map)
plt.show()

test_image.jpg