Update app.py

app.py

@@ -321,35 +321,44 @@ with col_right:
             self.X, self.y = X, y
             self.losses = {"Epoch": [], "Train Loss": [], "Val Loss": []}
             self.placeholder = st.empty()
+            self.current_epoch = 0  # Track current epoch
 
         def on_train_begin(self, logs=None):
             self.model = self.model  # Use the model passed implicitly by Keras
+            self.current_epoch = 0
 
         def on_epoch_end(self, epoch, logs=None):
             try:
-                self.losses["Epoch"].append(epoch + 1)
+                self.current_epoch = epoch + 1  # Update current epoch
+                self.losses["Epoch"].append(self.current_epoch)
                 self.losses["Train Loss"].append(logs["loss"])
                 self.losses["Val Loss"].append(logs.get("val_loss", logs["loss"]))
                 with self.placeholder.container():
                     # Single column for vertical stacking
                     st.subheader("Decision Region & Loss")
-                    # Decision region plot (3x3 size)
+                    # Display epoch count above decision region
+                    st.write(f"Epoch: {self.current_epoch}")
+                    # Decision region plot (3x3 size, improved accuracy)
                     fig1, ax1 = plt.subplots(figsize=(3, 3))  # Match dataset scatterplot size
                     if problem_type == "Classification":
                         X_2d = self.X[:, :2]  # Use only first two features for 2D
-                        y_pred_proba = self.model.predict(X_2d, verbose=0) if self.model else np.zeros((len(X_2d), 1))
+                        # Ensure model prediction for decision boundary
+                        y_pred_proba = self.model.predict(X_2d, verbose=0)
                         y_pred = (y_pred_proba > 0.5).astype(int).ravel()
                         try:
+                            # Use mlxtend for decision regions
                             plot_decision_regions(X_2d, self.y, clf=self.model, legend=2, colors='blue,red')
                             plt.scatter(X_2d[:, 0], X_2d[:, 1], c=self.y, cmap='coolwarm', edgecolors='k', alpha=0.7)
+                            # Add precise decision boundary using contour
                             xx, yy = np.meshgrid(np.linspace(X_2d[:, 0].min(), X_2d[:, 0].max(), 100),
                                                  np.linspace(X_2d[:, 1].min(), X_2d[:, 1].max(), 100))
                             grid = np.c_[xx.ravel(), yy.ravel()]
-                            Z = self.model.predict(grid, verbose=0) if self.model else np.zeros((len(grid), 1))
+                            Z = self.model.predict(grid, verbose=0)
                             Z = (Z > 0.5).astype(int).reshape(xx.shape)
                             plt.contour(xx, yy, Z, levels=[0.5], colors='black', linewidths=2)
                         except Exception as e:
                             st.warning(f"Decision region plot failed: {e}")
+                            # Fallback: Use contourf for decision regions
                             xx, yy = np.meshgrid(np.linspace(X_2d[:, 0].min(), X_2d[:, 0].max(), 100),
                                                  np.linspace(X_2d[:, 1].min(), X_2d[:, 1].max(), 100))
                             grid = np.c_[xx.ravel(), yy.ravel()]