Update app.py

app.py

@@ -216,7 +216,7 @@ def predict(image, brightness, contrast, hue, overlay_image, alpha, adversarial_
 
     # Apply preprocessing
     processed = apply_filters(image, brightness, contrast, hue)
-    final_image = superimpose_images(processed, overlay_image, alpha)  # Overlay the image properly
+    final_image = superimpose_images(processed, overlay_image, alpha)
     
     # Generate adversarial if enabled
     if adversarial_switch:
@@ -228,25 +228,26 @@ def predict(image, brightness, contrast, hue, overlay_image, alpha, adversarial_
         final_display = final_display.resize(orig_size)  # Resize back to original size
         model_input = transform_image(final_display)
     else:
-        # Step 1: Resize to 224x224
         resized_image = final_image.resize((224, 224))
-        # Step 2: Resize back to original size
         final_display = resized_image.resize(orig_size)
-        # Model gets the 224x224 resized version
         model_input = transform_image(resized_image)
 
     # Get predictions
     with torch.no_grad():
         output = model(model_input)
         probs = F.softmax(output, dim=1).cpu().numpy()[0]
+    
+    # Reset Grad-CAM activations before generating new heatmap
+    gradcam.activations.clear()
+    gradcam.gradients.clear()
 
-    # Generate Grad-CAM on the properly transformed image with overlay
+    # 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]
@@ -257,7 +258,6 @@ def predict(image, brightness, contrast, hue, overlay_image, alpha, adversarial_
 
     return final_display, Image.fromarray(superimposed), fig
 
-
 # Gradio Interface
 with gr.Blocks() as interface:
     gr.Markdown("<h2 style='text-align: center;'>ResNet Classifier with Adversarial Attacks</h2>")