41 lines
No EOL
1.3 KiB
Python
41 lines
No EOL
1.3 KiB
Python
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) |