added support for saving output in output folder, (#439)

* Update cam.py

added support for output_dir argument

* Update README.md

updated redme & added support for storing output in customoutputdir.

python cam.py --image-path input.jpg --method gradcam --output-dir custom_output

* Create cam_custom.py

add your own custom_model

* Delete cam_custom.py

README.md

@@ -238,10 +238,11 @@ two smoothing methods are supported:
 
 # Running the example script:
 
-Usage: `python cam.py --image-path <path_to_image> --method <method>`
+Usage: `python cam.py --image-path <path_to_image> --method <method> --output-dir <output_dir_path> `
+
 
 To use with CUDA:
-`python cam.py --image-path <path_to_image> --use-cuda`
+`python cam.py --image-path <path_to_image> --use-cuda --output-dir <output_dir_path> `
 
 ----------
 

cam.py

@@ -1,25 +1,18 @@
 import argparse
+import os
 import cv2
 import numpy as np
 import torch
 from torchvision import models
-from pytorch_grad_cam import GradCAM, \
- HiResCAM, \
- ScoreCAM, \
- GradCAMPlusPlus, \
- AblationCAM, \
- XGradCAM, \
- EigenCAM, \
- EigenGradCAM, \
- LayerCAM, \
- FullGrad, \
- GradCAMElementWise
-
-
+from pytorch_grad_cam import (
+ GradCAM, HiResCAM, ScoreCAM, GradCAMPlusPlus,
+ AblationCAM, XGradCAM, EigenCAM, EigenGradCAM,
+ LayerCAM, FullGrad, GradCAMElementWise
+)
 from pytorch_grad_cam import GuidedBackpropReLUModel
-from pytorch_grad_cam.utils.image import show_cam_on_image, \
- deprocess_image, \
- preprocess_image
+from pytorch_grad_cam.utils.image import (
+ show_cam_on_image, deprocess_image, preprocess_image
+)
 from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
 
 
@@ -32,21 +25,24 @@ def get_args():
  type=str,
  default='./examples/both.png',
  help='Input image path')
- parser.add_argument('--aug_smooth', action='store_true',
+ parser.add_argument('--aug-smooth', action='store_true',
  help='Apply test time augmentation to smooth the CAM')
  parser.add_argument(
- '--eigen_smooth',
+ '--eigen-smooth',
  action='store_true',
- help='Reduce noise by taking the first principle componenet'
+ help='Reduce noise by taking the first principle component'
  'of cam_weights*activations')
  parser.add_argument('--method', type=str, default='gradcam',
- choices=['gradcam', 'hirescam', 'gradcam++',
- 'scorecam', 'xgradcam',
- 'ablationcam', 'eigencam',
- 'eigengradcam', 'layercam', 'fullgrad'],
- help='Can be gradcam/gradcam++/scorecam/xgradcam'
- '/ablationcam/eigencam/eigengradcam/layercam')
-
+ choices=[
+ 'gradcam', 'hirescam', 'gradcam++',
+ 'scorecam', 'xgradcam', 'ablationcam',
+ 'eigencam', 'eigengradcam', 'layercam',
+ 'fullgrad', 'gradcamelementwise'
+ ],
+ help='CAM method')
+
+ parser.add_argument('--output-dir', type=str, default='output',
+ help='Output directory to save the images')
  args = parser.parse_args()
  args.use_cuda = args.use_cuda and torch.cuda.is_available()
  if args.use_cuda:
@@ -59,25 +55,26 @@ def get_args():
 
 if __name__ == '__main__':
  """ python cam.py -image-path <path_to_image>
- Example usage of loading an image, and computing:
+ Example usage of loading an image and computing:
  1. CAM
  2. Guided Back Propagation
  3. Combining both
  """
 
  args = get_args()
- methods = \
- {"gradcam": GradCAM,
- "hirescam": HiResCAM,
- "scorecam": ScoreCAM,
- "gradcam++": GradCAMPlusPlus,
- "ablationcam": AblationCAM,
- "xgradcam": XGradCAM,
- "eigencam": EigenCAM,
- "eigengradcam": EigenGradCAM,
- "layercam": LayerCAM,
- "fullgrad": FullGrad,
- "gradcamelementwise": GradCAMElementWise}
+ methods = {
+ "gradcam": GradCAM,
+ "hirescam": HiResCAM,
+ "scorecam": ScoreCAM,
+ "gradcam++": GradCAMPlusPlus,
+ "ablationcam": AblationCAM,
+ "xgradcam": XGradCAM,
+ "eigencam": EigenCAM,
+ "eigengradcam": EigenGradCAM,
+ "layercam": LayerCAM,
+ "fullgrad": FullGrad,
+ "gradcamelementwise": GradCAMElementWise
+ }
 
  model = models.resnet50(pretrained=True)
 
@@ -93,6 +90,7 @@ def get_args():
  # You can also try selecting all layers of a certain type, with e.g:
  # from pytorch_grad_cam.utils.find_layers import find_layer_types_recursive
  # find_layer_types_recursive(model, [torch.nn.ReLU])
+
  target_layers = [model.layer4]
 
  rgb_img = cv2.imread(args.image_path, 1)[:, :, ::-1]
@@ -115,6 +113,7 @@ def get_args():
  target_layers=target_layers,
  use_cuda=args.use_cuda) as cam:
 
+
  # AblationCAM and ScoreCAM have batched implementations.
  # You can override the internal batch size for faster computation.
  cam.batch_size = 32
@@ -123,12 +122,9 @@ def get_args():
  aug_smooth=args.aug_smooth,
  eigen_smooth=args.eigen_smooth)
 
- # Here grayscale_cam has only one image in the batch
  grayscale_cam = grayscale_cam[0, :]
 
  cam_image = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)
-
- # cam_image is RGB encoded whereas "cv2.imwrite" requires BGR encoding.
  cam_image = cv2.cvtColor(cam_image, cv2.COLOR_RGB2BGR)
 
  gb_model = GuidedBackpropReLUModel(model=model, use_cuda=args.use_cuda)
@@ -138,6 +134,12 @@ def get_args():
  cam_gb = deprocess_image(cam_mask * gb)
  gb = deprocess_image(gb)
 
- cv2.imwrite(f'{args.method}_cam.jpg', cam_image)
- cv2.imwrite(f'{args.method}_gb.jpg', gb)
- cv2.imwrite(f'{args.method}_cam_gb.jpg', cam_gb)
+ os.makedirs(args.output_dir, exist_ok=True)
+
+ cam_output_path = os.path.join(args.output_dir, f'{args.method}_cam.jpg')
+ gb_output_path = os.path.join(args.output_dir, f'{args.method}_gb.jpg')
+ cam_gb_output_path = os.path.join(args.output_dir, f'{args.method}_cam_gb.jpg')
+
+ cv2.imwrite(cam_output_path, cam_image)
+ cv2.imwrite(gb_output_path, gb)
+ cv2.imwrite(cam_gb_output_path, cam_gb)