Commit 98361753 authored by josh's avatar josh
Browse files

pymeanshift implemented and program added

parent df68d081
{
"cells": [],
"metadata": {},
"nbformat": 4,
"nbformat_minor": 4
}
......@@ -33,7 +33,36 @@ python sliceImage.py ./sliceImageResults/GeoEye_Original.jpg GeoEye_Slice 512 51
## Slice 4: 512x512
![](sliceImageResults/GeoEye_Slice-512-512.jpg)
## "meanShift.py" program:
## "meanShiftSeg.py" program:
Uses the pymeanshift library to create a segmented images
Install pymeanshift: pip install git+git://github.com/fjean/pymeanshift.git
## Command line arguments:
- arg1: image file used for segmentation
- arg2: directory to save results
- arg3: base name for results files
- arg4: spatial radius for ms algo
- arg5: range radius for ms algo
- arg6: min density for ms algo
## Example run:
```bash
python meanShiftSeg.py ./meanShiftResults/GeoEye_MS_Original.jpg ./meanShiftResults/ pyms 6 6 50
```
## Example run result:
## Original Image:
![](meanShiftResults/GeoEye_MS_Original.jpg)
## Segmentation Image:
![](meanShiftResults/pyms-color-seg.jpg)
## Labels Image:
![](meanShiftResults/pyms-labels-image.jpg)
## "meanShift3Channel.py" program:
Uses the mean shift algorithm to create a segmentation mask for an image.
## Command line arguments:
......@@ -42,7 +71,7 @@ Uses the mean shift algorithm to create a segmentation mask for an image.
## Example run:
```bash
python meanShift.py meanShiftResults/GeoEye_MS_Original.jpg meanShiftResults/GeoEye_MS_Segmented
python meanShift3Channel.py meanShiftResults/GeoEye_MS_Original.jpg meanShiftResults/GeoEye_MS_Segmented
```
## Example run result:
......@@ -50,7 +79,9 @@ python meanShift.py meanShiftResults/GeoEye_MS_Original.jpg meanShiftResults/Geo
![](meanShiftResults/GeoEye_MS_Original.jpg)
## Segmentation Image:
![](meanShiftResults/GeoEye_MS_Segmented.jpg)
![](meanShiftResults/three_channel_seg.jpg)
## Program run time:
This program takes roughly 20 seconds to run, tested on an AMD FX(tm)-4300 Quad-Core Processor
This diff is collapsed.
import cv2
import pathlib
import numpy as np
import os
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import pymeanshift as pms
from matplotlib.colors import to_rgb
from sklearn.cluster import MeanShift, estimate_bandwidth
from skimage import color
def sliceImage( imagePath: str, newFileBase: str, newSize: tuple, newPath: str ) -> dict:
def sliceImage( imagePath: str, newSize: tuple ) -> dict:
'''
Creates slices of an image.
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
{
"pixelsLost": int,
"sliceCount": int,
"slicePaths": list
"slicedImages": [ ( xStartPos, yStartPos, imgArr ) ],
"pixelsLost" : int,
"sliceCount" : int
}
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 )
......@@ -47,25 +46,22 @@ def sliceImage( imagePath: str, newFileBase: str, newSize: tuple, newPath: str )
# return data
pixelsLost = ( width - maxWidth ) + ( height - maxHeight )
newSliceCount = 0
newPaths = []
imageList = []
# process image
for i in range( 0, maxWidth, newWidth ):
for j in range( 0, maxHeight, newHeight ):
slicePath = newPath.joinpath( "{}-{}-{}.{}".format( newFileBase, i, j, fileExt ) )
newSliceCount += 1
newPaths.append( slicePath )
cv2.imwrite( str( slicePath ),
image[i:i+newWidth, j:j+newHeight, :] )
imageList.append( ( i, j, image[i:i+newWidth, j:j+newHeight, :] ) )
return {
"pixelsLost": pixelsLost,
"sliceCount": newSliceCount,
"slicePaths": newPaths
"slicedImages": imageList,
"pixelsLost" : pixelsLost,
"sliceCount" : newSliceCount
}
def meanshiftSegmentation( imagePath: str, outPath: str, quantile=0.2, samples=500, classColors: list=None ) -> dict:
def meanshift3Channel( imagePath: str, outPath: str, quantile=0.2, samples=500, classColors: list=None ) -> dict:
'''
SOURCE: https://www.efavdb.com/mean-shift
......@@ -87,7 +83,6 @@ def meanshiftSegmentation( imagePath: str, outPath: str, quantile=0.2, samples=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
'''
......@@ -141,3 +136,68 @@ def meanshiftSegmentation( imagePath: str, outPath: str, quantile=0.2, samples=5
2: ( classColors[2], color3 )
}
def meanShiftSegmentation( imagePath: str, outPath: str, spatialRadius: int=6, rangeRadius: int=6, minDensity: int=50 ):
'''
SOURCE - https://github.com/fjean/pymeanshift
return: ( segmented_image, labels_image, number_regions )
'''
original_image = cv2.imread( imagePath )
return pms.segment( original_image,
spatial_radius=spatialRadius,
range_radius=rangeRadius,
min_density=minDensity )
'''
Database functions
Database file structure example:
|-- ImageDatabase
│   |-- MeanShift
| | |--image_name-id
| | | |--128x128
| | | | |--image_name-128-smap-0-0.<ext> <- segmented map
| | | | |--image_name-128-simage-0-0.<ext> <- segmented image
| | | |--256x256
| | | |--512x512
│   |-- Original
| | |--image_name-id.<ext> <- original image, id = auto increment integer
│   |-- Sliced
| | |--image_name-id
| | | |--128x128
| | | | |--image_name-128-0-0.<ext> <- sliced image
| | | |--256x256
| | | |--512x512
'''
def initDB( baseDir: str, name: str ) -> bool:
'''
Initialized the image database at a given directory
File Struture:
|-- ImageDatabase
│   |-- MeanShift
│   |-- Original
│   |-- Sliced
'''
baseDir = pathlib.Path( baseDir )
dbDir = baseDir.joinpath( name )
if dbDir.is_dir():
print( "Failed to init: a database with this name already exists." )
return False
else:
originalImagesDir = dbDir.joinpath( "Original" )
slicedImagesDir = dbDir.joinpath( "Sliced" )
meanShiftSegDir = dbDir.joinpath( "MeanShift" )
os.mkdir( dbDir )
os.mkdir( originalImagesDir )
os.mkdir( slicedImagesDir )
os.mkdir( meanShiftSegDir )
if __name__ == "__main__":
baseDir = "../"
initDB( baseDir, "ImageDatabase" )
\ No newline at end of file
import sys
from ai4hdr_utils import meanshiftSegmentation
from ai4hdr_utils import meanshift3Channel
if __name__ == "__main__":
'''
......@@ -19,7 +19,7 @@ if __name__ == "__main__":
samples = 500
classColors = [ "darkgreen", "indigo", "gold" ]
result = meanshiftSegmentation( imagePath=imagePath,
result = meanshift3Channel( imagePath=imagePath,
outPath=outPath,
quantile=quantile,
samples=samples,
......
import sys
import pathlib
import cv2
import numpy as np
from ai4hdr_utils import meanShiftSegmentation
if __name__ == "__main__":
'''
Required arguments:
- arg1: image file used for segmentation
- arg2: directory to save results
- arg3: base name for results files
- arg4: spatial radius for ms algo
- arg5: range radius for ms algo
- arg6: min density for ms algo
'''
if len( sys.argv ) == 7:
imagePath = sys.argv[1]
outPath = pathlib.Path( sys.argv[2] )
baseFileName = sys.argv[3]
spatialRadius = int( sys.argv[4] )
rangeRadius = int( sys.argv[5] )
minDensity = int( sys.argv[6] )
( segmentedImage, labelsImage, numberRegions ) = meanShiftSegmentation( imagePath=imagePath,
outPath=outPath,
spatialRadius=spatialRadius,
rangeRadius=rangeRadius,
minDensity=minDensity )
fileExt = pathlib.Path( imagePath ).name.split( "." )[-1]
segColorPath = outPath.joinpath( "{}-color-seg.{}".format( baseFileName, fileExt ) )
labelImgPath = outPath.joinpath( "{}-labels-image.{}".format( baseFileName, fileExt ) )
# scale labelsImage to [0,255]
print( labelsImage.max() )
labelsImage = labelsImage.astype( float )
labelsImage /= labelsImage.max()
labelsImage *= 255
# save results to image files
cv2.imwrite( str( segColorPath ), segmentedImage )
cv2.imwrite( str( labelImgPath ), labelsImage )
else:
print( "Invalid arg count" )
\ No newline at end of file
import sys
import cv2
import pathlib
from ai4hdr_utils import sliceImage
......@@ -17,7 +19,21 @@ if __name__ == "__main__":
newFileBase = sys.argv[2]
newSize = ( int( sys.argv[3] ), int( sys.argv[4] ) )
newPath = sys.argv[5]
sliceImage( imagePath, newFileBase, newSize, newPath )
# get sliced images
result = sliceImage( imagePath, newSize )
fileExt = pathlib.Path( imagePath ).name.split( "." )[-1]
newPath = pathlib.Path( newPath )
# save each image to DB
for res in result["slicedImages"]:
i = res[0]
j = res[1]
imgArr = res[2]
slicePath = newPath.joinpath( "{}-{}-{}.{}".format( newFileBase, i, j, fileExt ) )
cv2.imwrite( str( slicePath ), imgArr )
else:
print( "Invalid arg count" )
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