#! /usr/bin/env python3
#
def exercise4 ( ):

#*****************************************************************************80
#
## exercise4() uses K-Means to simplify an image.
#
#  Discussion:
#
#    Suppose an image is a (w x h x d) array of unsigned integers.
#    The w and h are the width and height of the image in pixels.
#    d is 3, and represents the values of red, green, and blue
#    that specify the color of each pixel.
#
#    We would like to regard the colors as data to be clustered.
#    In particular, we want to create K color clusters, such that
#    each color used in the image is assigned to one of these clusters.
#
#    Each cluster can be represented by its center value Z.
#
#    We reason that if each color is close to its cluster center Z,
#    we might be able to simplify the picture so that only K colors 
#    are used, namely, the K cluster centers Z.
#
#    We can try to create a new image, using just these K colors.
#
#    This is a standard process in image compression and image analysis.    
#
#    Because the image is stored as a triple index array of unsigned
#    integers, we need to do some mysterious manipulations to convert
#    the image to a numpy array, process it with kmeans(), and then
#    convert the result back into something that looks like an image
#    that we can display!
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    25 March 2025
#
#  Author:
#
#    John Burkardt
#
  from scipy.cluster.vq import kmeans2
  import imageio.v2 as imageio
  import matplotlib.pyplot as plt
  import numpy as np
  import platform

  print ( '' )
  print ( 'exercise4():' )
  print ( '  Python version: ' + platform.python_version ( ) )
  print ( '  Use K-Means algorithm to reduce the colors in an image.' )
#
#  Read an image, using the imageio library.
#
  filename = 'swim.jpg'
  print ( '' )
  print ( '  Reading "' + filename + '"' )
  image = imageio.imread ( filename )
#
#  Display the image, print its shape, and type.
#
  plt.imshow ( image )
  plt.show ( )
  print ( image.shape )
  print ( type ( image ) )
  print ( type ( image[0,0,0] ) )
#
#  Convert the numeric type from uint8 to float64.
#  uint8 is an 8 bit unsigned integer between 0 and 255.
#
  data = np.array ( image, dtype = np.float64 ) / 255
#
#  Reshape the image from a 3D object to 2D, so that kmeans can handle it.
#
  w, h, d = data.shape
  print ( '  Image is ', w, ' wide by ', h, ' high with ', d, ' RGB intensities' )
  data = np.reshape ( data, ( w * h, d ) )
#
#  Try to count unique colors.
#
  colors_unique = np.unique ( data, axis = 0 )
  colors_count = len ( colors_unique )
  print ( '  Original image uses ', colors_count, ' unique colors.' )
#
#  Use K-Means to create k clusters of colors.
#
  k = 10
  Z, C = kmeans2 ( data, k )
#
#  For each pixel, 
#    C is the index of the cluster it belongs to.
#    Z is the center of a cluster (an "average" color)
#
#  Z[C] is a 2D array replacing each color by its cluster color.
#  We need to reshape this array back to (w x h x d).
#
  data2 = Z[C]

  colors_unique = np.unique ( data2, axis = 0 )
  colors_count = len ( colors_unique )
  print ( '  Clustered image uses ', colors_count, ' unique colors.' )

  data2 = data2.reshape ( w, h, d )
  plt.imshow ( data2 )
  plt.show ( )
#
#  Convert image back to uint8 arithmetic.
#
  image2 = np.array ( 255 * data2, dtype = np.uint8 )
#
#  Write the modified image.
#
# filename = 'swim_' + str(k) + '_colors.png'
  filename = 'swim_' + str(k) + '_colors.jpg'
  imageio.imwrite ( filename, image2 )
  print ( '  Graphics saved as "' + filename + '"' )

  return

if ( __name__ == "__main__" ):
  exercise4 ( )