Update app.py

app.py

@@ -34,28 +34,27 @@ def create_model():
     model.compile(optimizer="adam", loss="categorical_crossentropy", metrics=["accuracy"])
     return model
 
+# Streamlit UI
+st.title("CNN for MNIST Classification")
+
+# Check if model is saved
+model_path = "mnist_cnn_model.h5"
+
 # Custom callback for logging
 class StreamlitLogger(keras.callbacks.Callback):
     def on_epoch_end(self, epoch, logs=None):
-        if logs is None:
+        if logs is none:
             logs = {}
-        
+
         st.write(f"Epoch {epoch + 1}:")
         st.write(f"  Train Loss: {logs.get('loss'):.4f}")
         st.write(f"  Train Accuracy: {logs.get('accuracy'):.4f}")
         st.write(f"  Val Loss: {logs.get('val_loss'):.4f}")
         st.write(f"  Val Accuracy: {logs.get('val_accuracy'):.4f}")
 
-# 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()
     
-    # Create logger instance
     logger = StreamlitLogger()
 
     with st.spinner("Training..."):
@@ -72,9 +71,44 @@ if st.button("Train Model"):
     ax1.set_title("Training and Validation Loss")
     ax1.set_xlabel("Epoch")
     ax1.set_ylabel("Loss")
-    ax1.legend()
 
     ax2.plot(history.history["accuracy"], label="Train Accuracy")
     ax2.plot(history.history["val_accuracy"], label="Val Accuracy")
     ax2.set_title("Training and Validation Accuracy")
-    ax2.set_xlabel
+    ax2.set_xlabel("Epoch")
+    ax2.set_ylabel("Accuracy")
+
+    ax1.legend()
+    ax2.legend()
+
+    st.pyplot(fig)
+
+    # Evaluate the model on test data
+    test_preds = np.argmax(model.predict(test_images), axis=1)
+    true_labels = np.argmax(test_labels, axis=1)
+
+    st.session_state['true_labels'] = true_labels
+
+    report = classification_report(true_labels, test_preds, digits=4)
+    st.text("Classification Report:")
+    st.text(report)
+
+index = st.number_input("Enter an index (0-9999) to test:", min_value=0, max_value=9999, step=1)
+
+def test_index_prediction(index):
+    image = test_images[index].reshape(28, 28)
+    st.image(image, caption=f"True Label: {st.session_state['true_labels'][index]}", use_column_width=True)
+
+    # Reload the model
+    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"):
+    test_index_prediction(index)