Spaces:

kerjaaa
/

trash

No application file

trash / src /predict_image.py

Synced repo using 'sync_with_huggingface' Github Action

9076fcb verified almost 2 years ago

2.4 kB

# src/predict_image.py

import tensorflow as tf # type: ignore
import numpy as np
from tensorflow.keras.preprocessing import image # type: ignore
import matplotlib.pyplot as plt
import sys

# Load the saved model
def load_model(model_path='simpan_resnet_model.h5'):
    try:
        model = tf.keras.models.load_model(model_path)
        print(f"Loaded model from {model_path}")
    except Exception as e:
        print(f"Error loading model: {e}")
        sys.exit(1)
    return model

# Preprocess the image for prediction
def preprocess_image(img_path, target_size=(224, 224)):
    img = image.load_img(img_path, target_size=target_size)  # Load and resize image
    img_array = image.img_to_array(img)                      # Convert to array
    img_array = np.expand_dims(img_array, axis=0)            # Add batch dimension
    img_array /= 255.0                                       # Normalize
    return img_array

# Predict the class of the image
def predict_image(model, img_path, class_names):
    img_array = preprocess_image(img_path)
    predictions = model.predict(img_array)
    predicted_class = np.argmax(predictions, axis=1)[0]      # Get the predicted class index
    confidence = predictions[0][predicted_class]             # Confidence score for predicted class
    return class_names[predicted_class], confidence

# Plot the image with prediction
def display_prediction(img_path, predicted_class, confidence):
    img = image.load_img(img_path)
    plt.imshow(img)
    plt.title(f"Predicted: {predicted_class} ({confidence:.2f})")
    plt.axis('off')
    plt.show()

# Main function for running the prediction
def main(img_path, model_path='simpan_resnet_model.h5'):
    # Define the class names (modify as per your dataset)
    class_names = ['Class 0', 'Class 1', 'Class 2', 'Class 3', 'Class 4', 'Class 5']

    # Load the model
    model = load_model(model_path)

    # Predict and display the result
    predicted_class, confidence = predict_image(model, img_path, class_names)
    print(f"Predicted Class: {predicted_class}, Confidence: {confidence:.4f}")
    display_prediction(img_path, predicted_class, confidence)

if __name__ == "__main__":
    # Ensure the image path is passed as a command-line argument
    if len(sys.argv) < 2:
        print("Usage: python predict_image.py <path_to_image>")
        sys.exit(1)
    
    img_path = sys.argv[1]
    main(img_path)

	# src/predict_image.py

	import tensorflow as tf # type: ignore
	import numpy as np
	from tensorflow.keras.preprocessing import image # type: ignore
	import matplotlib.pyplot as plt
	import sys

	# Load the saved model
	def load_model(model_path='simpan_resnet_model.h5'):
	try:
	model = tf.keras.models.load_model(model_path)
	print(f"Loaded model from {model_path}")
	except Exception as e:
	print(f"Error loading model: {e}")
	sys.exit(1)
	return model

	# Preprocess the image for prediction
	def preprocess_image(img_path, target_size=(224, 224)):
	img = image.load_img(img_path, target_size=target_size) # Load and resize image
	img_array = image.img_to_array(img) # Convert to array
	img_array = np.expand_dims(img_array, axis=0) # Add batch dimension
	img_array /= 255.0 # Normalize
	return img_array

	# Predict the class of the image
	def predict_image(model, img_path, class_names):
	img_array = preprocess_image(img_path)
	predictions = model.predict(img_array)
	predicted_class = np.argmax(predictions, axis=1)[0] # Get the predicted class index
	confidence = predictions[0][predicted_class] # Confidence score for predicted class
	return class_names[predicted_class], confidence

	# Plot the image with prediction
	def display_prediction(img_path, predicted_class, confidence):
	img = image.load_img(img_path)
	plt.imshow(img)
	plt.title(f"Predicted: {predicted_class} ({confidence:.2f})")
	plt.axis('off')
	plt.show()

	# Main function for running the prediction
	def main(img_path, model_path='simpan_resnet_model.h5'):
	# Define the class names (modify as per your dataset)
	class_names = ['Class 0', 'Class 1', 'Class 2', 'Class 3', 'Class 4', 'Class 5']

	# Load the model
	model = load_model(model_path)

	# Predict and display the result
	predicted_class, confidence = predict_image(model, img_path, class_names)
	print(f"Predicted Class: {predicted_class}, Confidence: {confidence:.4f}")
	display_prediction(img_path, predicted_class, confidence)

	if __name__ == "__main__":
	# Ensure the image path is passed as a command-line argument
	if len(sys.argv) < 2:
	print("Usage: python predict_image.py <path_to_image>")
	sys.exit(1)

	img_path = sys.argv[1]
	main(img_path)