Commit 7840e5a4 authored by josh's avatar josh
Browse files

mean shift function implemented

parent cb7abc06
This diff is collapsed.
......@@ -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)
\ No newline at end of file
![](sliceImage_ex/GeoEye_Slice-512-512.jpg)
\ No newline at end of file
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
This diff is collapsed.
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment