darpanaswal commited on
Commit
5e448d3
·
verified ·
1 Parent(s): ff0f581

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -11
app.py CHANGED
@@ -208,16 +208,11 @@ def predict(image, brightness, contrast, hue, overlay_image, alpha, adversarial_
208
  if image is None:
209
  return None, None, None
210
 
 
 
211
  # Apply preprocessing
212
  processed = apply_filters(image, brightness, contrast, hue)
213
  final_image = superimpose_images(processed, overlay_image, alpha)
214
-
215
- # Store the original size
216
- original_size = final_image.size
217
-
218
- # Convert to tensor
219
- to_tensor = transforms.Compose([transforms.Resize(32), transforms.ToTensor()])
220
- # img_tensor_01 = to_tensor(final_image).unsqueeze(0).to(device)
221
 
222
  # Generate adversarial if enabled
223
  if adversarial_switch:
@@ -226,10 +221,15 @@ def predict(image, brightness, contrast, hue, overlay_image, alpha, adversarial_
226
  orig_pred = torch.argmax(orig_out).item()
227
  adv_tensor_01 = generate_adversarial(final_image, orig_pred)
228
  final_display = transforms.ToPILImage()(adv_tensor_01.squeeze().cpu().detach())
229
- # model_input = transform_image(final_display)
 
230
  else:
231
- final_display = final_image
232
- # model_input = transform_image(final_image)
 
 
 
 
233
 
234
  # Get predictions
235
  with torch.no_grad():
@@ -239,7 +239,7 @@ def predict(image, brightness, contrast, hue, overlay_image, alpha, adversarial_
239
  # Generate Grad-CAM
240
  heatmap, _ = gradcam.generate(model_input)
241
  final_np = np.array(final_display)
242
- heatmap = cv2.resize(heatmap, (final_np.shape[1], final_np.shape[0]))
243
  heatmap = cv2.applyColorMap(np.uint8(255 * heatmap), cv2.COLORMAP_JET)
244
  superimposed = cv2.addWeighted(cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB), 0.5, final_np, 0.5, 0)
245
 
 
208
  if image is None:
209
  return None, None, None
210
 
211
+ orig_size = image.size # Get original size (width, height)
212
+
213
  # Apply preprocessing
214
  processed = apply_filters(image, brightness, contrast, hue)
215
  final_image = superimpose_images(processed, overlay_image, alpha)
 
 
 
 
 
 
 
216
 
217
  # Generate adversarial if enabled
218
  if adversarial_switch:
 
221
  orig_pred = torch.argmax(orig_out).item()
222
  adv_tensor_01 = generate_adversarial(final_image, orig_pred)
223
  final_display = transforms.ToPILImage()(adv_tensor_01.squeeze().cpu().detach())
224
+ final_display = final_display.resize(orig_size) # Resize back to original size
225
+ model_input = transform_image(final_display)
226
  else:
227
+ # Step 1: Resize to 224x224
228
+ resized_image = final_image.resize((224, 224))
229
+ # Step 2: Resize back to original size
230
+ final_display = resized_image.resize(orig_size)
231
+ # Model gets the 224x224 resized version
232
+ model_input = transform_image(resized_image)
233
 
234
  # Get predictions
235
  with torch.no_grad():
 
239
  # Generate Grad-CAM
240
  heatmap, _ = gradcam.generate(model_input)
241
  final_np = np.array(final_display)
242
+ heatmap = cv2.resize(heatmap, final_np.shape[:2][::-1])
243
  heatmap = cv2.applyColorMap(np.uint8(255 * heatmap), cv2.COLORMAP_JET)
244
  superimposed = cv2.addWeighted(cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB), 0.5, final_np, 0.5, 0)
245