Update pages/13_FFNN.py

pages/13_FFNN.py

@@ -59,6 +59,8 @@ def train_network(net, trainloader, criterion, optimizer, epochs):
 def test_network(net, testloader):
     correct = 0
     total = 0
+    all_labels = []
+    all_predicted = []
     with torch.no_grad():
         for data in testloader:
             images, labels = data
@@ -66,9 +68,11 @@ def test_network(net, testloader):
             _, predicted = torch.max(outputs.data, 1)
             total += labels.size(0)
             correct += (predicted == labels).sum().item()
+            all_labels.extend(labels.numpy())
+            all_predicted.extend(predicted.numpy())
     accuracy = 100 * correct / total
     st.write(f'Accuracy of the network on the 10000 test images: {accuracy:.2f}%')
-    return accuracy
+    return accuracy, all_labels, all_predicted
 
 # Load the data
 trainloader, testloader = load_data()
@@ -102,11 +106,22 @@ if st.sidebar.button('Train Network'):
     plt.ylabel('Loss')
     plt.grid(True)
     st.pyplot(plt)
+    
+    # Store the trained model in the session state
+    st.session_state['trained_model'] = net
 
 # Test the network
-if st.sidebar.button('Test Network'):
-    accuracy = test_network(net, testloader)
+if 'trained_model' in st.session_state and st.sidebar.button('Test Network'):
+    accuracy, all_labels, all_predicted = test_network(st.session_state['trained_model'], testloader)
     st.write(f'Test Accuracy: {accuracy:.2f}%')
+    
+    # Display results in a table
+    st.write('GroundTruth vs Predicted')
+    results = pd.DataFrame({
+        'Ground Truth': all_labels,
+        'Predicted': all_predicted
+    })
+    st.table(results.head(50))  # Display first 50 results for brevity
 
 # Visualize some test results
 def imshow(img):
@@ -115,12 +130,12 @@ def imshow(img):
     plt.imshow(np.transpose(npimg, (1, 2, 0)))
     plt.show()
 
-if st.sidebar.button('Show Test Results'):
+if 'trained_model' in st.session_state and st.sidebar.button('Show Test Results'):
     dataiter = iter(testloader)
     images, labels = next(dataiter)  # Use next function
     imshow(torchvision.utils.make_grid(images))
     
-    outputs = net(images)
+    outputs = st.session_state['trained_model'](images)
     _, predicted = torch.max(outputs, 1)
     
     st.write('GroundTruth vs Predicted')