update app

app.py

@@ -12,6 +12,7 @@ import time
 from datetime import datetime
 import gradio as gr
 
+# LeNet for MNIST
 class LeNet(nn.Module):
     def __init__(self):
         super(LeNet, self).__init__()
@@ -41,6 +42,7 @@ class LeNet(nn.Module):
         else:
             return x5
 
+# Noise functions
 def salt_pepper_noise(images, prob=0.01, device='cuda'):
     batch_smap = torch.rand_like(images) < prob / 2
     pepper = torch.rand_like(images) < prob / 2
@@ -55,6 +57,7 @@ def pepper_statistical_noise(images, prob=0.01, device='cuda'):
     noisy[pepper] = 0.0
     return torch.clamp(noisy, 0, 1)
 
+# MVL calculation with hooks fallback
 def get_layer_outputs(model, input_tensor):
     outputs = []
     def hook(module, input, output):
@@ -128,6 +131,7 @@ def get_models_for_dataset(dataset_name):
     else:
         return []
 
+
 def get_dataset_and_transform(dataset_name):
     if dataset_name == 'MNIST':
         transform = transforms.Compose([
@@ -175,16 +179,20 @@ def layer_sustainability_analysis(dataset_name, model_name, selected_attacks, nu
     logs = []
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
+    # Prepare dataset & loader
     dataset, transform = get_dataset_and_transform(dataset_name)
     testloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=False)
     logs.append(f"{dataset_name} dataset loaded")
 
+    # Init model
     model = initialize_model(model_name, device)
     logs.append(f"Model {model_name} loaded on {device}")
 
+    # Model stats
     param_count, layer_count = get_model_stats(model)
     logs.append(f"Model stats: {param_count} parameters, {layer_count} layers")
 
+    # Setup attacks
     all_attacks = {
         'FGSM': FGSM(model, eps=0.03),
         'PGD': PGD(model, eps=0.03, alpha=0.01, steps=40, random_start=True),
@@ -198,41 +206,109 @@ def layer_sustainability_analysis(dataset_name, model_name, selected_attacks, nu
         return ["No valid attacks selected", None] + [None]*6 + ["", '\n'.join(logs)]
     logs.append(f"Selected attacks: {', '.join(attacks.keys())}")
 
-    results = {attack_name: {'cm': [], 'mvl': []} for attack_name in attacks}
+    # Prepare output dir for plots
+    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
+    output_dir = os.path.join(output_dir_base, f"{model_name}_{timestamp}")
+    os.makedirs(output_dir, exist_ok=True)
+    logs.append(f"Output directory: {output_dir}")
+
+    # Collect results
+    results = {atk: {'cm': [], 'mvl': []} for atk in attacks}
 
+    # Process batches
     for i, (images, labels) in enumerate(testloader):
         if i >= num_batches:
             break
         images, labels = images.to(device), labels.to(device)
         logs.append(f"Processing batch {i+1}/{num_batches}...")
 
-        for attack_name, attack in attacks.items():
-            adv_images = attack(images, labels)
-            mvl_list = compute_mvl(model, images, adv_images, device)
-            results[attack_name]['mvl'].append(mvl_list)
-            cm = np.mean(mvl_list)
-            results[attack_name]['cm'].append(cm)
-
-    # Placeholders for plots (add your plot generation here)
-    cm_plot_path = None
-    mvl_plot_paths = [None]*5
-    integrated_mvl_plot_path = None
+        for atk_name, atk in attacks.items():
+            adv_images = atk(images, labels)
+            mvl_vals = compute_mvl(model, images, adv_images, device)
+            results[atk_name]['mvl'].append(mvl_vals)
+            results[atk_name]['cm'].append(np.mean(mvl_vals))
+
+    # Compute mean/std CM per attack
+    cm_means = {atk: np.mean(results[atk]['cm']) for atk in attacks}
+    cm_stds = {atk: np.std(results[atk]['cm']) for atk in attacks}
+
+    # Plot CM bar
+    plt.figure(figsize=(8,6))
+    names = list(attacks.keys())
+    means = [cm_means[n] for n in names]
+    stds = [cm_stds[n] for n in names]
+    x = np.arange(len(names))
+    plt.bar(x, means, yerr=stds, capsize=5)
+    plt.xticks(x, names, rotation=45)
+    plt.ylabel("CM (Relative Error)")
+    plt.title(f"CM for {model_name} ({dataset_name})")
+    plt.tight_layout()
+    cm_plot_path = os.path.join(output_dir, "cm_plot.png")
+    plt.savefig(cm_plot_path)
+    plt.close()
+    logs.append(f"Saved CM plot to {cm_plot_path}")
+
+    # Plot MVL per attack
+    mvl_plot_paths = []
+    colors = ['skyblue', 'lightgreen', 'coral', 'lightgray', 'purple']
+    for idx, atk in enumerate(names):
+        mvl_arr = np.array(results[atk]['mvl'])
+        mean_vals = np.mean(mvl_arr, axis=0)
+        std_vals = np.std(mvl_arr, axis=0)
+        layers = [f"Layer {i+1}" for i in range(len(mean_vals))]
+        plt.figure(figsize=(8,6))
+        plt.plot(layers, mean_vals, marker='o', color=colors[idx % len(colors)], label=atk)
+        plt.fill_between(layers, mean_vals - std_vals, mean_vals + std_vals, color=colors[idx % len(colors)], alpha=0.3)
+        plt.title(f"MVL per Layer - {atk}")
+        plt.ylabel("MVL (Mean ± Std)")
+        plt.xticks(rotation=45)
+        plt.grid(True)
+        plt.tight_layout()
+        path = os.path.join(output_dir, f"mvl_{atk.lower().replace(' ', '_')}.png")
+        plt.savefig(path)
+        plt.close()
+        mvl_plot_paths.append(path)
+        logs.append(f"Saved MVL plot for {atk} to {path}")
+
+    # Integrated MVL plot
+    plt.figure(figsize=(10,6))
+    for idx, atk in enumerate(names):
+        mvl_arr = np.array(results[atk]['mvl'])
+        mean_vals = np.mean(mvl_arr, axis=0)
+        std_vals = np.std(mvl_arr, axis=0)
+        layers = [f"Layer {i+1}" for i in range(len(mean_vals))]
+        plt.plot(layers, mean_vals, marker='o', color=colors[idx % len(colors)], label=atk)
+        plt.fill_between(layers, mean_vals - std_vals, mean_vals + std_vals, color=colors[idx % len(colors)], alpha=0.3)
+    plt.title(f"Integrated MVL - {model_name}")
+    plt.ylabel("MVL (Mean ± Std)")
+    plt.xticks(rotation=45)
+    plt.legend()
+    plt.grid(True)
+    plt.tight_layout()
+    integrated_mvl_plot_path = os.path.join(output_dir, "integrated_mvl.png")
+    plt.savefig(integrated_mvl_plot_path)
+    plt.close()
+    logs.append(f"Saved integrated MVL plot to {integrated_mvl_plot_path}")
 
     processing_time = time.time() - start_time
 
     stats = {
         'Dataset': dataset_name,
         'Model': model_name,
-        'Parameter Count': param_count,
-        'Layer Count': layer_count,
-        'Processing Time (s)': round(processing_time, 2),
-        'Number of Batches': num_batches,
-        'Attacks Used': ', '.join(attacks.keys())
+        'Parameters': param_count,
+        'Layers': layer_count,
+        'Batches': num_batches,
+        'Attacks': ', '.join(names),
+        'Time (s)': round(processing_time, 2)
     }
-    stats_text = "## Model Statistics\n\n| Metric | Value |\n|--------|-------|\n"
+    stats_text = "## Model Statistics\n\n| Metric | Value |\n|---|---|\n"
     for k,v in stats.items():
         stats_text += f"| {k} | {v} |\n"
 
+    # Pad MVL plot paths to length 5 for UI consistency
+    while len(mvl_plot_paths) < 5:
+        mvl_plot_paths.append(None)
+
     return [None, cm_plot_path] + mvl_plot_paths[:5] + [integrated_mvl_plot_path, stats_text, '\n'.join(logs)]
 
 paper_info_html = """
@@ -260,7 +336,7 @@ paper_info_html = """
 def update_models(dataset_name):
     models = get_models_for_dataset(dataset_name)
     default_value = models[0] if models else None
-    return models, default_value  # Return choices and default value as a tuple
+    return models, default_value  # Return choices and default value for older gradio versions
 
 def create_interface():
     datasets = ['MNIST', 'CIFAR-10']
@@ -292,11 +368,10 @@ def create_interface():
             with gr.Tab("Logs"):
                 log_output = gr.Textbox(label="Processing Logs")
 
-        # Return choices and value separately for older gradio versions
         dataset_input.change(
             fn=update_models,
             inputs=dataset_input,
-            outputs=[model_input, model_input]
+            outputs=[model_input, model_input]  # updates choices and default value
         )
 
         run_button.click(