Update app.py

app.py

@@ -34,6 +34,18 @@ def create_model():
     model.compile(optimizer="adam", loss="categorical_crossentropy", metrics=["accuracy"])
     return model
 
+# 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}")
+
 # Streamlit UI
 st.title("CNN for MNIST Classification")
 
@@ -42,8 +54,12 @@ model_path = "mnist_cnn_model.h5"
 
 if st.button("Train Model"):
     model = create_model()
+    
+    # Create logger instance
+    logger = StreamlitLogger()
+
     with st.spinner("Training..."):
-        history = model.fit(train_images, train_labels, validation_data=(test_images, test_labels), epochs=10, batch_size=64)
+        history = model.fit(train_images, train_labels, validation_data=(test_images, test_labels), epochs=10, batch_size=64, callbacks=[logger])
 
     # Save the model
     model.save(model_path)
@@ -61,38 +77,4 @@ if st.button("Train Model"):
     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")
-    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)
-
-    # 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: {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"):
-    test_index_prediction(index)
+    ax2.set_xlabel