from sklearn.cluster import KMeans
from skimage import io, color
import numpy as np
import matplotlib.pyplot as plt


# Function to convert RGB to Lab color space
def rgb_to_lab(rgb):
    return color.rgb2lab(np.array([[rgb]]))[0][0]


# Load the image
image = io.imread("/home/cory/Extra/Pictures/cube.png")

# Reshape the image to a list of RGB values
pixels = image.reshape((-1, 4))
pixels = [pixel[0:3] for pixel in pixels]
# Convert RGB to Lab
lab_pixels = np.apply_along_axis(rgb_to_lab, 1, pixels)

# Choose the number of clusters (groups)
num_clusters = 6  # You can adjust this based on your needs

# Apply k-means clustering
kmeans = KMeans(n_clusters=num_clusters, random_state=42)
kmeans.fit(lab_pixels)

# Assign each pixel to its cluster
labels = kmeans.labels_

# Reshape labels back to the shape of the original image
labels = labels.reshape(image.shape[:2])

# Visualize the clustered image
plt.imshow(labels, cmap='viridis')
plt.colorbar()
plt.savefig('filename.svg')
plt.show()