pymeanshift implemented and program added

.ipynb_checkpoints/Untitled-checkpoint.ipynb

+{
+ "cells": [],
+ "metadata": {},
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

README.md

@@ -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
\ No newline at end of file
+ This program takes roughly 20 seconds to run, tested on an AMD FX(tm)-4300 Quad-Core Processor
+
+

ai4hdr_utils.py

 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
     '''
@@ -140,4 +135,69 @@ def meanshiftSegmentation( imagePath: str, outPath: str, quantile=0.2, samples=5
         1: ( classColors[1], color2 ),
         2: ( classColors[2], color3 )
     }
-    
\ No newline at end of file
+    
+
+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

meanShift.py → meanShift3Channel.py

 import sys
-from ai4hdr_utils import meanshiftSegmentation
+from ai4hdr_utils import meanshift3Channel
 
 if __name__ == "__main__":
     '''
@@ -19,11 +19,11 @@ if __name__ == "__main__":
         samples     = 500
         classColors = [ "darkgreen", "indigo", "gold" ]
 
-        result = meanshiftSegmentation( imagePath=imagePath,
-                                        outPath=outPath,
-                                        quantile=quantile,
-                                        samples=samples,
-                                        classColors=classColors )
+        result = meanshift3Channel( imagePath=imagePath,
+                                    outPath=outPath,
+                                    quantile=quantile,
+                                    samples=samples,
+                                    classColors=classColors )
 
     else:
         print( "Invalid arg count" )

meanShiftSeg.py

+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

sliceImage.py

 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" )