Update app.py

app.py

@@ -5,6 +5,7 @@ from tensorflow import keras
 from tensorflow.keras import layers
 from tensorflow.keras.datasets import mnist
 import streamlit as st
+import os
 
 # Load the MNIST dataset
 (train_images, train_labels), (test_images, test_labels) = mnist.load_data()
@@ -36,11 +37,17 @@ def create_model():
 # Streamlit UI
 st.title("CNN for MNIST Classification")
 
+# Check if model is saved
+model_path = "mnist_cnn_model.h5"
+
 if st.button("Train Model"):
     model = create_model()
     with st.spinner("Training..."):
         history = model.fit(train_images, train_labels, validation_data=(test_images, test_labels), epochs=10, batch_size=64)
-        
+
+    # Save the model
+    model.save(model_path)
+
     # Plot training loss and accuracy
     fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
 
@@ -76,13 +83,16 @@ def test_index_prediction(index):
     image = test_images[index].reshape(28, 28)
     st.image(image, caption=f"True Label: {true_labels[index]}", use_column_width=True)
 
+    # Reload the model if needed
+    if not os.path.exists(model_path):
+        st.error("Train the model first.")
+        return
+
+    model = keras.models.load_model(model_path)
+
     prediction = model.predict(test_images[index].reshape(1, 28, 28, 1))
     predicted_class = np.argmax(prediction)
     st.write(f"Predicted Class: {predicted_class}")
 
 if st.button("Test Index"):
-    if 'model' in locals() and model:
-        test_index_prediction(index)
-    else:
-        st.error("Train the model first.")
-
+    test_index_prediction(index)