Confusion matrix plot

app.py

@@ -59,6 +59,7 @@ MODEL_INFO = [
 ]
 
 DEMO_LEARNING_RATE = 0.05 # don't use default; use something more fun
+DEMO_ALPHA = 0.25
 
 # load image metadata
 images_data = []
@@ -305,51 +306,116 @@ def create_probability_chart():
     return temp_fig
 
 def create_accuracy_chart():
-    """Create a bar chart showing true accuracy of each model"""
-    global oracle, dataset
+    """Create confusion matrix estimates for each model side by side"""
+    global coda_selector, iteration_count
 
-    if oracle is None or dataset is None:
-        # Fallback for initial state
-        model_labels = [info['name'] for info in MODEL_INFO]
-        accuracies = np.random.random(len(MODEL_INFO))  # Random accuracies for now
-    else:
-        true_losses = oracle.true_losses(dataset.preds)
-        # Convert losses to accuracies (assuming loss is 1 - accuracy)
-        accuracies = (1 - true_losses).detach().cpu().numpy().flatten()
-        model_labels = ["   " + info['name'] for info in MODEL_INFO[:len(accuracies)]]
-
-    # Find the index of the highest accuracy
-    best_idx = np.argmax(accuracies)
-
-    fig, ax = plt.subplots(figsize=(8, 2.8), dpi=150)
-
-    # Create colors array - highlight the best model
-    colors = ['red' if i == best_idx else 'forestgreen' for i in range(len(model_labels))]
-    bars = ax.bar(range(len(model_labels)), accuracies, color=colors, alpha=0.7)
-
-    # Add text above the highest bar
-    ax.text(best_idx, accuracies[best_idx] + 0.005, 'True best model',
-            ha='center', va='bottom', fontsize=12, fontweight='bold')
+    if coda_selector is None:
+        # Fallback for initial state - return empty figure
+        fig, ax = plt.subplots(figsize=(8, 2.8), dpi=150)
+        ax.text(0.5, 0.5, 'Start demo to see confusion matrices',
+                ha='center', va='center', fontsize=12)
+        ax.axis('off')
+        plt.tight_layout()
+        temp_fig = fig
+        plt.close(fig)
+        return temp_fig
+
+    # Get confusion matrix estimates from CODA's Dirichlet distributions
+    dirichlets = coda_selector.dirichlets  # Shape: [num_models, num_classes, num_classes]
+    num_models = dirichlets.shape[0]
+    num_classes = dirichlets.shape[1]
+
+    # Convert Dirichlet parameters to expected confusion matrices
+    # The expected value of a Dirichlet is alpha / sum(alpha)
+    confusion_matrices = []
+    for model_idx in range(num_models):
+        alpha = dirichlets[model_idx].detach().cpu().numpy()
+        # Normalize each row to get probabilities
+        conf_matrix = alpha / alpha.sum(axis=1, keepdims=True)
+        confusion_matrices.append(conf_matrix)
+
+    # Create subplots for each model
+    fig, axes = plt.subplots(1, num_models, figsize=(8, 2.8), dpi=150)
+    if num_models == 1:
+        axes = [axes]
+
+    for model_idx, (ax, conf_matrix) in enumerate(zip(axes, confusion_matrices)):
+        # Apply square root scaling to make small values more visible
+        # This expands small values while still showing large values
+        sqrt_conf_matrix = np.sqrt(np.sqrt(np.sqrt(np.sqrt(conf_matrix))))
+
+        # Plot confusion matrix as heatmap with sqrt-scaled values
+        im = ax.imshow(sqrt_conf_matrix, cmap='Blues', aspect='auto')#, vmin=0, vmax=1)
+
+        # Add model name as title
+        model_info = MODEL_INFO[model_idx]
+        ax.set_title(f"{model_info['name']}", fontsize=10, pad=5)
 
-    ax.set_ylabel('True (oracle) \naccuracy of model', fontsize=12)
-    ax.set_title('True Model Accuracies', fontsize=12)
-    ax.set_ylim(np.min(accuracies) - 0.025, np.max(accuracies) + 0.05)
+        # Set axis labels
+        if model_idx == 0:
+            ax.set_ylabel('True class', fontsize=9)
+        ax.set_xlabel('Predicted', fontsize=9)
 
-    # Set x-axis labels and ticks
-    ax.set_xticks(range(len(model_labels)))
-    ax.set_xticklabels(model_labels, fontsize=12, ha='center')
+        # Set ticks
+        ax.set_xticks(range(num_classes))
+        ax.set_yticks(range(num_classes))
+        ax.set_xticklabels(range(num_classes), fontsize=8)
+        ax.set_yticklabels(range(num_classes), fontsize=8)
 
-    # Add logos to x-axis
-    for i, model_info in enumerate(MODEL_INFO[:len(accuracies)]):
-        add_logo_to_x_axis(ax, i, model_info['logo'], model_info['name'])
-    plt.yticks(fontsize=12)
+    plt.suptitle(f"CODA's Confusion Matrix Estimates (Iteration {iteration_count})", fontsize=12, y=0.98)
     plt.tight_layout()
 
-    # Save the figure and close it to prevent memory leaks
     temp_fig = fig
     plt.close(fig)
     return temp_fig
 
+# OLD CODE - True Model Accuracies Bar Chart (kept for easy reversion)
+# def create_accuracy_chart():
+#     """Create a bar chart showing true accuracy of each model"""
+#     global oracle, dataset
+#
+#     if oracle is None or dataset is None:
+#         # Fallback for initial state
+#         model_labels = [info['name'] for info in MODEL_INFO]
+#         accuracies = np.random.random(len(MODEL_INFO))  # Random accuracies for now
+#     else:
+#         true_losses = oracle.true_losses(dataset.preds)
+#         # Convert losses to accuracies (assuming loss is 1 - accuracy)
+#         accuracies = (1 - true_losses).detach().cpu().numpy().flatten()
+#         model_labels = ["   " + info['name'] for info in MODEL_INFO[:len(accuracies)]]
+#
+#     # Find the index of the highest accuracy
+#     best_idx = np.argmax(accuracies)
+#
+#     fig, ax = plt.subplots(figsize=(8, 2.8), dpi=150)
+#
+#     # Create colors array - highlight the best model
+#     colors = ['red' if i == best_idx else 'forestgreen' for i in range(len(model_labels))]
+#     bars = ax.bar(range(len(model_labels)), accuracies, color=colors, alpha=0.7)
+#
+#     # Add text above the highest bar
+#     ax.text(best_idx, accuracies[best_idx] + 0.005, 'True best model',
+#             ha='center', va='bottom', fontsize=12, fontweight='bold')
+#
+#     ax.set_ylabel('True (oracle) \naccuracy of model', fontsize=12)
+#     ax.set_title('True Model Accuracies', fontsize=12)
+#     ax.set_ylim(np.min(accuracies) - 0.025, np.max(accuracies) + 0.05)
+#
+#     # Set x-axis labels and ticks
+#     ax.set_xticks(range(len(model_labels)))
+#     ax.set_xticklabels(model_labels, fontsize=12, ha='center')
+#
+#     # Add logos to x-axis
+#     for i, model_info in enumerate(MODEL_INFO[:len(accuracies)]):
+#         add_logo_to_x_axis(ax, i, model_info['logo'], model_info['name'])
+#     plt.yticks(fontsize=12)
+#     plt.tight_layout()
+#
+#     # Save the figure and close it to prevent memory leaks
+#     temp_fig = fig
+#     plt.close(fig)
+#     return temp_fig
+
 # Create the Gradio interface
 with gr.Blocks(title="CODA: Wildlife Photo Classification Challenge", 
                theme=gr.themes.Base(),
@@ -659,9 +725,12 @@ with gr.Blocks(title="CODA: Wildlife Photo Classification Challenge",
                 A shaggy, dark brown antelope recognized by its white rump ring and backward-curving horns in males. Smaller and darker than the common eland, waterbuck prefer wet habitats and lack the eland's throat dewlap.
                 
                 ----
-                            
+
                 """)
 
+            # Add spacing before buttons
+            gr.HTML("<div style='margin-top: 0.1em;'></div>")
+
             with gr.Row():
                 back_button = gr.Button("← Back to Intro", variant="secondary", size="lg", visible=False)
                 guide_button = gr.Button("View Species Classification Guide", variant="secondary", size="lg")
@@ -810,7 +879,8 @@ with gr.Blocks(title="CODA: Wildlife Photo Classification Challenge",
         # Create oracle and CODA selector for this user
         oracle = Oracle(dataset, loss_fn=loss_fn)
         coda_selector = CODA(dataset, 
-                             learning_rate=DEMO_LEARNING_RATE)
+                             learning_rate=DEMO_LEARNING_RATE,
+                             alpha=DEMO_ALPHA)
 
         image, status, predictions = get_next_coda_image()
         prob_plot = create_probability_chart()
@@ -849,7 +919,8 @@ with gr.Blocks(title="CODA: Wildlife Photo Classification Challenge",
         # Create oracle and CODA selector for this user
         oracle = Oracle(dataset, loss_fn=loss_fn)
         coda_selector = CODA(dataset,
-                             learning_rate=DEMO_LEARNING_RATE)
+                             learning_rate=DEMO_LEARNING_RATE,
+                             alpha=DEMO_ALPHA)
 
         # Reset all displays
         prob_plot = create_probability_chart()