analysis on all possible files in a provided directory, writes to JSON
This commit is contained in:
2
.gitignore
vendored
Normal file
2
.gitignore
vendored
Normal file
@@ -0,0 +1,2 @@
|
||||
__pycache__
|
||||
predictions
|
||||
31
main.py
31
main.py
@@ -3,23 +3,46 @@
|
||||
# most of this application adapted from the following walkthrough:
|
||||
# https://towardsdatascience.com/how-to-use-a-pre-trained-model-vgg-for-image-classification-8dd7c4a4a517
|
||||
|
||||
import sys, os
|
||||
import sys, os, json, time
|
||||
from predict import predict
|
||||
from keras.applications.vgg16 import VGG16
|
||||
|
||||
print("\n\n\n")
|
||||
print("Imports successful! Running startup processes...")
|
||||
|
||||
# generate current time for use in identifying outfiles
|
||||
cur_time = str(int(time.time()))
|
||||
|
||||
# create the target directory if it doesn't exist
|
||||
if (not os.path.exists("./predictions")):
|
||||
print("Did not find predictions directory, creating...")
|
||||
os.makedirs("./predictions")
|
||||
|
||||
# declare model to be used for each prediction
|
||||
model = VGG16(weights='imagenet')
|
||||
|
||||
# receive directory path as CLI argument and get a list of all files in path
|
||||
path = sys.argv[1]
|
||||
if (path[-1] != "/"):
|
||||
path += "/"
|
||||
|
||||
files = os.listdir(path)
|
||||
|
||||
# store all results in one list
|
||||
all_results = []
|
||||
|
||||
print("Running image analysis. This may take some time")
|
||||
|
||||
# for each file in directory, append its prediction result to main list
|
||||
for file in files:
|
||||
result = predict(model, file)
|
||||
all_results.append({ path: file, result: result })
|
||||
result = predict(model, path + file)
|
||||
if result is not None:
|
||||
all_results.append({ "path": file, "prediction": result })
|
||||
|
||||
print(all_results)
|
||||
print("Analysis complete! Writing JSON to ./predictions/predictions" + cur_time + ".json")
|
||||
|
||||
# convert object to JSON and write to JSON file
|
||||
with open("./predictions/predictions" + cur_time + ".json", "w") as outfile:
|
||||
json.dump(all_results, outfile)
|
||||
|
||||
print("Process complete!")
|
||||
|
||||
10
predict.py
10
predict.py
@@ -3,6 +3,10 @@ from keras.utils import load_img, img_to_array
|
||||
from keras.applications.vgg16 import preprocess_input, decode_predictions
|
||||
|
||||
def predict(model, path):
|
||||
# only allow valid file types
|
||||
if not (".jpg" in path or ".jpeg" in path):
|
||||
return None
|
||||
|
||||
# receive image path as CLI argument
|
||||
img = load_img(path, color_mode='rgb', target_size=(224, 224))
|
||||
|
||||
@@ -16,4 +20,10 @@ def predict(model, path):
|
||||
features = model.predict(x)
|
||||
p = decode_predictions(features)
|
||||
|
||||
for predict in p:
|
||||
i = 0
|
||||
while i < len(predict):
|
||||
predict[i] = (predict[i][0], predict[i][1], str(predict[i][2]))
|
||||
i = i + 1
|
||||
|
||||
return p
|
||||
|
||||
24
readresult.py
Normal file
24
readresult.py
Normal file
@@ -0,0 +1,24 @@
|
||||
import sys, json
|
||||
|
||||
path = sys.argv[1]
|
||||
|
||||
with open(path) as file:
|
||||
contents = json.load(file)
|
||||
|
||||
for line in contents:
|
||||
prediction = line['prediction']
|
||||
for section in prediction:
|
||||
for guess in section:
|
||||
if (float(guess[2]) > 0.75):
|
||||
print(line['path'])
|
||||
print("Probable match: " + guess[1])
|
||||
print(guess)
|
||||
print("\n")
|
||||
elif (float(guess[2]) > 0.3):
|
||||
print(line['path'])
|
||||
print("Potential match: " + guess[1])
|
||||
print(guess)
|
||||
print("\n")
|
||||
# else:
|
||||
# print(line['path'] + ": inconclusive")
|
||||
# print("\n")
|
||||
Reference in New Issue
Block a user