mean shift function implemented

README.md

@@ -14,21 +14,21 @@ Note: if the slice size doesn't fit evenly into the size of
 
 ## Example run
 ```bash
-python sliceImage.py ./example_images/GeoEye_After011310.jpg GeoEye_Slice 512 512 ./example_images/
+python sliceImage.py ./sliceImage_ex/GeoEye_After011310.jpg GeoEye_Slice 512 512 ./sliceImage_ex/
 ```
 
 ## Example run result
 ## Original Image: 1180x1180
-![](example_images/GeoEye_After011310.jpg)
+![](sliceImage_ex/GeoEye_After011310.jpg)
 
 ## Slice 1: 512x512
-![](example_images/GeoEye_Slice-0-0.jpg)
+![](sliceImage_ex/GeoEye_Slice-0-0.jpg)
 
 ## Slice 2: 512x512
-![](example_images/GeoEye_Slice-0-512.jpg)
+![](sliceImage_ex/GeoEye_Slice-0-512.jpg)
 
 ## Slice 3: 512x512
-![](example_images/GeoEye_Slice-512-0.jpg)
+![](sliceImage_ex/GeoEye_Slice-512-0.jpg)
 
 ## Slice 4: 512x512
-![](example_images/GeoEye_Slice-512-512.jpg)
+![](sliceImage_ex/GeoEye_Slice-512-512.jpg)
\ No newline at end of file

ai4hdr_utils.py

 import cv2
 import pathlib
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.colors as mcolors
 
+from matplotlib.colors import to_rgb
+from sklearn.cluster import MeanShift, estimate_bandwidth
+from skimage import color
 
-
-
-def meanshiftSegmentation( image ):
-    pass
-
-
-def sliceImage( imagePath: str, newFileBase: str, newSize: tuple, newPath: str ) -> int:
+def sliceImage( imagePath: str, newFileBase: str, newSize: tuple, newPath: str ) -> dict:
     '''
-    Creates slices of images given 
+    Creates slices of an image.
 
-    imagePath  : path to image used for slicing
-    newFileBase: used to rename each slice
-    newSize    : width and height of new slices
-    newPath    : where to save new slicest
+    Parameters:
+        - imagePath  : path to image used for slicing
+        - newFileBase: used to rename each slice
+        - newSize    : width and height of new slices
+        - newPath    : where to save new slicest
     
     returns:  dictionary containing the following information
     {
@@ -25,6 +26,7 @@ def sliceImage( imagePath: str, newFileBase: str, newSize: tuple, newPath: str )
     }
     OR empty dictionary if slicing was unable to be completed
     '''
+    fileExt = pathlib.Path( imagePath ).name.split( "." )[-1] 
     newPath = pathlib.Path( newPath )
 
     image = cv2.imread( imagePath )
@@ -50,15 +52,99 @@ def sliceImage( imagePath: str, newFileBase: str, newSize: tuple, newPath: str )
     # process image
     for i in range( 0, maxWidth, newWidth ):
         for j in range( 0, maxHeight, newHeight ):
-            slicePath = newPath.joinpath( "{}-{}-{}.jpg".format( newFileBase, i, j ) ) 
+            slicePath = newPath.joinpath( "{}-{}-{}.{}".format( newFileBase, i, j, fileExt ) ) 
             newSliceCount += 1
             newPaths.append( slicePath )
             cv2.imwrite( str( slicePath ), 
-                         image[i:i+newWidth, 
-                         j:j+newHeight, :] )
+                         image[i:i+newWidth, j:j+newHeight, :] )
 
     return {
         "pixelsLost": pixelsLost,
         "sliceCount": newSliceCount,
         "slicePaths": newPaths
-    }
+    }
\ No newline at end of file
+
+
+def meanshiftSegmentation( imagePath: str, outPath: str, classColors: list=None, quantile=0.2, samples=500 ) -> dict:
+    '''
+    SOURCE: https://www.efavdb.com/mean-shift
+
+    Creates and saves a segmented version of image using the mean shift algorithm to determine pixel classes.
+    Supports 3 channel images, each color channel is used as a feature.
+
+    Parameters:
+        - imagePath   : path to image used
+        - outPath     : path to save segmented image ( image type extension automatically added )
+        - classColors : custom list of colors to use for classes, default: [ "darkgreen", "indigo", "gold" ]
+                        Note: list index = class label
+        - quantile    : used for estimate_bandwidth function, should be between [0, 1], default: 0.2
+        - samples     : used for estimate_bandwidth function, number of samples to use, defualt: 500
+
+    Returns: a dictionary mapping the class integer to a tuple containing a color name and rgb value
+    {
+        0: ( "darkgreen", ( 0.0, 0.3, 0.0 ) ), 
+        1: ( "indigo",    ( 0.2, 0.0, 0.5 ) ), 
+        2: ( "gold",      ( 1.0, 0.8, 0.0 ) )
+    }
+
+
+    Quantile and samples are used for sklearn estimate_bandwidth function. 
+    For more info: https://scikit-learn.org/stable/modules/generated/sklearn.cluster.estimate_bandwidth.html
+    '''
+    if classColors is None:
+        classColors = [ "darkgreen", "indigo", "gold" ]
+
+    # create mapping for color name to rgb
+    color1 = to_rgb( classColors[0] )
+    color2 = to_rgb( classColors[1] )
+    color3 = to_rgb( classColors[2] )
+    
+    # cv2 uses BGR order, flip for RGB
+    colorToBGR = {
+        classColors[0]: ( color1[2], color1[1], color1[0] ),
+        classColors[1]: ( color2[2], color2[1], color2[0] ),
+        classColors[2]: ( color3[2], color3[1], color3[0] ),
+    }
+
+    fileExt = pathlib.Path( imagePath ).name.split( "." )[-1] 
+    
+    # read image and generate feature arrays
+    img     = cv2.imread( imagePath )
+    imgArr  = np.array( img )
+    flatArr = np.reshape( imgArr, [-1, 3] )
+
+    # run mean shift
+    bandwidth = estimate_bandwidth( X=flatArr, quantile=quantile, n_samples=samples )
+    meanShift = MeanShift( bandwidth, bin_seeding=True, n_jobs=-1 ) # n_jobs=-1 uses all processors
+    meanShift.fit( flatArr )
+
+    # use labels to construct a class array with classes: 0, 1, 2
+    labels = meanShift.labels_
+    classArr = np.reshape( labels, [img.shape[0], img.shape[1]] )
+
+    # use classColors to generate a 3 channel image from class array
+    rgbValues = []
+    for i in range( classArr.shape[0] ):
+        for j in range( classArr.shape[1] ):
+            color = classColors[classArr[i][j]]
+            rgbValues.append( colorToBGR[color] )
+    finalArr = np.asarray( rgbValues )
+    finalArr = np.reshape( finalArr, imgArr.shape )
+
+    # scale values to range [0, 255] to save with opencv
+    finalArr *= 255.0
+    cv2.imwrite( "{}.{}".format( outPath, fileExt ), finalArr )
+    
+    return {
+        0: ( classColors[0], color1 ),
+        1: ( classColors[1], color2 ),
+        2: ( classColors[2], color3 )
+    }
+    
+
+if __name__ == "__main__":
+    cwd       =  pathlib.Path.cwd()
+    imagePath = cwd.joinpath( "meanShiftSeg_ex/GeoEye_MS_Original.jpg" )
+    outPath   = cwd.joinpath( "meanShiftSeg_ex/GeoEye_MS_Segmented" )
+
+    result = meanshiftSegmentation( str( imagePath ), str( outPath ) )
\ No newline at end of file