File size: 3,320 Bytes
848ad3f 3391b54 848ad3f 3391b54 848ad3f | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | import gradio as gr
import tensorflow as tf
from tensorflow import keras
import numpy as np
from PIL import Image
# --- Configuration ---
MODEL_PATH = "cats-vs-dogs-finetuned.keras"
IMAGE_SIZE = (150, 150) # Adjust this to match the input size your model expects!
CLASS_LABELS = ['Cat', 'Dog']
# --- Load the Model ---
# We load the Keras model. Hugging Face Spaces will automatically find this file
# if you upload it to your repository.
try:
model = keras.models.load_model(MODEL_PATH)
print(f"Model loaded successfully from {MODEL_PATH}")
except Exception as e:
# If the model fails to load (e.g., during initial setup before it's uploaded),
# we use a placeholder function. This helps the app start.
print(f"Error loading model: {e}. Using a placeholder function.")
model = None
# --- Prediction Function ---
def predict_image(input_img_pil):
"""
Predicts the class (Cat or Dog) given a PIL Image object.
Args:
input_img_pil: A PIL Image object received from Gradio's Image input.
Returns:
A dictionary of class labels and their probabilities (for Gradio's Label output).
"""
if model is None:
# Placeholder behavior if model loading failed
return {"Error": 1.0}
# 1. Preprocessing: Resize and convert to NumPy array
img_resized = input_img_pil.resize(IMAGE_SIZE)
print("image resized")
img_array = keras.preprocessing.image.img_to_array(img_resized)
print(" image converted to array")
# 2. Rescaling and Batch dimension:
# Keras models usually expect input shapes like (Batch_Size, Height, Width, Channels)
# and often expect pixel values to be normalized (e.g., 0-1 range).
# Please adjust the normalization based on how your model was trained!
img_array = img_array / 255.0 # Common normalization step
img_array = np.expand_dims(img_array, axis=0) # Add batch dimension
# 3. Prediction
predictions = model.predict(img_array)[0] # Get the single prediction result
# 4. Format the output for Gradio's Label component
# The output is expected to be a dictionary: {'label': probability, ...}
# Assuming predictions is a 2-element array: [prob_cat, prob_dog]
output_dict = {
CLASS_LABELS[0]: float(predictions[0]),
CLASS_LABELS[1]: float(predictions[1])
}
return output_dict
# --- Gradio Interface Setup ---
# Define the input component (Image) and output component (Label)
image_input = gr.Image(type="pil", label="Upload a Cat or Dog Image")
label_output = gr.Label(num_top_classes=2, label="Prediction")
# Example images for users to try (place these in your Space if you use them)
examples = [
# To use these, you would need to upload files named 'example_cat.jpg' and 'example_dog.jpg'
# 'example_cat.jpg',
# 'example_dog.jpg'
]
# Create the Gradio interface
demo = gr.Interface(
fn=predict_image,
inputs=image_input,
outputs=label_output,
title="Keras Cat vs Dog Classifier",
description="Upload an image of a cat or dog to see the model's prediction. The model is loaded from cat-vs-dog.keras.",
theme=gr.themes.Soft(),
# Optional: Add examples if you upload them
# examples=examples
)
# Launch the app
if __name__ == "__main__":
demo.launch()
|