How to perform k means clustering for images in a tf dataset?

Question

I have a custom function like this to perform k means clustering for images on tf dataset(only on training dataset)

def segment_image(images,labels):

      print(type(images))
      print('hello')
      vectorized = images.reshape((-1,3))
      kmeans = KMeans(n_clusters= 5, random_state = 0, n_init='auto').fit(images)
      centers = np.uint8(kmeans.cluster_centers_)
      segmented_data = centers[kmeans.labels_.flatten()]
      seg_image = segmented_data.reshape((image.shape))
      segmented_image=segmented_image. convert_to_tensor()
  
  return segmented_image,labels 

for reshaping i need to convert the argument "images " to a numpy array. I tried various method and iterations but not happening. I am able to convert my train_ds to numpy array but not the images separately

After this i can call this function on the dataset as train_ds = train_ds.map(segment_image)

Any suggestion please ??