Update app.py

app.py

@@ -6,6 +6,7 @@ import certifi
 import numpy as np
 import gradio as gr
 import torch.nn as nn
+import torch.optim as optim
 from torchvision import models
 import torch.nn.functional as F
 import matplotlib.pyplot as plt
@@ -32,7 +33,7 @@ with urllib.request.urlopen(url) as f:
 cifar10_classes = {int(k): v[1] for k, v in imagenet_classes.items()}
 
 # Load the pre-trained ResNet model
-model = models.resnet18(pretrained=True)
+model = models.resnet152(pretrained=True)
 
 # Modify the classifier for 6 classes with an additional hidden layer
 # model.fc = nn.Sequential(
@@ -182,83 +183,97 @@ def superimpose_images(base_image, overlay_image, alpha):
 
     return Image.fromarray(blended_array)
 
+def generate_adversarial(input_image, orig_pred, epsilon=12/255, steps=50):
+    """Generate adversarial example"""
+    input_image = transform_image(input_image)
+    delta = torch.zeros_like(input_image, requires_grad=True)
+    opt = optim.SGD([delta], lr=1e-1, momentum=0.9)
+
+    for _ in range(steps):
+        perturbed = torch.clamp(input_image + delta, 0.0, 1.0)
+        output = model(perturbed)
+        loss = -F.cross_entropy(output, torch.tensor([orig_pred], device=device))
+        print(loss.item())  
+        
+        opt.zero_grad()
+        loss.backward()
+        opt.step()
+        delta.data.clamp_(-epsilon, epsilon)
+
+    return torch.clamp(input_image + delta, 0.0, 1.0)
+
 # Prediction function
-def predict(image, brightness, contrast, hue, overlay_image, alpha):
-    """Apply filters, superimpose, classify image, and visualize results."""
+def predict(image, brightness, contrast, hue, overlay_image, alpha, adversarial_switch):
+    """Main prediction function"""
     if image is None:
         return None, None, None
 
-    # Apply feature filters
-    processed_image = apply_filters(image, brightness, contrast, hue)
-
-    # Superimpose overlay image
-    final_image = superimpose_images(processed_image, overlay_image, alpha)
+    # Apply preprocessing
+    processed = apply_filters(image, brightness, contrast, hue)
+    final_image = superimpose_images(processed, overlay_image, alpha)
+    
+    # Convert to tensor
+    to_tensor = transforms.Compose([transforms.Resize(32), transforms.ToTensor()])
+    # img_tensor_01 = to_tensor(final_image).unsqueeze(0).to(device)
+    
+    # Generate adversarial if enabled
+    if adversarial_switch:
+        with torch.no_grad():
+            orig_out = model(transform_image(final_image))
+            orig_pred = torch.argmax(orig_out).item()
+        adv_tensor_01 = generate_adversarial(final_image, orig_pred)
+        final_display = transforms.ToPILImage()(adv_tensor_01.squeeze().cpu().detach())
+        model_input = transform_image(final_display)
+    else:
+        final_display = final_image
+        model_input = transform_image(final_image)
+
+    # Get predictions
+    with torch.no_grad():
+        output = model(model_input)
+        probs = F.softmax(output, dim=1).cpu().numpy()[0]
+
+    # Generate Grad-CAM
+    heatmap, _ = gradcam.generate(model_input)
+    final_np = np.array(final_display)
+    heatmap = cv2.resize(heatmap, final_np.shape[:2][::-1])
+    heatmap = cv2.applyColorMap(np.uint8(255 * heatmap), cv2.COLORMAP_JET)
+    superimposed = cv2.addWeighted(cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB), 0.5, final_np, 0.5, 0)
+    
+    # Create plot
+    fig, ax = plt.subplots(figsize=(6, 4))
+    top5_idx = np.argsort(probs)[-5:][::-1]
+    ax.bar([class_labels[i] for i in top5_idx], probs[top5_idx], color='skyblue')
+    ax.set_ylabel("Probability"), ax.set_title("Class Probabilities")
+    plt.xticks(rotation=45, ha='right', fontsize=8)
+    plt.tight_layout()
 
-    # Convert PIL Image to Tensor
-    image_tensor = transform_image(final_image)
+    return final_display, Image.fromarray(superimposed), fig
 
-    with torch.no_grad():
-        output = model(image_tensor)
-        probabilities = F.softmax(output, dim=1).cpu().numpy()[0]
-        # pred_class = np.argmax(probabilities)
-        # top_5 = torch.topk(probabilities, 5)
-
-    heatmap, _ = gradcam.generate(image_tensor)
-
-    # Create GradCAM overlay
-    final_np = np.array(final_image)
-    heatmap = cv2.resize(heatmap, (final_np.shape[1], final_np.shape[0]))
-    heatmap = np.uint8(255 * heatmap)
-    heatmap_colored = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
-    heatmap_rgb = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)
-    superimposed = cv2.addWeighted(heatmap_rgb, 0.5, final_np, 0.5, 0)
-    gradcam_image = Image.fromarray(superimposed)
-
-
-    # Generate Bar Chart
-    with plt.xkcd():
-        fig, ax = plt.subplots(figsize=(6, 4))
-        top5_indices = np.argsort(probabilities)[-5:][::-1]  # Indices of top 5 probabilities
-        top5_probs = probabilities[top5_indices]
-        top5_labels = [class_labels[i] for i in top5_indices]
-        ax.bar(top5_labels, top5_probs, color='skyblue')
-        ax.set_ylabel("Probability")
-        ax.set_title("Class Probabilities")
-        ax.set_ylim([0, 1])
-        plt.tight_layout(pad=3)
-        ax.set_xticklabels(top5_labels, rotation=45, ha="right", fontsize=8)
-        for i, v in enumerate(top5_probs):
-            ax.text(i, v + 0.02, f"{v:.2f}", ha='center', fontsize=10)
-
-    return final_image, gradcam_image, fig
 
 # Gradio Interface
 with gr.Blocks() as interface:
-    gr.Markdown("<h2 style='text-align: center;'>Image Classifier with Superimposition & Adjustable Filters</h2>")
+    gr.Markdown("<h2 style='text-align: center;'>CIFAR-10 Classifier with Adversarial Attacks</h2>")
     
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(type="pil", label="Upload Base Image")
-            overlay_input = gr.Image(type="pil", label="Upload Overlay Image (Optional)")
+            image_input = gr.Image(type="pil", label="Input Image")
+            overlay_input = gr.Image(type="pil", label="Overlay Image (Optional)")
             brightness = gr.Slider(0.5, 2.0, value=1.0, label="Brightness")
             contrast = gr.Slider(0.5, 2.0, value=1.0, label="Contrast")
             hue = gr.Slider(-0.5, 0.5, value=0.0, label="Hue")
-            alpha = gr.Slider(0.0, 1.0, value=0.5, label="Overlay Weight (Alpha)")
+            alpha = gr.Slider(0.0, 1.0, value=0.5, label="Overlay Alpha")
+            adversarial_switch = gr.Checkbox(label="Add Adversarial Noise")
 
         with gr.Column():
-            processed_image = gr.Image(label="Final Processed Image")
+            processed_image = gr.Image(label="Processed Image")
             gradcam_output = gr.Image(label="GradCAM Overlay")
-            bar_chart = gr.Plot(label="Class Probabilities")
+            bar_chart = gr.Plot(label="Predictions")
 
-    inputs = [image_input, brightness, contrast, hue, overlay_input, alpha]
+    inputs = [image_input, brightness, contrast, hue, overlay_input, alpha, adversarial_switch]
     outputs = [processed_image, gradcam_output, bar_chart]
-
-    # Event listeners for real-time updates
-    image_input.change(predict, inputs=inputs, outputs=outputs)
-    overlay_input.change(predict, inputs=inputs, outputs=outputs)
-    brightness.change(predict, inputs=inputs, outputs=outputs)
-    contrast.change(predict, inputs=inputs, outputs=outputs)
-    hue.change(predict, inputs=inputs, outputs=outputs)
-    alpha.change(predict, inputs=inputs, outputs=outputs)
+    
+    for component in [image_input, overlay_input, brightness, contrast, hue, alpha, adversarial_switch]:
+        component.change(predict, inputs=inputs, outputs=outputs)
 
 interface.launch()