Update pages/13_FFNN.py

pages/13_FFNN.py

@@ -5,6 +5,7 @@ import torch.optim as optim
 import torchvision
 import torchvision.transforms as transforms
 import matplotlib.pyplot as plt
+import seaborn as sns
 import pandas as pd
 import numpy as np
 
@@ -77,7 +78,14 @@ def test_network(net, testloader):
 # Load the data
 trainloader, testloader = load_data()
 
-# Streamlit sidebar for input parameters
+# Streamlit interface
+st.title("Feedforward Neural Network for MNIST Classification")
+
+st.write("""
+This application demonstrates how to build and train a Feedforward Neural Network (FFNN) for image classification using the MNIST dataset. You can adjust hyperparameters, visualize sample images, and see the model's performance.
+""")
+
+# Sidebar for input parameters
 st.sidebar.header('Model Hyperparameters')
 hidden1_size = st.sidebar.slider('Hidden Layer 1 Size', 128, 1024, 512)
 hidden2_size = st.sidebar.slider('Hidden Layer 2 Size', 128, 1024, 256)
@@ -86,6 +94,25 @@ learning_rate = st.sidebar.slider('Learning Rate', 0.001, 0.1, 0.01, step=0.001)
 momentum = st.sidebar.slider('Momentum', 0.0, 1.0, 0.9, step=0.1)
 epochs = st.sidebar.slider('Epochs', 1, 20, 5)
 
+# Display some sample images
+st.write("## Sample Images from MNIST Dataset")
+sample_images, sample_labels = next(iter(trainloader))
+fig, axes = plt.subplots(1, 6, figsize=(15, 5))
+for i in range(6):
+    axes[i].imshow(sample_images[i].numpy().squeeze(), cmap='gray')
+    axes[i].set_title(f'Label: {sample_labels[i].item()}')
+    axes[i].axis('off')
+st.pyplot(fig)
+
+# Class distribution
+st.write("## Class Distribution in MNIST Dataset")
+class_counts = np.bincount(sample_labels.numpy())
+fig, ax = plt.subplots()
+sns.barplot(x=list(range(10)), y=class_counts, ax=ax)
+ax.set_ylabel('Count')
+ax.set_title('Class Distribution')
+st.pyplot(fig)
+
 # Create the network
 net = FeedforwardNeuralNetwork(784, hidden1_size, hidden2_size, hidden3_size, 10)
 criterion = nn.CrossEntropyLoss()
@@ -116,7 +143,7 @@ if 'trained_model' in st.session_state and st.sidebar.button('Test Network'):
     st.write(f'Test Accuracy: {accuracy:.2f}%')
     
     # Display results in a table
-    st.write('GroundTruth vs Predicted')
+    st.write('Ground Truth vs Predicted')
     results = pd.DataFrame({
         'Ground Truth': all_labels,
         'Predicted': all_predicted
@@ -138,9 +165,9 @@ if 'trained_model' in st.session_state and st.sidebar.button('Show Test Results'
     outputs = st.session_state['trained_model'](images)
     _, predicted = torch.max(outputs, 1)
     
-    st.write('GroundTruth vs Predicted')
+    st.write('Ground Truth vs Predicted')
     results = pd.DataFrame({
         'Ground Truth': labels.numpy(),
         'Predicted': predicted.numpy()
     })
-    st.table(results)
+    st.table(results.head(50))  # Display first 50 results for brevity