hb-setosys's picture
Update app.py
af590ff verified
raw
history blame
2.26 kB
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from tensorflow.keras.applications.efficientnet_v2 import preprocess_input
import numpy as np
from google.colab import files
# Step 1: Upload the model to Colab (run this in a Colab cell)
uploaded = files.upload() # Upload the .h5 model file
# Step 2: Load the trained model
MODEL_PATH = "setosys_dogs_model.h5" # Adjust if needed
model = load_model(MODEL_PATH)
# Step 3: Define the class labels manually (as per your model's training setup)
# You need to know the classes that were used during model training
class_labels = ["Labrador Retriever", "German Shepherd", "Golden Retriever", "Bulldog", "Poodle"] # Example, update this list
# Step 4: Define image preprocessing function for EfficientNetV2
def preprocess_image(image_path):
"""Preprocess the image to match the model's input requirements."""
img = load_img(image_path, target_size=(224, 224)) # Resize image to model input size
img_array = img_to_array(img)
img_array = preprocess_input(img_array) # EfficientNetV2 preprocessing
img_array = np.expand_dims(img_array, axis=0) # Add batch dimension
return img_array
# Step 5: Define prediction function
def predict_dog_breed(image_path):
"""Predict the breed of the dog in the uploaded image."""
img_array = preprocess_image(image_path)
predictions = model.predict(img_array)
# Check the shape of the predictions to make sure the output is correct
print("Predictions Shape:", predictions.shape)
class_idx = np.argmax(predictions) # Index of the highest prediction probability
confidence = float(np.max(predictions)) # Confidence score
# Get predicted breed and its confidence score
predicted_breed = class_labels[class_idx] if class_idx < len(class_labels) else "Unknown"
return predicted_breed, confidence
# Step 6: Upload and test with an image
uploaded_image = files.upload() # Upload a test image
image_path = list(uploaded_image.keys())[0] # Get the filename of the uploaded image
# Step 7: Run prediction
breed, confidence = predict_dog_breed(image_path)
print(f"Predicted Breed: {breed}, Confidence: {confidence:.2f}")