hb-setosys's picture
Update app.py
1540714 verified
raw
history blame
3.69 kB
import os
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.applications.efficientnet_v2 import preprocess_input
from tensorflow.keras.models import load_model
import gradio as gr
# List of class labels from the Stanford Dogs dataset
class_labels = [
'Affenpinscher', 'African_hunting_dog', 'Airedale', 'American_Staffordshire_terrier', 'American_water_spaniel',
'Anatolian_shepherd_dog', 'Australian_terrier', 'Basenji', 'Basset_hound', 'Beagle', 'Bearded_collie',
'Beauceron', 'Bedlington_terrier', 'Belgian_malinois', 'Belgian_sheepdog', 'Bernese_mountain_dog',
'Biewer', 'Black-and-tan_coonhound', 'Black_russian_terrier', 'Border_collie', 'Border_terrier',
'Borzoi', 'Boston_bull', 'Bouvier_des_Flandres', 'Boxer', 'Boykin_spaniel', 'Briard', 'Brittany',
'Bull_mastiff', 'Cairn_terrier', 'Canaan_dog', 'Cavalier_king_charles_spaniel', 'Chihuahua',
'Chinese_crested', 'Chow', 'Clumber_spaniel', 'Cocker_spaniel', 'Collie', 'Curly-coated_retriever',
'Dachshund', 'Dalmatian', 'Dandie_Dinmont_terrier', 'Doberman', 'English_cocker_spaniel',
'English_setter', 'English_springer_spaniel', 'Entlebucher_mountain_dog', 'Field_spaniel', 'Finnish_spitz',
'Flat-coated_retriever', 'French_bulldog', 'German_pinscher', 'German_shepherd', 'German_short-haired_pointer',
'Giant_schnauzer', 'Glen_of_imaal_terrier', 'Golden_retriever', 'Goldendoodle', 'Great_dane', 'Great_pyrenees',
'Greater_swiss_mountain_dog', 'Havanese', 'Irish_setter', 'Irish_terrier', 'Irish_water_spaniel',
'Italian_greyhound', 'Japanese_chin', 'Keeshond', 'Kerry_blue_terrier', 'King_charles_spaniel',
'Klee_kai', 'Labrador_retriever', 'Lakeland_terrier', 'Lhasa', 'Maltese_dog', 'Manchester_terrier',
'Mastiff', 'Miniature_pinscher', 'Miniature_schnauzer', 'Newfoundland', 'Norfolk_terrier', 'Norwegian_elkhound',
'Norwich_terrier', 'Old_english_sheepdog', 'Otterhound', 'Papillon', 'Pekingese', 'Pembroke', 'Pharaoh_hound',
'Plott', 'Pointer', 'Pomeranian', 'Poodle', 'Portuguese_water_dog', 'Rottweiler', 'Saint_bernard',
'Saluki', 'Samoyed', 'Schipperke', 'Scotch_terrier', 'Shiba_inu', 'Shih-tzu', 'Siberian_husky', 'Silky_terrier',
'Staffordshire_bullterrier', 'Standard_schnauzer', 'Tibetan_mastiff', 'Tibetan_terrier', 'Weimaraner',
'Welsh_springer_spaniel', 'West_highland_white_terrier', 'Whippet', 'Yorkshire_terrier'
]
# Load the trained model
model = load_model("setosys_dogs_model.h5")
# Preprocessing function
def preprocess_image(image):
# No need to use load_img, the image is already a PIL image from Gradio
img = image.resize((224, 224)) # Resize image to 224x224
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
# Prediction function
def predict_dog_breed(image):
img_array = preprocess_image(image)
predictions = model.predict(img_array)
class_idx = np.argmax(predictions) # Get index of the class with the highest probability
breed = class_labels[class_idx] # Map the index to the breed name
confidence = predictions[0][class_idx] # Confidence score
return [breed, confidence] # Return as a list
# Create a Gradio interface
iface = gr.Interface(fn=predict_dog_breed,
inputs=gr.Image(type="pil"), # Gradio automatically handles PIL images
outputs=gr.JSON(), # Outputs will be a JSON response (list of results)
live=True)
# Launch the Gradio app
iface.launch()