Update app.py

app.py

@@ -123,7 +123,7 @@ class MultiLayerGradCAM:
         for handle in self.handles:
             handle.remove()
 
-gradcam = MultiLayerGradCAM(model, target_layers=['layer3', 'layer4'])
+gradcam = MultiLayerGradCAM(model, target_layers=['layer4'])
 
 # Image transformation function
 def transform_image(image):
@@ -131,8 +131,8 @@ def transform_image(image):
     mean, std = [0.4914, 0.4822, 0.4465], [0.247, 0.243, 0.261]
     img_size=224
     transform = transforms.Compose([ #IMAGENET
-    transforms.Resize(256),        # Resize shorter side to 256, keeping aspect ratio
-    transforms.CenterCrop(224),    # Crop the center 224x224 region
+    transforms.Resize((224,224)),        # Resize shorter side to 256, keeping aspect ratio
+    # transforms.CenterCrop(224),    # Crop the center 224x224 region
     transforms.ToTensor(),         # Convert to tensor (scales to [0,1])
     transforms.Normalize(          # Normalize using ImageNet mean & std
         mean=[0.485, 0.456, 0.406],
@@ -183,11 +183,11 @@ 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):
+def generate_adversarial(input_image, orig_pred, epsilon=20/255, steps=500):
     """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)
+    opt = optim.SGD([delta], lr=2e-1, momentum=0.9)
 
     for _ in range(steps):
         perturbed = torch.clamp(input_image + delta, 0.0, 1.0)
@@ -199,8 +199,13 @@ def generate_adversarial(input_image, orig_pred, epsilon=12/255, steps=50):
         loss.backward()
         opt.step()
         delta.data.clamp_(-epsilon, epsilon)
+    x = input_image + delta  # Compute the raw sum
+    x_min = x.amin(dim=(1, 2, 3), keepdim=True)  # Per-image min
+    x_max = x.amax(dim=(1, 2, 3), keepdim=True)  # Per-image max
 
-    return torch.clamp(input_image + delta, 0.0, 1.0)
+    output = (x - x_min) / (x_max - x_min + 1e-8)  # Avoid division by zero
+    # return output
+    return output
 
 # Prediction function
 def predict(image, brightness, contrast, hue, overlay_image, alpha, adversarial_switch):