inti

app.py

@@ -253,29 +253,25 @@ def main():
         st.success("Training complete!")
         st.line_chart({"Loss": losses, "Val Loss": val_losses, "Accuracy": accuracies})
 
-        # 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)
-
-    logdir = "logs"
-    if os.path.exists(logdir) and os.listdir(logdir):
-        ea = event_accumulator.EventAccumulator(logdir)
-        ea.Reload()
-        scalars = ea.Tags()['scalars']
-        for tag in ['loss', 'val_loss', 'accuracy']:
-            if tag in scalars:
-                values = [s.value for s in ea.Scalars(tag)]
-                st.line_chart({tag: values})
-        # Show confusion matrix images if available
-        if 'confusion_matrix' in ea.Tags()['images']:
-            st.subheader("TensorBoard Confusion Matrices")
-            for img in ea.Images('confusion_matrix'):
-                st.image(img.encoded_image_string)
-    else:
-        st.info("No TensorBoard logs found. Please upload logs to the 'logs' directory.")
+        # Wait a moment to ensure logs are written
+        time.sleep(1)
+
+        logdir = "logs"
+        if os.path.exists(logdir) and os.listdir(logdir):
+            ea = event_accumulator.EventAccumulator(logdir)
+            ea.Reload()
+            scalars = ea.Tags()['scalars']
+            for tag in ['loss', 'val_loss', 'accuracy']:
+                if tag in scalars:
+                    values = [s.value for s in ea.Scalars(tag)]
+                    st.line_chart({tag: values})
+            # Show confusion matrix images if available
+            if 'confusion_matrix' in ea.Tags()['images']:
+                st.subheader("TensorBoard Confusion Matrices")
+                for img in ea.Images('confusion_matrix'):
+                    st.image(img.encoded_image_string)
+        else:
+            st.info("No TensorBoard logs found. Please upload logs to the 'logs' directory.")
 
 if __name__ == "__main__":
     main()