Update app.py

app.py

@@ -211,6 +211,9 @@ 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)
+
+    # Store the original size
+    original_size = final_image.size
     
     # Convert to tensor
     to_tensor = transforms.Compose([transforms.Resize(32), transforms.ToTensor()])
@@ -223,10 +226,10 @@ def predict(image, brightness, contrast, hue, overlay_image, alpha, adversarial_
             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)
+        # model_input = transform_image(final_display)
     else:
         final_display = final_image
-        model_input = transform_image(final_image)
+        # model_input = transform_image(final_image)
 
     # Get predictions
     with torch.no_grad():
@@ -236,7 +239,7 @@ def predict(image, brightness, contrast, hue, overlay_image, alpha, adversarial_
     # 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.resize(heatmap, (final_np.shape[1], final_np.shape[0]))
     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)