Update app.py

app.py

@@ -63,16 +63,52 @@ def plot_confusion_matrix_beamPred(cm, classes, title, save_path):
     plt.close()
     
 # Function to compute the average confusion matrix across CSV files in a folder
+#def compute_average_confusion_matrix(folder):
+#    confusion_matrices = []
+#   for file in os.listdir(folder):
+#        if file.endswith(".csv"):
+#            data = pd.read_csv(os.path.join(folder, file))
+#            y_true = data["Target"]
+#            y_pred = data["Top-1 Prediction"]
+#            num_labels = len(np.unique(y_true))
+#            cm = confusion_matrix(y_true, y_pred, labels=np.arange(num_labels))
+#            confusion_matrices.append(cm)
+#
+#    if confusion_matrices:
+#        avg_cm = np.mean(confusion_matrices, axis=0)
+#        return avg_cm
+#    else:
+#        return None
+
 def compute_average_confusion_matrix(folder):
     confusion_matrices = []
+    max_num_labels = 0
+
+    # First pass to determine the maximum number of labels
+    for file in os.listdir(folder):
+        if file.endswith(".csv"):
+            data = pd.read_csv(os.path.join(folder, file))
+            num_labels = len(np.unique(data["Target"]))
+            max_num_labels = max(max_num_labels, num_labels)
+
+    # Second pass to calculate the confusion matrices and pad if necessary
     for file in os.listdir(folder):
         if file.endswith(".csv"):
             data = pd.read_csv(os.path.join(folder, file))
             y_true = data["Target"]
             y_pred = data["Top-1 Prediction"]
             num_labels = len(np.unique(y_true))
-            cm = confusion_matrix(y_true, y_pred, labels=np.arange(num_labels))
-            confusion_matrices.append(cm)
+            
+            # Compute confusion matrix
+            cm = confusion_matrix(y_true, y_pred, labels=np.arange(max_num_labels))
+
+            # If the confusion matrix is smaller, pad it to match the largest size
+            if cm.shape[0] < max_num_labels:
+                padded_cm = np.zeros((max_num_labels, max_num_labels))
+                padded_cm[:cm.shape[0], :cm.shape[1]] = cm
+                confusion_matrices.append(padded_cm)
+            else:
+                confusion_matrices.append(cm)
 
     if confusion_matrices:
         avg_cm = np.mean(confusion_matrices, axis=0)
@@ -84,6 +120,7 @@ def compute_average_confusion_matrix(folder):
 
 
 
+
 # Custom class to capture print output
 class PrintCapture(io.StringIO):
     def __init__(self):