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

Author: akashAD98

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)