Spaces:
Running
on
T4
Running
on
T4
| # Copyright (c) OpenMMLab. All rights reserved. | |
| """This script is in the experimental verification stage and cannot be | |
| guaranteed to be completely correct. Currently Grad-based CAM and Grad-free CAM | |
| are supported. | |
| The target detection task is different from the classification task. It not | |
| only includes the AM map of the category, but also includes information such as | |
| bbox and mask, so this script is named bboxam. | |
| """ | |
| import argparse | |
| import os.path | |
| import warnings | |
| from functools import partial | |
| import cv2 | |
| import mmcv | |
| from mmengine import Config, DictAction, MessageHub | |
| from mmengine.utils import ProgressBar | |
| try: | |
| from pytorch_grad_cam import AblationCAM, EigenCAM | |
| except ImportError: | |
| raise ImportError('Please run `pip install "grad-cam"` to install ' | |
| 'pytorch_grad_cam package.') | |
| from mmyolo.utils.boxam_utils import (BoxAMDetectorVisualizer, | |
| BoxAMDetectorWrapper, DetAblationLayer, | |
| DetBoxScoreTarget, GradCAM, | |
| GradCAMPlusPlus, reshape_transform) | |
| from mmyolo.utils.misc import get_file_list | |
| GRAD_FREE_METHOD_MAP = { | |
| 'ablationcam': AblationCAM, | |
| 'eigencam': EigenCAM, | |
| # 'scorecam': ScoreCAM, # consumes too much memory | |
| } | |
| GRAD_BASED_METHOD_MAP = {'gradcam': GradCAM, 'gradcam++': GradCAMPlusPlus} | |
| ALL_SUPPORT_METHODS = list(GRAD_FREE_METHOD_MAP.keys() | |
| | GRAD_BASED_METHOD_MAP.keys()) | |
| IGNORE_LOSS_PARAMS = { | |
| 'yolov5': ['loss_obj'], | |
| 'yolov6': ['loss_cls'], | |
| 'yolox': ['loss_obj'], | |
| 'rtmdet': ['loss_cls'], | |
| 'yolov7': ['loss_obj'], | |
| 'yolov8': ['loss_cls'], | |
| 'ppyoloe': ['loss_cls'], | |
| } | |
| # This parameter is required in some algorithms | |
| # for calculating Loss | |
| message_hub = MessageHub.get_current_instance() | |
| message_hub.runtime_info['epoch'] = 0 | |
| def parse_args(): | |
| parser = argparse.ArgumentParser(description='Visualize Box AM') | |
| parser.add_argument( | |
| 'img', help='Image path, include image file, dir and URL.') | |
| parser.add_argument('config', help='Config file') | |
| parser.add_argument('checkpoint', help='Checkpoint file') | |
| parser.add_argument( | |
| '--method', | |
| default='gradcam', | |
| choices=ALL_SUPPORT_METHODS, | |
| help='Type of method to use, supports ' | |
| f'{", ".join(ALL_SUPPORT_METHODS)}.') | |
| parser.add_argument( | |
| '--target-layers', | |
| default=['neck.out_layers[2]'], | |
| nargs='+', | |
| type=str, | |
| help='The target layers to get Box AM, if not set, the tool will ' | |
| 'specify the neck.out_layers[2]') | |
| parser.add_argument( | |
| '--out-dir', default='./output', help='Path to output file') | |
| parser.add_argument( | |
| '--show', action='store_true', help='Show the CAM results') | |
| parser.add_argument( | |
| '--device', default='cuda:0', help='Device used for inference') | |
| parser.add_argument( | |
| '--score-thr', type=float, default=0.3, help='Bbox score threshold') | |
| parser.add_argument( | |
| '--topk', | |
| type=int, | |
| default=-1, | |
| help='Select topk predict resutls to show. -1 are mean all.') | |
| parser.add_argument( | |
| '--max-shape', | |
| nargs='+', | |
| type=int, | |
| default=-1, | |
| help='max shapes. Its purpose is to save GPU memory. ' | |
| 'The activation map is scaled and then evaluated. ' | |
| 'If set to -1, it means no scaling.') | |
| parser.add_argument( | |
| '--preview-model', | |
| default=False, | |
| action='store_true', | |
| help='To preview all the model layers') | |
| parser.add_argument( | |
| '--norm-in-bbox', action='store_true', help='Norm in bbox of am image') | |
| parser.add_argument( | |
| '--cfg-options', | |
| nargs='+', | |
| action=DictAction, | |
| help='override some settings in the used config, the key-value pair ' | |
| 'in xxx=yyy format will be merged into config file. If the value to ' | |
| 'be overwritten is a list, it should be like key="[a,b]" or key=a,b ' | |
| 'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" ' | |
| 'Note that the quotation marks are necessary and that no white space ' | |
| 'is allowed.') | |
| # Only used by AblationCAM | |
| parser.add_argument( | |
| '--batch-size', | |
| type=int, | |
| default=1, | |
| help='batch of inference of AblationCAM') | |
| parser.add_argument( | |
| '--ratio-channels-to-ablate', | |
| type=int, | |
| default=0.5, | |
| help='Making it much faster of AblationCAM. ' | |
| 'The parameter controls how many channels should be ablated') | |
| args = parser.parse_args() | |
| return args | |
| def init_detector_and_visualizer(args, cfg): | |
| max_shape = args.max_shape | |
| if not isinstance(max_shape, list): | |
| max_shape = [args.max_shape] | |
| assert len(max_shape) == 1 or len(max_shape) == 2 | |
| model_wrapper = BoxAMDetectorWrapper( | |
| cfg, args.checkpoint, args.score_thr, device=args.device) | |
| if args.preview_model: | |
| print(model_wrapper.detector) | |
| print('\n Please remove `--preview-model` to get the BoxAM.') | |
| return None, None | |
| target_layers = [] | |
| for target_layer in args.target_layers: | |
| try: | |
| target_layers.append( | |
| eval(f'model_wrapper.detector.{target_layer}')) | |
| except Exception as e: | |
| print(model_wrapper.detector) | |
| raise RuntimeError('layer does not exist', e) | |
| ablationcam_extra_params = { | |
| 'batch_size': args.batch_size, | |
| 'ablation_layer': DetAblationLayer(), | |
| 'ratio_channels_to_ablate': args.ratio_channels_to_ablate | |
| } | |
| if args.method in GRAD_BASED_METHOD_MAP: | |
| method_class = GRAD_BASED_METHOD_MAP[args.method] | |
| is_need_grad = True | |
| else: | |
| method_class = GRAD_FREE_METHOD_MAP[args.method] | |
| is_need_grad = False | |
| boxam_detector_visualizer = BoxAMDetectorVisualizer( | |
| method_class, | |
| model_wrapper, | |
| target_layers, | |
| reshape_transform=partial( | |
| reshape_transform, max_shape=max_shape, is_need_grad=is_need_grad), | |
| is_need_grad=is_need_grad, | |
| extra_params=ablationcam_extra_params) | |
| return model_wrapper, boxam_detector_visualizer | |
| def main(): | |
| args = parse_args() | |
| # hard code | |
| ignore_loss_params = None | |
| for param_keys in IGNORE_LOSS_PARAMS: | |
| if param_keys in args.config: | |
| print(f'The algorithm currently used is {param_keys}') | |
| ignore_loss_params = IGNORE_LOSS_PARAMS[param_keys] | |
| break | |
| cfg = Config.fromfile(args.config) | |
| if args.cfg_options is not None: | |
| cfg.merge_from_dict(args.cfg_options) | |
| if not os.path.exists(args.out_dir) and not args.show: | |
| os.mkdir(args.out_dir) | |
| model_wrapper, boxam_detector_visualizer = init_detector_and_visualizer( | |
| args, cfg) | |
| # get file list | |
| image_list, source_type = get_file_list(args.img) | |
| progress_bar = ProgressBar(len(image_list)) | |
| for image_path in image_list: | |
| image = cv2.imread(image_path) | |
| model_wrapper.set_input_data(image) | |
| # forward detection results | |
| result = model_wrapper()[0] | |
| pred_instances = result.pred_instances | |
| # Get candidate predict info with score threshold | |
| pred_instances = pred_instances[pred_instances.scores > args.score_thr] | |
| if len(pred_instances) == 0: | |
| warnings.warn('empty detection results! skip this') | |
| continue | |
| if args.topk > 0: | |
| pred_instances = pred_instances[:args.topk] | |
| targets = [ | |
| DetBoxScoreTarget( | |
| pred_instances, | |
| device=args.device, | |
| ignore_loss_params=ignore_loss_params) | |
| ] | |
| if args.method in GRAD_BASED_METHOD_MAP: | |
| model_wrapper.need_loss(True) | |
| model_wrapper.set_input_data(image, pred_instances) | |
| boxam_detector_visualizer.switch_activations_and_grads( | |
| model_wrapper) | |
| # get box am image | |
| grayscale_boxam = boxam_detector_visualizer(image, targets=targets) | |
| # draw cam on image | |
| pred_instances = pred_instances.numpy() | |
| image_with_bounding_boxes = boxam_detector_visualizer.show_am( | |
| image, | |
| pred_instances, | |
| grayscale_boxam, | |
| with_norm_in_bboxes=args.norm_in_bbox) | |
| if source_type['is_dir']: | |
| filename = os.path.relpath(image_path, args.img).replace('/', '_') | |
| else: | |
| filename = os.path.basename(image_path) | |
| out_file = None if args.show else os.path.join(args.out_dir, filename) | |
| if out_file: | |
| mmcv.imwrite(image_with_bounding_boxes, out_file) | |
| else: | |
| cv2.namedWindow(filename, 0) | |
| cv2.imshow(filename, image_with_bounding_boxes) | |
| cv2.waitKey(0) | |
| # switch | |
| if args.method in GRAD_BASED_METHOD_MAP: | |
| model_wrapper.need_loss(False) | |
| boxam_detector_visualizer.switch_activations_and_grads( | |
| model_wrapper) | |
| progress_bar.update() | |
| if not args.show: | |
| print(f'All done!' | |
| f'\nResults have been saved at {os.path.abspath(args.out_dir)}') | |
| if __name__ == '__main__': | |
| main() | |