Restore legacy plot outputs (match 42bf0f)

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import os
-from PIL import Image, ImageDraw
+from PIL import Image
 import numpy as np
 import pickle
 import io
@@ -79,53 +79,41 @@ def _ensure_hf_repo_cloned(repo_url: str, repo_dir: str) -> None:
 
     subprocess.run(["git", "clone", "--depth", "1", clone_url, repo_dir], check=True, env=env)
 
-
-def _make_error_image(message: str, size=(420, 420)) -> Image.Image:
-    # A simple fallback to avoid blank/failed Gradio image renders.
-    img = Image.new("RGB", size, color=(255, 255, 255))
-    draw = ImageDraw.Draw(img)
-    draw.text((12, 12), message[:800], fill=(0, 0, 0))
-    return img
-
 #################### BEAM PREDICTION #########################}
 def beam_prediction_task(data_percentage, task_complexity, theme='Dark'):
     # Folder naming convention based on input_type, data_percentage, and task_complexity
     raw_folder = f"images/raw_{data_percentage/100:.1f}_{task_complexity}"
     embeddings_folder = f"images/embedding_{data_percentage/100:.1f}_{task_complexity}"
 
-    # Process raw confusion matrix
-    raw_cm = compute_average_confusion_matrix(raw_folder, num_labels=int(task_complexity))
+    # Process raw confusion matrix (match legacy behavior)
+    raw_cm = compute_average_confusion_matrix(raw_folder)
     if raw_cm is not None:
-        raw_img = plot_confusion_matrix_beamPred(
+        raw_cm_path = os.path.join(raw_folder, "confusion_matrix_raw.png")
+        plot_confusion_matrix_beamPred(
             raw_cm,
             classes=np.arange(raw_cm.shape[0]),
             title=f"Confusion Matrix (Raw Channels)\n{data_percentage}% data, {task_complexity} beams",
-            save_path=None,
+            save_path=raw_cm_path,
             theme=theme,
         )
+        raw_img = Image.open(raw_cm_path)
     else:
-        raw_img = _make_error_image(
-            f"No data found for Raw Channels.\n\n"
-            f"Expected CSVs under: {raw_folder}\n"
-            f"Try a different data percentage / beam count."
-        )
+        raw_img = None
 
-    # Process embeddings confusion matrix
-    embeddings_cm = compute_average_confusion_matrix(embeddings_folder, num_labels=int(task_complexity))
+    # Process embeddings confusion matrix (match legacy behavior)
+    embeddings_cm = compute_average_confusion_matrix(embeddings_folder)
     if embeddings_cm is not None:
-        embeddings_img = plot_confusion_matrix_beamPred(
+        embeddings_cm_path = os.path.join(embeddings_folder, "confusion_matrix_embeddings.png")
+        plot_confusion_matrix_beamPred(
             embeddings_cm,
             classes=np.arange(embeddings_cm.shape[0]),
             title=f"Confusion Matrix (LWM Embeddings)\n{data_percentage}% data, {task_complexity} beams",
-            save_path=None,
+            save_path=embeddings_cm_path,
             theme=theme,
         )
+        embeddings_img = Image.open(embeddings_cm_path)
     else:
-        embeddings_img = _make_error_image(
-            f"No data found for LWM Embeddings.\n\n"
-            f"Expected CSVs under: {embeddings_folder}\n"
-            f"Try a different data percentage / beam count."
-        )
+        embeddings_img = None
 
     return raw_img, embeddings_img
 
@@ -148,7 +136,7 @@ def compute_f1_score(cm):
     f1 = np.nan_to_num(f1)  # Replace NaN with 0
     return np.mean(f1)  # Return the mean F1-score across all classes
 
-def plot_confusion_matrix_beamPred(cm, classes, title, save_path=None, theme='Dark'):
+def plot_confusion_matrix_beamPred(cm, classes, title, save_path, theme='Dark'):
     # Compute the average F1-score
     avg_f1 = compute_f1_score(cm)
 
@@ -185,67 +173,47 @@ def plot_confusion_matrix_beamPred(cm, classes, title, save_path=None, theme='Da
     plt.xlabel('Predicted label', color=text_color, fontsize=20)
     plt.tight_layout()
 
-    buf = io.BytesIO()
-    plt.savefig(buf, format="png", transparent=True)
+    plt.savefig(save_path, transparent=True)  # Transparent to blend with the site background
     plt.close()
-    buf.seek(0)
-    return Image.open(buf)
 
-def compute_average_confusion_matrix(folder, num_labels=None):
-    """Average confusion matrices across CSVs in a folder.
-
-    For Beam Prediction, the correct class set is defined by the selected beam count
-    (task_complexity). Inferring the number of labels from unique targets can shrink
-    the matrix when some beams never appear in a particular split.
-    """
-
-    if not os.path.isdir(folder):
-        return None
-
-    csv_files = [f for f in os.listdir(folder) if f.endswith(".csv")]
-    if not csv_files:
-        return None
-
-    # If num_labels isn't specified, infer from max label index seen.
-    if num_labels is None:
-        inferred = 0
-        for file in csv_files:
-            data = pd.read_csv(os.path.join(folder, file))
-            y_true = pd.to_numeric(data.get("Target"), errors="coerce").fillna(-1).astype(int)
-            y_pred = pd.to_numeric(data.get("Top-1 Prediction"), errors="coerce").fillna(-1).astype(int)
-            true_max = int(y_true.max()) if len(y_true) else -1
-            pred_max = int(y_pred.max()) if len(y_pred) else -1
-            local_max = max(true_max, pred_max)
-            inferred = max(inferred, local_max + 1)
-        num_labels = inferred
-
-    try:
-        num_labels = int(num_labels)
-    except Exception:
-        return None
-    if num_labels <= 0:
-        return None
+    # Return the saved image
+    return Image.open(save_path)
 
+def compute_average_confusion_matrix(folder):
     confusion_matrices = []
-    for file in csv_files:
-        data = pd.read_csv(os.path.join(folder, file))
-        y_true = pd.to_numeric(data.get("Target"), errors="coerce").fillna(-1).astype(int)
-        y_pred = pd.to_numeric(data.get("Top-1 Prediction"), errors="coerce").fillna(-1).astype(int)
-
-        # Filter invalid/out-of-range labels to avoid skew.
-        valid = (y_true >= 0) & (y_true < num_labels) & (y_pred >= 0) & (y_pred < num_labels)
-        y_true = y_true[valid]
-        y_pred = y_pred[valid]
-
-        if len(y_true) == 0:
-            continue
+    max_num_labels = 0
 
-        cm = confusion_matrix(y_true, y_pred, labels=np.arange(num_labels))
-        confusion_matrices.append(cm)
+    # 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)
 
-    if not confusion_matrices:
+    # 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))
+
+            # 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)
+        return avg_cm
+    else:
         return None
-    return np.mean(confusion_matrices, axis=0)
 
 ########################## LOS/NLOS CLASSIFICATION #############################3
 
@@ -305,11 +273,12 @@ def plot_confusion_matrix_from_csv(csv_file_path, title, save_path, light_mode=F
     plt.xlabel('Predicted label', color=text_color, fontsize=14)
     plt.tight_layout()
 
-    buf = io.BytesIO()
-    plt.savefig(buf, format="png", transparent=True)
+    # Save the plot as an image
+    plt.savefig(save_path, transparent=True)  # Use transparent to blend with the website
     plt.close()
-    buf.seek(0)
-    return Image.open(buf)
+
+    # Return the saved image
+    return Image.open(save_path)
 
 # Function to load confusion matrix based on percentage and input_type
 def display_confusion_matrices_los(percentage):
@@ -322,28 +291,20 @@ def display_confusion_matrices_los(percentage):
     # Process raw confusion matrix
     raw_csv_file = os.path.join(raw_folder, f"test_predictions_raw_{percentage/100:.3f}_los.csv")
     raw_cm_img_path = os.path.join(raw_folder, "confusion_matrix_raw.png")
-    try:
-        raw_img = plot_confusion_matrix_from_csv(
-            raw_csv_file,
-            f"Confusion Matrix (Raw Channels)\n{percentage:.1f}% data",
-            raw_cm_img_path,
-        )
-    except Exception as exc:
-        raw_img = _make_error_image(f"Failed to load Raw CSV:\n{raw_csv_file}\n\n{exc}")
+    raw_img = plot_confusion_matrix_from_csv(
+        raw_csv_file,
+        f"Confusion Matrix (Raw Channels)\n{percentage:.1f}% data",
+        raw_cm_img_path,
+    )
 
     # Process embeddings confusion matrix
     embeddings_csv_file = os.path.join(embeddings_folder, f"test_predictions_embedding_{percentage/100:.3f}_los.csv")
     embeddings_cm_img_path = os.path.join(embeddings_folder, "confusion_matrix_embeddings.png")
-    try:
-        embeddings_img = plot_confusion_matrix_from_csv(
-            embeddings_csv_file,
-            f"Confusion Matrix (LWM Embeddings)\n{percentage:.1f}% data",
-            embeddings_cm_img_path,
-        )
-    except Exception as exc:
-        embeddings_img = _make_error_image(
-            f"Failed to load Embedding CSV:\n{embeddings_csv_file}\n\n{exc}"
-        )
+    embeddings_img = plot_confusion_matrix_from_csv(
+        embeddings_csv_file,
+        f"Confusion Matrix (LWM Embeddings)\n{percentage:.1f}% data",
+        embeddings_cm_img_path,
+    )
 
     return raw_img, embeddings_img