Update app.py

app.py

@@ -40,25 +40,10 @@ 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:
-            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}")
-
 if st.button("Train Model"):
     model = create_model()
-    
-    logger = StreamlitLogger()
-
     with st.spinner("Training..."):
-        history = model.fit(train_images, train_labels, validation_data=(test_images, test_labels), epochs=10, batch_size=64, callbacks=[logger])
+        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)
@@ -71,14 +56,13 @@ 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("Epoch")
     ax2.set_ylabel("Accuracy")
-
-    ax1.legend()
     ax2.legend()
 
     st.pyplot(fig)
@@ -87,19 +71,22 @@ if st.button("Train Model"):
     test_preds = np.argmax(model.predict(test_images), axis=1)
     true_labels = np.argmax(test_labels, axis=1)
 
+    # Store the test labels globally for later use
     st.session_state['true_labels'] = true_labels
 
+    # Classification report
     report = classification_report(true_labels, test_preds, digits=4)
     st.text("Classification Report:")
     st.text(report)
 
+# Testing with a specific index
 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
+    # Reload the model if needed
     if not os.path.exists(model_path):
         st.error("Train the model first.")
         return