Remove post-training TensorBoard log charts

app.py

@@ -237,7 +237,6 @@ def main():
             # Simulate predictions and labels for confusion matrix
             y_true = np.random.randint(0, num_classes, 100)
             y_pred = y_true.copy()
-            # Add some noise to predictions
             y_pred[np.random.choice(100, 10, replace=False)] = np.random.randint(0, num_classes, 10)
             cm = confusion_matrix(y_true, y_pred, labels=range(num_classes))
 
@@ -249,13 +248,15 @@ def main():
                 plt.close(fig)
                 writer.add_figure("confusion_matrix", fig, epoch)
 
-            time.sleep(0.5)
-
         writer.close()
         st.success("Training complete!")
 
-        # Wait a moment to ensure logs are written
-        time.sleep(1)
+        # Show last confusion matrix in Streamlit
+        st.subheader("Confusion Matrix (Last Epoch)")
+        fig, ax = plt.subplots()
+        disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=[f"Class {i}" for i in range(num_classes)])
+        disp.plot(ax=ax)
+        st.pyplot(fig)
 
 if __name__ == "__main__":
     main()