File size: 2,769 Bytes
ffd6b68 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import os
import argparse
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing import image
from tensorflow.keras.models import load_model
def load_and_preprocess_image(img_path, target_size):
# Load and preprocess the image for prediction.
"""Load and preprocess the image for prediction."""
img = image.load_img(img_path, target_size=target_size)
img_array = image.img_to_array(img)
img_array = np.expand_dims(img_array, axis=0) # Create batch axis
img_array = img_array / 255.0 # Normalize the image
return img_array
def load_model_from_file(model_path):
# Load the pre-trained model from the specified path.
"""Load the pre-trained model from the specified path."""
model = load_model(model_path)
print(f"Model loaded from {model_path}")
return model
def make_predictions(model, img_array):
# Make predictions using the loaded model.
"""Make predictions using the loaded model."""
predictions = model.predict(img_array)
return predictions
def get_class_names(train_dir):
"""Get class names from training directory."""
class_names = os.listdir(train_dir) # Assuming subfolder names are the class labels
class_names.sort() # Ensure consistent ordering
return class_names
def main(model_path, img_path, train_dir):
# Main function to load model, preprocess image, make predictions, and display results.
# Define target image size based on model requirements
target_size = (224, 224) # Adjust if needed
# Load the model
model = load_model_from_file(model_path)
# Get class names from train directory
class_names = get_class_names(train_dir)
# Load and preprocess the image
img_array = load_and_preprocess_image(img_path, target_size)
# Make predictions
predictions = make_predictions(model, img_array)
predicted_label_index = np.argmax(predictions, axis=1)[0]
predicted_label = class_names[predicted_label_index]
probability_score = predictions[0][predicted_label_index]
print(f"Predicted label: {predicted_label}, Probability: {probability_score:.4f}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Load a pre-trained model and make a prediction on a new image")
parser.add_argument('--model_path', type=str, required=True, help='Path to the saved model')
parser.add_argument('--img_path', type=str, required=True, help='Path to the image to be predicted')
parser.add_argument('--train_dir', type=str, required=True, help='Directory containing training dataset for inferring class names')
args = parser.parse_args()
main(args.model_path, args.img_path, args.train_dir)
|