Commit 77a69531 authored by josh's avatar josh
Browse files

cleaning up dir structure

parent dbb1c3df
......@@ -9,7 +9,7 @@
## Example run:
```bash
python Scripts/initDB.py ../ ImageDatabase
python ./initDB.py ../ TestDatabase
```
## Directory Structure Result:
......@@ -28,7 +28,7 @@ python Scripts/initDB.py ../ ImageDatabase
## Example run:
```bash
python Scripts/importAndProcessImage.py TestResults/MS/GeoEye_MS_Original.jpg ../ImageDatabase/ GeoEye
python ./importAndProcessImage.py ../GeoEye_Original.jpg ../TestDatabase/ TestGeo
```
## Database Directory Result:
......@@ -126,5 +126,15 @@ python ./Scripts/meanShift3Channel.py ./TestResults/MS/GeoEye_MS_Original.jpg ./
## createModel.py
```bash
python createModel.py 128 128 3 m128x128 ../ImageDatabase/Images/0-jpg-GeoEye_Original-OR/Models/
python createModel.py 128 128 3 m128x128 ../TestDatabase/Images/0-jpg-TestGeo-OR/Models/
```
## trainModel.py
```bash
python trainModel.py ../TestDatabase/Images/0-jpg-TestGeo-OR/Models/m128x128/ ../roadSamples/
```
## modelPredict.py
```bash
python modelPredict.py ../TestDatabase/Images/0-jpg-TestGeo-OR/Models/m128x128/ ../roadPredict/
```
\ No newline at end of file
......@@ -7,51 +7,20 @@ import numpy as np
from PIL import Image
from pathlib import Path
from ai4hdr.model import HDRModel
from sample.model import HDRModel
INPUT_ID = "input"
MASK_ID = "mask"
def getSamples(dataDir: Path) -> [(np.array, np.array)]:
samples = []
sampleId = 0
for sampleDir in sorted(dataDir.iterdir()):
inputPath = sampleDir.joinpath("{}-sample-{}.jpg".format(sampleId, INPUT_ID))
maskPath = sampleDir.joinpath("{}-sample-{}.jpg".format(sampleId, MASK_ID))
inputImage = Image.open(inputPath)
maskImage = Image.open(maskPath)
newSize = (128, 128)
resizeInput = inputImage.resize(newSize)
resizeMask = maskImage.resize(newSize)
inputArr = np.array(resizeInput)
maskArr = np.array(resizeMask.convert("L"))
samples.append((inputArr, maskArr))
sampleId += 1
SIZE_X = samples[0][0].shape[0]
SIZE_Y = samples[0][0].shape[1]
X_TRAIN = []
Y_TRAIN = []
for sample in samples:
X_TRAIN.append(sample[0])
Y_TRAIN.append(sample[1])
X_TRAIN = np.array(X_TRAIN)
Y_TRAIN = np.array(Y_TRAIN)
return (X_TRAIN, Y_TRAIN)
def createAndSaveModel(shape: tuple, saveName: str, saveDir: Path):
saveDir = saveDir.joinpath(modelName)
os.mkdir(saveDir)
hdrModel = HDRModel(shape=shape)
hdrModel.save(saveDir)
if __name__ == "__main__":
inputShape = (int(sys.argv[1]), int(sys.argv[2]), int(sys.argv[3]))
modelName = sys.argv[4]
saveLoc = Path(sys.argv[5])
saveLoc = saveLoc.joinpath(modelName)
os.mkdir(saveLoc)
hdrModel = HDRModel(shape=inputShape)
hdrModel.save(saveLoc)
\ No newline at end of file
createAndSaveModel(inputShape, modelName, saveLoc)
\ No newline at end of file
import sys
import pathlib
from ai4hdrDatabase import importAndProcess
# TO DO::POTENTIAL PROBLEM - this function holds all the results in memory, may be too much for large images
# consider option to do limited what is produced
def sliceAndMeanShift(imgArr: np.array):
'''
Returns dictionary:
{
"slices":
{
"128": [ (x, y, imgArr) ], <- x, y are pixel start positions in original image
"256": ...,
"512":
},
"meanShift":
{
"128": [ (x, y, numRegions, segImage, labelImage) ],
"256": ...,
"512":
}
}
'''
finalResults = {
"slices":
{
"128": [],
"256": [],
"512": []
},
"meanShift":
{
"128": [],
"256": [],
"512": []
}
}
sizes = [ 128, 256, 512 ]
# Loop: 128, 256, 512
for size in sizes:
print("Processing Size:", size)
if imgArr.shape[0] >= size and imgArr.shape[1] >= size:
# Step 4: Slice
sliceResults = sliceImage(imgArr, (size, size))
# Step 5: Save slices
finalResults["slices"][str(size)] = sliceResults["slicedImages"]
# Step 6: Mean Shift on all slices
for res in sliceResults["slicedImages"]:
i = res[0]
j = res[1]
slicedImg = res[2].astype(np.uint8)
# mean shift
(segmentedImage, labelsImage, numberRegions) = meanShiftSegmentation(slicedImg)
# add to results
finalResults["meanShift"][str(size)].append((i, j, numberRegions, segmentedImage, labelsImage))
import numpy as np
return finalResults
from sample.database import importAndProcess
if __name__ == "__main__":
'''
......@@ -86,4 +25,4 @@ if __name__ == "__main__":
importAndProcess(imgPath=imgPath, dbPath=dbPath, newImgName=newName)
else:
print("Invalid arg count")
raise Exception("Invalid arg count")
import sys
import pathlib
from ai4hdrDatabase import initDB
from sample.database import initDB
if __name__ == "__main__":
'''
......@@ -17,9 +17,9 @@ if __name__ == "__main__":
dbAbsPath = pathlib.Path(dbLoc).joinpath(dbName).absolute()
if not success:
print("FAILED: Database {} already exsits at {}".format(dbName, dbAbsPath))
raise Exception("FAILED: Database {} already exsits at {}".format(dbName, dbAbsPath))
else:
print("SUCCESS: Database {} created at {}".format(dbName, dbAbsPath))
else:
print("Invalid arg count")
raise Exception("Invalid arg count")
import sys
from ai4hdr_utils import meanshift3Channel
if __name__ == "__main__":
'''
Required arguments:
- arg1: image file used for segmentation
- arg2: file name to save segmentation
'''
if len(sys.argv) == 3:
imagePath = sys.argv[1]
outPath = sys.argv[2]
'''
TO DO: Need customization of following parameters?
For now, using same defaults as the function
'''
quantile = 0.2
samples = 500
classColors = [ "darkgreen", "indigo", "gold" ]
result = meanshift3Channel(imagePath=imagePath,
outPath=outPath,
quantile=quantile,
samples=samples,
classColors=classColors)
else:
print("Invalid arg count")
\ No newline at end of file
......@@ -2,7 +2,7 @@ import sys
import pathlib
import cv2
import numpy as np
from ai4hdr_utils import meanShiftSegmentation
from sample.model import meanShiftSegmentation
if __name__ == "__main__":
'''
......@@ -26,20 +26,15 @@ if __name__ == "__main__":
# run algorithm
imgArr = cv2.imread(imagePath)
(segmentedImage, labelsImage, numberRegions) = meanShiftSegmentation(imgArr=imgArr,
spatialRadius=spatialRadius,
rangeRadius=rangeRadius,
minDensity=minDensity)
spatialRadius=spatialRadius,
rangeRadius=rangeRadius,
minDensity=minDensity)
# create save file paths
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]
#labelsImage = labelsImage.astype(float)
#labelsImage /= labelsImage.max()
#labelsImage *= 255
# save results to image files
cv2.imwrite(str(segColorPath), segmentedImage)
cv2.imwrite(str(labelImgPath), labelsImage)
......
import sys
import getopt
import pickle
import os
import numpy as np
from PIL import Image
from pathlib import Path
from sample.model import HDRModel
if __name__ == "__main__":
modelDir = Path(sys.argv[1])
imageDir = Path(sys.argv[2])
print(modelDir)
print(imageDir)
model = HDRModel(loadDir=modelDir)
model.featureModel.summary()
predictDir = imageDir.joinpath("MODEL_GENERATED")
predictions = {
}
for imagePath in imageDir.iterdir():
pathParts = imagePath.name.split(".")
imgName = pathParts[0]
imgExt = pathParts[-1]
imgInput = np.array(Image.open(imagePath))
print(imgInput.shape)
prediction = model.predict(imgInput)
predImage = Image.fromarray(prediction)
predictions[str(predictDir.joinpath("{}-PREDICTION.{}".format(imgName, imgExt)))] = predImage
os.mkdir(predictDir)
for key in predictions.keys():
predictions[key].save(key)
\ No newline at end of file
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 pprint import pprint
from matplotlib.colors import to_rgb
from sklearn.cluster import MeanShift, estimate_bandwidth
from skimage import color
'''
OTHER IMPLEMENTATIONS
'''
def meanshift3Channel(imagePath: str, outPath: str, quantile=0.2, samples=500, classColors: list=None) -> 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.
Arguments:
- imagePath : path to image used
- outPath : path to save segmented image (image type extension automatically added)
- 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
- classColors : custom list of colors to use for classes, default: [ "darkgreen", "indigo", "gold" ]
Note: list index = class label
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)
}
......@@ -2,6 +2,10 @@ import cv2
import pathlib
import os
import numpy as np
from sample.preprocess import sliceAndMeanShift
'''
Database file structure example:
......@@ -79,9 +83,11 @@ def importAndProcess(imgPath: str, dbPath: str, newImgName: str):
baseImageFilePath = baseImagePath.joinpath("{}.{}".format(baseImageName, fileExt))
cv2.imwrite(str(baseImageFilePath), imgArr)
# Create dir for sliced images
# Create dir for sliced images and models
slicedImagePath = baseImagePath.joinpath("SL")
os.mkdir(slicedImagePath)
modelsPath = baseImagePath.joinpath("Models")
os.mkdir(modelsPath)
# Run processing
processedResult = sliceAndMeanShift(imgArr)
......
......@@ -74,8 +74,6 @@ class HDRModel:
'''
Supports only a single input: shape=(x, y, z)
'''
plt.imshow(xInput)
plt.show()
# Extract features from input
xExpand = np.expand_dims(xInput, axis=0)
features = self.__getFeatures(xExpand)
......@@ -98,6 +96,6 @@ class HDRModel:
self.inputChannels = shape[-1]
vgg = VGG16(weights="imagenet", include_top=False, input_shape=shape)
self.featureModel = Model(inputs=vgg.input, outputs=vgg.get_layer("block1_conv2").output)
self.rfClassifer = RandomForestClassifier(n_estimators=50, random_state=42, verbose=1, n_jobs=1)
self.rfClassifer = RandomForestClassifier(n_estimators=50, random_state=42, verbose=1, n_jobs=-1)
......@@ -64,3 +64,69 @@ def meanShiftSegmentation(imgArr: np.array, spatialRadius: int=6, rangeRadius: i
spatial_radius=spatialRadius,
range_radius=rangeRadius,
min_density=minDensity)
# TO DO::POTENTIAL PROBLEM - this function holds all the results in memory, may be too much for large images
# consider option to do limited what is produced
def sliceAndMeanShift( imgArr: np.array ):
'''
Returns dictionary:
{
"slices":
{
"128": [ ( x, y, imgArr ) ], <- x, y are pixel start positions in original image
"256": ...,
"512":
},
"meanShift":
{
"128": [ ( x, y, numRegions, segImage, labelImage ) ],
"256": ...,
"512":
}
}
'''
finalResults = {
"slices":
{
"128": [],
"256": [],
"512": []
},
"meanShift":
{
"128": [],
"256": [],
"512": []
}
}
sizes = [ 128, 256, 512 ]
# Loop: 128, 256, 512
for size in sizes:
print( "Processing Size:", size )
if imgArr.shape[0] >= size and imgArr.shape[1] >= size:
# Step 4: Slice
sliceResults = sliceImage( imgArr, ( size, size ) )
# Step 5: Save slices
finalResults["slices"][str( size )] = sliceResults["slicedImages"]
# Step 6: Mean Shift on all slices
for res in sliceResults["slicedImages"]:
i = res[0]
j = res[1]
slicedImg = res[2].astype( np.uint8 )
#print( slicedImg.shape )
# mean shift
( segmentedImage, labelsImage, numberRegions ) = meanShiftSegmentation( slicedImg )
# add to results
finalResults["meanShift"][str( size )].append( ( i, j, numberRegions, segmentedImage, labelsImage ) )
return finalResults
\ No newline at end of file
import sys
import cv2
import pathlib
from ai4hdr_utils import sliceImage
from sample.preprocess import sliceImage
if __name__ == "__main__":
......
import sys
import shutil
import os
import numpy as np
from pathlib import Path
from PIL import Image
from sample.model import HDRModel
def getSamples(dataDir: Path) -> [(np.array, np.array)]:
samples = []
for sampleDir in dataDir.iterdir():
inputPath = sampleDir.joinpath("{}.jpg".format("input"))
maskPath = sampleDir.joinpath("{}.jpg".format("mask"))
inputImage = Image.open(inputPath)
maskImage = Image.open(maskPath)
#newSize = (128, 128)
#resizeInput = inputImage.resize(newSize)
#resizeMask = maskImage.resize(newSize)
inputArr = np.array(inputImage)
maskArr = np.array(maskImage.convert("L"))
samples.append((inputArr, maskArr))
X_TRAIN = []
Y_TRAIN = []
for sample in samples:
X_TRAIN.append(sample[0])
Y_TRAIN.append(sample[1])
return (np.array(X_TRAIN), np.array(Y_TRAIN))
if __name__ == "__main__":
'''
Train Model and save, overwriting loaded model
'''
modelDir = Path(sys.argv[1])
sampleDir = Path(sys.argv[2])
print(modelDir)
print(sampleDir)
(xTrain, yTrain) = getSamples(sampleDir)
xTrain = xTrain[0:3,:,:,:]
yTrain = yTrain[0:3,:,:]
print(xTrain.shape)
print(yTrain.shape)
model = HDRModel(loadDir=modelDir)
model.featureModel.summary()
model.fit(xTrain, yTrain)
# Remove loaded model and save trained model
shutil.rmtree(modelDir)
os.mkdir(modelDir)
model.save(modelDir)
\ No newline at end of file
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