Update app.py

app.py

@@ -3,48 +3,43 @@ import numpy as np
 from keras.preprocessing import image
 from keras.models import load_model
 
-# Define a dictionary of classes
-classes = {'french_bulldog': 0,
-           'german_shepherd': 1,
-           'golden_retriever': 2,
-           'poodle': 3,
-           'yorkshire_terrier': 4}
 
-# Load the saved model from the disk
-model = load_model('best_model.h5')
-
-# Define the function for predicting the dog breed
-def predict_dog_breed(image_path):
-    # Load the image from the specified path and preprocess it
-    img = image.load_img(image_path, target_size=(256, 256))
+# 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.
-
-    # Make the prediction using the trained model
     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 the predicted dog breed
+    predicted_class = [k for k, v in CLASSES.items() if v == class_idx][0]
     return predicted_class
 
-# Define the Streamlit app
-def app():
-    st.title("Dog Breed Classification App")
-    st.write("Upload an image of a dog to predict its breed.")
 
-    # Allow the user to upload an image file
-    uploaded_file = st.file_uploader("Choose a dog image...", type=["jpg", "jpeg", "png"])
-
-    # Make the prediction when the user clicks the "Predict" button
+# 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_path = f"tmp/{uploaded_file.name}"
-        with open(image_path, "wb") as f:
+        image_file = uploaded_file.name
+        with open(image_file, 'wb') as f:
             f.write(uploaded_file.getbuffer())
-        predicted_class = predict_dog_breed(image_path)
-        st.write(f"Predicted dog breed: {predicted_class}")
-
-# Run the Streamlit app
+        predicted_class = predict_breed(image_file)
+        st.image(uploaded_file, caption=f'Predicted class: {predicted_class}', use_column_width=True)
+        
+        
 if __name__ == '__main__':
-    app()
+    main()