init
This commit is contained in:
7
README.md
Normal file
7
README.md
Normal 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
25
main.py
Normal 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
19
predict.py
Normal 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
|
||||
Reference in New Issue
Block a user