import cupy as cp
import numpy as np
from sklearn.cluster import KMeans
from skimage import io
import time

# Load the image using skimage
image = io.imread('fruits.jpg')

# Convert the image to a CuPy array
image_cp = cp.asarray(image)

# Flatten the image into a 2D array of pixels
image_flat = image_cp.get().reshape(image_cp.shape[0] * image_cp.shape[1], image_cp.shape[2])

def Kmeans_cuda(K=1):
    # Use KMeans to cluster the pixels into a specified number of clusters
    kmeans = KMeans(n_clusters=K, random_state=0).fit(image_flat)

    # Predict the cluster for each pixel
    clusters = kmeans.predict(image_flat)

    # Create a new CuPy array to hold the modified image
    new_image_cp = cp.empty_like(image_cp)

    # Iterate over each pixel and assign its value to the corresponding cluster center
    for i, cluster in enumerate(clusters):
        new_image_cp[i // image_cp.shape[1], i % image_cp.shape[1]] = cp.asarray(kmeans.cluster_centers_[cluster])

    # Convert the CuPy array back to a NumPy array
    new_image = cp.asnumpy(new_image_cp)

    # Save the modified image using skimage
    io.imsave("fruits" + "_%d" % K + "_cuda.jpg", new_image)


for K in range(1,256):
    start_time = time.time()
    Kmeans_cuda(K=K)
    end_time = time.time()
    print(f"It took {end_time-start_time:.2f} seconds to compute for K =",K)