M2_SETI/D3/TP/TP_SETI_Kmeans/Kmeans_skcuda.py
2022-12-22 19:38:17 +01:00

23 lines
715 B
Python

import numpy as np
import cv2
import cupy as cp
# Load the image and convert it to a NumPy array
image = cv2.imread("fruits.jpg")
image = image.astype(np.float32)
# Use cupy to transfer the image to the GPU
image_gpu = cp.asarray(image)
# Perform k-means clustering on the GPU
cluster_centers_gpu, labels_gpu, _ = cp.cluster.kmeans(image_gpu.reshape(-1, 3), k=8)
# Transfer the cluster centers and labels back to the CPU
cluster_centers = cp.asnumpy(cluster_centers_gpu)
labels = cp.asnumpy(labels_gpu)
# Convert the image pixels to the closest cluster
clustered_image = cluster_centers[labels].reshape(image.shape)
# Save the clustered image as a PNG file
cv2.imwrite("clustered_image.png", clustered_image)