This commit is contained in:
Mikayla Dobson
2023-01-04 20:24:02 -06:00
commit 8ce90346d9
3 changed files with 51 additions and 0 deletions

7
README.md Normal file
View File

@@ -0,0 +1,7 @@
# Simple Image Recognition Module
A simple implementation of image recognition software using the pretrained TensorFlow VGG16 model and Python.
## Usage
This tool runs from the command line and expects a single argument referring to a relative path. This program will execute within the provided directory, and will attempt to classify each image in the provided directory.

25
main.py Normal file
View File

@@ -0,0 +1,25 @@
# IMAGE RECOGNITION UTIL USING TF/KERAS
#
# 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
from predict import predict
from keras.applications.vgg16 import VGG16
# 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]
files = os.listdir(path)
# store all results in one list
all_results = []
# 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 })
print(all_results)

19
predict.py Normal file
View File

@@ -0,0 +1,19 @@
import numpy as np
from keras.utils import load_img, img_to_array
from keras.applications.vgg16 import preprocess_input, decode_predictions
def predict(model, path):
# receive image path as CLI argument
img = load_img(path ,color_mode='rgb', target_size=(224, 224))
# loaded image to np array for model to read
x = img_to_array(img)
x.shape
x = np.expand_dims(x, axis=0)
# process array and make predictions
x = preprocess_input(x)
features = model.predict(x)
p = decode_predictions(features)
return p