Spaces:
Runtime error
Runtime error
import streamlit as st | |
import numpy as np | |
from tensorflow.keras.preprocessing import image | |
from keras.models import load_model | |
# Define the dictionary of classes and load the model | |
CLASSES = { | |
'french_bulldog': 0, | |
'german_shepherd': 1, | |
'golden_retriever': 2, | |
'poodle': 3, | |
'yorkshire_terrier': 4 | |
} | |
MODEL_PATH = 'best_model.h5' | |
model = load_model(MODEL_PATH) | |
# Define a function to make predictions on a given image | |
def predict_breed(image_file): | |
img = image.load_img(image_file, target_size=(256, 256)) | |
x = image.img_to_array(img) | |
x = np.expand_dims(x, axis=0) | |
x = x / 255. | |
preds = model.predict(x) | |
class_idx = np.argmax(preds[0]) | |
predicted_class = [k for k, v in CLASSES.items() if v == class_idx][0] | |
return predicted_class | |
# Create the Streamlit app | |
def main(): | |
st.title('Dog Breed Classification') | |
uploaded_file = st.file_uploader('Choose an image of a dog', type=['jpg', 'jpeg', 'png']) | |
if uploaded_file is not None: | |
image_file = uploaded_file.name | |
with open(image_file, 'wb') as f: | |
f.write(uploaded_file.getbuffer()) | |
predicted_class = predict_breed(image_file) | |
st.image(uploaded_file, caption=f'Predicted class: {predicted_class}', use_column_width=True) | |
# Add a button to trigger image upload | |
if st.button('Upload Image'): | |
uploaded_file = st.file_uploader('Choose an image of a dog', type=['jpg', 'jpeg', 'png']) | |
if uploaded_file is not None: | |
image_file = uploaded_file.name | |
with open(image_file, 'wb') as f: | |
f.write(uploaded_file.getbuffer()) | |
predicted_class = predict_breed(image_file) | |
st.image(uploaded_file, caption=f'Predicted class: {predicted_class}', use_column_width=True) | |
if __name__ == '__main__': | |
main() | |