|
|
|
|
|
import argparse
|
|
|
import os
|
|
|
import os.path as osp
|
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
|
|
import mmcv
|
|
|
import numpy as np
|
|
|
import torch.nn as nn
|
|
|
from mmengine import Config, DictAction
|
|
|
from mmengine.dataset import Compose, pseudo_collate
|
|
|
|
|
|
from mmaction.apis import init_recognizer
|
|
|
from mmaction.utils import GradCAM
|
|
|
|
|
|
|
|
|
def parse_args():
|
|
|
parser = argparse.ArgumentParser(
|
|
|
description='MMAction2 GradCAM Visualization')
|
|
|
parser.add_argument('config', help='test config file path')
|
|
|
parser.add_argument('checkpoint', help='checkpoint file/url')
|
|
|
parser.add_argument('video', help='video file/url or rawframes directory')
|
|
|
parser.add_argument(
|
|
|
'--use-frames',
|
|
|
default=False,
|
|
|
action='store_true',
|
|
|
help='whether to use rawframes as input')
|
|
|
parser.add_argument(
|
|
|
'--device', type=str, default='cuda:0', help='CPU/CUDA device option')
|
|
|
parser.add_argument(
|
|
|
'--target-layer-name',
|
|
|
type=str,
|
|
|
default='backbone/layer4/1/relu',
|
|
|
help='GradCAM target layer name')
|
|
|
parser.add_argument('--out-filename', default=None, help='output filename')
|
|
|
parser.add_argument('--fps', default=5, type=int)
|
|
|
parser.add_argument(
|
|
|
'--cfg-options',
|
|
|
nargs='+',
|
|
|
action=DictAction,
|
|
|
default={},
|
|
|
help='override some settings in the used config, the key-value pair '
|
|
|
'in xxx=yyy format will be merged into config file. For example, '
|
|
|
"'--cfg-options model.backbone.depth=18 model.backbone.with_cp=True'")
|
|
|
parser.add_argument(
|
|
|
'--target-resolution',
|
|
|
nargs=2,
|
|
|
default=None,
|
|
|
type=int,
|
|
|
help='Target resolution (w, h) for resizing the frames when using a '
|
|
|
'video as input. If either dimension is set to -1, the frames are '
|
|
|
'resized by keeping the existing aspect ratio')
|
|
|
parser.add_argument(
|
|
|
'--resize-algorithm',
|
|
|
default='bilinear',
|
|
|
help='resize algorithm applied to generate video & gif')
|
|
|
|
|
|
args = parser.parse_args()
|
|
|
return args
|
|
|
|
|
|
|
|
|
def build_inputs(model: nn.Module,
|
|
|
video_path: str,
|
|
|
use_frames: bool = False) -> Dict:
|
|
|
"""build inputs for GradCAM.
|
|
|
|
|
|
Note that, building inputs for GradCAM is exactly the same as building
|
|
|
inputs for Recognizer test stage. Codes from `inference_recognizer`.
|
|
|
|
|
|
Args:
|
|
|
model (nn.Module): Recognizer model.
|
|
|
video_path (str): video file/url or rawframes directory.
|
|
|
use_frames (bool): whether to use rawframes as input.
|
|
|
Defaults to False.
|
|
|
|
|
|
Returns:
|
|
|
dict: Both GradCAM inputs and Recognizer test stage inputs,
|
|
|
including two keys, ``inputs`` and ``data_samples``.
|
|
|
"""
|
|
|
if not (osp.exists(video_path) or video_path.startswith('http')):
|
|
|
raise RuntimeError(f"'{video_path}' is missing")
|
|
|
|
|
|
if osp.isfile(video_path) and use_frames:
|
|
|
raise RuntimeError(
|
|
|
f"'{video_path}' is a video file, not a rawframe directory")
|
|
|
if osp.isdir(video_path) and not use_frames:
|
|
|
raise RuntimeError(
|
|
|
f"'{video_path}' is a rawframe directory, not a video file")
|
|
|
|
|
|
cfg = model.cfg
|
|
|
|
|
|
|
|
|
test_pipeline = cfg.test_pipeline
|
|
|
test_pipeline = Compose(test_pipeline)
|
|
|
|
|
|
if use_frames:
|
|
|
filename_tmpl = cfg.test_dataloader.dataset.get(
|
|
|
'filename_tmpl', 'img_{:05}.jpg')
|
|
|
start_index = cfg.test_dataloader.dataset.get('start_index', 1)
|
|
|
data = dict(
|
|
|
frame_dir=video_path,
|
|
|
total_frames=len(os.listdir(video_path)),
|
|
|
label=-1,
|
|
|
start_index=start_index,
|
|
|
filename_tmpl=filename_tmpl,
|
|
|
modality='RGB')
|
|
|
else:
|
|
|
start_index = cfg.test_dataloader.dataset.get('start_index', 0)
|
|
|
data = dict(
|
|
|
filename=video_path,
|
|
|
label=-1,
|
|
|
start_index=start_index,
|
|
|
modality='RGB')
|
|
|
data = test_pipeline(data)
|
|
|
data = pseudo_collate([data])
|
|
|
|
|
|
return data
|
|
|
|
|
|
|
|
|
def _resize_frames(frame_list: List[np.ndarray],
|
|
|
scale: Optional[Tuple[int]] = None,
|
|
|
keep_ratio: bool = True,
|
|
|
interpolation: str = 'bilinear') -> List[np.ndarray]:
|
|
|
"""Resize frames according to given scale.
|
|
|
|
|
|
Codes are modified from `mmaction/datasets/transforms/processing.py`,
|
|
|
`Resize` class.
|
|
|
|
|
|
Args:
|
|
|
frame_list (list[np.ndarray]): Frames to be resized.
|
|
|
scale (tuple[int]): If keep_ratio is True, it serves as scaling
|
|
|
factor or maximum size: the image will be rescaled as large
|
|
|
as possible within the scale. Otherwise, it serves as (w, h)
|
|
|
of output size.
|
|
|
keep_ratio (bool): If set to True, Images will be resized without
|
|
|
changing the aspect ratio. Otherwise, it will resize images to a
|
|
|
given size. Defaults to True.
|
|
|
interpolation (str): Algorithm used for interpolation:
|
|
|
'nearest' | 'bilinear'. Defaults to ``'bilinear'``.
|
|
|
|
|
|
Returns:
|
|
|
list[np.ndarray]: Resized frames.
|
|
|
"""
|
|
|
if scale is None or (scale[0] == -1 and scale[1] == -1):
|
|
|
return frame_list
|
|
|
scale = tuple(scale)
|
|
|
max_long_edge = max(scale)
|
|
|
max_short_edge = min(scale)
|
|
|
if max_short_edge == -1:
|
|
|
scale = (np.inf, max_long_edge)
|
|
|
|
|
|
img_h, img_w, _ = frame_list[0].shape
|
|
|
|
|
|
if keep_ratio:
|
|
|
new_w, new_h = mmcv.rescale_size((img_w, img_h), scale)
|
|
|
else:
|
|
|
new_w, new_h = scale
|
|
|
|
|
|
frame_list = [
|
|
|
mmcv.imresize(img, (new_w, new_h), interpolation=interpolation)
|
|
|
for img in frame_list
|
|
|
]
|
|
|
|
|
|
return frame_list
|
|
|
|
|
|
|
|
|
def main():
|
|
|
args = parse_args()
|
|
|
|
|
|
cfg = Config.fromfile(args.config)
|
|
|
cfg.merge_from_dict(args.cfg_options)
|
|
|
|
|
|
|
|
|
model = init_recognizer(cfg, args.checkpoint, device=args.device)
|
|
|
|
|
|
inputs = build_inputs(model, args.video, use_frames=args.use_frames)
|
|
|
gradcam = GradCAM(model, args.target_layer_name)
|
|
|
results = gradcam(inputs)
|
|
|
|
|
|
if args.out_filename is not None:
|
|
|
try:
|
|
|
from moviepy.editor import ImageSequenceClip
|
|
|
except ImportError:
|
|
|
raise ImportError('Please install moviepy to enable output file.')
|
|
|
|
|
|
|
|
|
frames_batches = (results[0] * 255.).numpy().astype(np.uint8)
|
|
|
frames = frames_batches.reshape(-1, *frames_batches.shape[-3:])
|
|
|
|
|
|
frame_list = list(frames)
|
|
|
frame_list = _resize_frames(
|
|
|
frame_list,
|
|
|
args.target_resolution,
|
|
|
interpolation=args.resize_algorithm)
|
|
|
|
|
|
video_clips = ImageSequenceClip(frame_list, fps=args.fps)
|
|
|
out_type = osp.splitext(args.out_filename)[1][1:]
|
|
|
if out_type == 'gif':
|
|
|
video_clips.write_gif(args.out_filename)
|
|
|
else:
|
|
|
video_clips.write_videofile(args.out_filename, remove_temp=True)
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
main()
|
|
|
|
