Commit dbb1c3df authored by josh's avatar josh

reorganizing file structure, adding model hdr model class, createModel script

parent 042b2d9a
__pycache__/
modules/__pycache__/
\ No newline at end of file
modules/__pycache__/
.vscode/
.ipynb_checkpoints/
\ No newline at end of file
......@@ -124,3 +124,7 @@ python ./Scripts/meanShift3Channel.py ./TestResults/MS/GeoEye_MS_Original.jpg ./
This program takes roughly 20 seconds to run, tested on an AMD FX(tm)-4300 Quad-Core Processor
## createModel.py
```bash
python createModel.py 128 128 3 m128x128 ../ImageDatabase/Images/0-jpg-GeoEye_Original-OR/Models/
```
\ No newline at end of file
This source diff could not be displayed because it is too large. You can view the blob instead.
from ai4hdr_backend.lib.ai4hdr import importAndProcess
\ No newline at end of file
from ai4hdr_backend.modules.ai4hdr import importAndProcess
from ai4hdr_backend.modules.ai4hdrModel import HDRModel
\ 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
def sliceImage(imgArr: np.array, newSize: tuple) -> dict:
'''
Creates slices of an image.
Arguments:
- imgArr : image array used for slicing
- newSize : width and height of new slices
returns: dictionary containing the following information
{
"slicedImages": [ (xStartPos, yStartPos, slicedImgArr) ],
"sliceCount" : int
}
OR empty dictionary if slicing was unable to be completed
'''
width = imgArr.shape[0]
height = imgArr.shape[1]
newWidth = newSize[0]
newHeight = newSize[1]
# Return empty dictionary if desired slice is larger than image
if newWidth > width or newHeight > height:
return {}
# Return data
newSliceCount = 0
imageList = []
# Process image
for i in range(0, width, newWidth):
for j in range(0, height, newHeight):
newSliceCount += 1
# Create array of zeros of slice size
base = np.zeros((newWidth, newHeight, imgArr.shape[2]))
# Slice image and broadcase to base array
# Slice will be blacked out when outside range of original image
sliceArr = imgArr[i:i+newWidth, j:j+newHeight, :]
base[ :sliceArr.shape[0], :sliceArr.shape[1]] = sliceArr
imageList.append((i, j, base))
return {
"slicedImages": imageList,
"sliceCount" : newSliceCount
}
def meanShiftSegmentation(imgArr: np.array, spatialRadius: int=6, rangeRadius: int=6, minDensity: int=50) -> tuple:
'''
SOURCE - https://github.com/fjean/pymeanshift
Arguments:
- imgArr : image array to use for algo
- spatialRadius: spatial radius for ms algo
- rangeRadius : range radius for ms algo
- minDensity : min density for ms algo
return: (segmented_image, labels_image, number_regions)
'''
return pms.segment(imgArr,
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
'''
Database file structure example:
......@@ -161,7 +18,6 @@ Database file structure example:
| | | | |--256x256
| | | | |--512x512
'''
def initDB(dbLoc: str, name: str) -> bool:
'''
Initialized the image database at a given directory
......@@ -180,12 +36,13 @@ def initDB(dbLoc: str, name: str) -> bool:
return True
def importAndProcess(imgPath: str, dbPath: str, newImgName: str):
'''
Slices and segments an image and imports it into the database
'''
imgPath = pathlib.Path(imgPath)#.absolute()
dbPath = pathlib.Path(dbPath)#.absolute()
imgPath = pathlib.Path(imgPath)
dbPath = pathlib.Path(dbPath)
# Ensure database exists
if not dbPath.is_dir():
......@@ -275,83 +132,4 @@ def importAndProcess(imgPath: str, dbPath: str, newImgName: str):
# Save processed image and label array
cv2.imwrite(str(msSegFileDir), msSegArr)
np.save(str(msLabFileDir), msLabImg)
'''
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)
}
np.save(str(msLabFileDir), msLabImg)
\ No newline at end of file
import numpy as np
import pickle
import os
import matplotlib.pyplot as plt
......@@ -6,6 +8,7 @@ from pathlib import Path
from PIL import Image
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import Input, Conv2D
from tensorflow.keras.models import Model
from tensorflow.keras.applications.vgg16 import VGG16
......@@ -15,14 +18,45 @@ from sklearn.ensemble import RandomForestClassifier
class HDRModel:
def __init__(self, modelShape):
def __init__(self, shape=None, loadDir=None):
self.featureModel = None
self.rfClassifer = None
self.inputChannels = 3
self.outputChannels = 1
self.__initModels(modelShape)
self.__fmSaveDir = "featureModel"
self.__fmSaveFile = "fm.h5"
self.__rfSaveDir = "rfClassifer"
self.__rfSaveFile = "classifier.sav"
if loadDir is not None:
self.load(loadDir)
elif shape is not None:
self.__initModels(shape)
def save(self, saveDir):
# Save tf model
fmDir = saveDir.joinpath(self.__fmSaveDir)
os.mkdir(fmDir)
self.featureModel.save(str(fmDir.joinpath(self.__fmSaveFile)))
# Save sklearn random forest
classifierDir = saveDir.joinpath(self.__rfSaveDir)
os.mkdir(classifierDir)
pickle.dump(self.rfClassifer, open(classifierDir.joinpath(self.__rfSaveFile), "wb"))
def load(self, modelDir):
# Load tf model
fmPath = modelDir.joinpath(self.__fmSaveDir).joinpath(self.__fmSaveFile)
self.featureModel = keras.models.load_model(str(fmPath))
self.featureModel.summary()
# Load sklean model
rfPath = modelDir.joinpath(self.__rfSaveDir).joinpath(self.__rfSaveFile)
self.rfClassifer = pickle.load(open(rfPath, "rb"))
def fit(self, xTrain, yTrain):
......@@ -32,8 +66,6 @@ class HDRModel:
# Reshape input for random forest
xForestTrain = features.reshape(-1, features.shape[-1])
yForestTrain = yTrain.reshape(-1)
print("Forest Train X:", xForestTrain.shape)
print("Forest Train Y:", yForestTrain.shape)
self.rfClassifer.fit(xForestTrain, yForestTrain)
......@@ -42,23 +74,19 @@ class HDRModel:
'''
Supports only a single input: shape=(x, y, z)
'''
print("X INPUT", xInput.shape)
plt.imshow(xInput)
plt.show()
# Extract features from input
xExpand = np.expand_dims(xInput, axis=0)
features = self.__getFeatures(xExpand)
# Reshape input for random forest
forestInput = features.reshape(-1, features.shape[-1])
print("FOREST INPUT", forestInput.shape)
forestOutput = self.rfClassifer.predict(forestInput)
# Reshape output back to image shape
outputShape = (xInput.shape[0], xInput.shape[1])
finalOutput = forestOutput.reshape(outputShape)
print("FINAL OUTPUT", finalOutput.shape)
plt.imshow(finalOutput)
plt.show()
return finalOutput
......@@ -70,56 +98,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)
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)
if __name__ == "__main__":
sampleDir = Path.cwd().parent.joinpath("roadTest")
(xTrain, yTrain) = getSamples(sampleDir)
inputShape = X_TRAIN[0].shape
hdrModel = HDRModel(inputShape)
hdrModel.featureModel.summary()
self.rfClassifer = RandomForestClassifier(n_estimators=50, random_state=42, verbose=1, n_jobs=1)
hdrModel.fit(X_TRAIN, Y_TRAIN)
prediction = hdrModel.predict(X_TRAIN[0])
\ No newline at end of file
import numpy as np
import pymeanshift as pms
def sliceImage(imgArr: np.array, newSize: tuple) -> dict:
'''
Creates slices of an image.
Arguments:
- imgArr : image array used for slicing
- newSize : width and height of new slices
returns: dictionary containing the following information
{
"slicedImages": [ (xStartPos, yStartPos, slicedImgArr) ],
"sliceCount" : int
}
OR empty dictionary if slicing was unable to be completed
'''
width = imgArr.shape[0]
height = imgArr.shape[1]
newWidth = newSize[0]
newHeight = newSize[1]
# Return empty dictionary if desired slice is larger than image
if newWidth > width or newHeight > height:
return {}
# Return data
newSliceCount = 0
imageList = []
# Process image
for i in range(0, width, newWidth):
for j in range(0, height, newHeight):
newSliceCount += 1
# Create array of zeros of slice size
base = np.zeros((newWidth, newHeight, imgArr.shape[2]))
# Slice image and broadcase to base array
# Slice will be blacked out when outside range of original image
sliceArr = imgArr[i:i+newWidth, j:j+newHeight, :]
base[ :sliceArr.shape[0], :sliceArr.shape[1]] = sliceArr
imageList.append((i, j, base))
return {
"slicedImages": imageList,
"sliceCount" : newSliceCount
}
def meanShiftSegmentation(imgArr: np.array, spatialRadius: int=6, rangeRadius: int=6, minDensity: int=50) -> tuple:
'''
SOURCE - https://github.com/fjean/pymeanshift
Arguments:
- imgArr : image array to use for algo
- spatialRadius: spatial radius for ms algo
- rangeRadius : range radius for ms algo
- minDensity : min density for ms algo
return: (segmented_image, labels_image, number_regions)
'''
return pms.segment(imgArr,
spatial_radius=spatialRadius,
range_radius=rangeRadius,
min_density=minDensity)
import sys
import getopt
import pickle
import os
import numpy as np
from PIL import Image
from pathlib import Path
from ai4hdr.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)
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
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