Update app.py

app.py

@@ -208,16 +208,11 @@ def predict(image, brightness, contrast, hue, overlay_image, alpha, adversarial_
     if image is None:
         return None, None, None
 
+    orig_size = image.size  # Get original size (width, height)
+
     # 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()])
-    # img_tensor_01 = to_tensor(final_image).unsqueeze(0).to(device)
     
     # Generate adversarial if enabled
     if adversarial_switch:
@@ -226,10 +221,15 @@ 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)
+        final_display = final_display.resize(orig_size)  # Resize back to original size
+        model_input = transform_image(final_display)
     else:
-        final_display = final_image
-        # model_input = transform_image(final_image)
+        # 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():
@@ -239,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[1], final_np.shape[0]))
+    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)