|
|
|
|
|
import os.path as osp
|
|
|
from typing import Dict, List, Optional, Sequence, Tuple, Union
|
|
|
|
|
|
import mmcv
|
|
|
import numpy as np
|
|
|
from mmengine.dist import master_only
|
|
|
from mmengine.fileio.io import isdir, isfile, join_path, list_dir_or_file
|
|
|
from mmengine.visualization import Visualizer
|
|
|
|
|
|
from mmaction.registry import VISBACKENDS, VISUALIZERS
|
|
|
from mmaction.structures import ActionDataSample
|
|
|
|
|
|
|
|
|
def _get_adaptive_scale(img_shape: Tuple[int, int],
|
|
|
min_scale: float = 0.3,
|
|
|
max_scale: float = 3.0) -> float:
|
|
|
"""Get adaptive scale according to frame shape.
|
|
|
|
|
|
The target scale depends on the the short edge length of the frame. If the
|
|
|
short edge length equals 224, the output is 1.0. And output linear scales
|
|
|
according the short edge length.
|
|
|
|
|
|
You can also specify the minimum scale and the maximum scale to limit the
|
|
|
linear scale.
|
|
|
|
|
|
Args:
|
|
|
img_shape (Tuple[int, int]): The shape of the canvas frame.
|
|
|
min_size (int): The minimum scale. Defaults to 0.3.
|
|
|
max_size (int): The maximum scale. Defaults to 3.0.
|
|
|
|
|
|
Returns:
|
|
|
int: The adaptive scale.
|
|
|
"""
|
|
|
short_edge_length = min(img_shape)
|
|
|
scale = short_edge_length / 224.
|
|
|
return min(max(scale, min_scale), max_scale)
|
|
|
|
|
|
|
|
|
@VISUALIZERS.register_module()
|
|
|
class ActionVisualizer(Visualizer):
|
|
|
"""Universal Visualizer for classification task.
|
|
|
|
|
|
Args:
|
|
|
name (str): Name of the instance. Defaults to 'visualizer'.
|
|
|
vis_backends (list, optional): Visual backend config list.
|
|
|
Defaults to None.
|
|
|
save_dir (str, optional): Save file dir for all storage backends.
|
|
|
If it is None, the backend storage will not save any data.
|
|
|
fig_save_cfg (dict): Keyword parameters of figure for saving.
|
|
|
Defaults to empty dict.
|
|
|
fig_show_cfg (dict): Keyword parameters of figure for showing.
|
|
|
Defaults to empty dict.
|
|
|
|
|
|
Examples:
|
|
|
>>> import torch
|
|
|
>>> import decord
|
|
|
>>> from pathlib import Path
|
|
|
>>> from mmaction.structures import ActionDataSample, ActionVisualizer
|
|
|
>>> from mmengine.structures import LabelData
|
|
|
>>> # Example frame
|
|
|
>>> video = decord.VideoReader('./demo/demo.mp4')
|
|
|
>>> video = video.get_batch(range(32)).asnumpy()
|
|
|
>>> # Example annotation
|
|
|
>>> data_sample = ActionDataSample()
|
|
|
>>> data_sample.gt_label = LabelData(item=torch.tensor([2]))
|
|
|
>>> # Setup the visualizer
|
|
|
>>> vis = ActionVisualizer(
|
|
|
... save_dir="./outputs",
|
|
|
... vis_backends=[dict(type='LocalVisBackend')])
|
|
|
>>> # Set classes names
|
|
|
>>> vis.dataset_meta = {'classes': ['running', 'standing', 'sitting']}
|
|
|
>>> # Save the visualization result by the specified storage backends.
|
|
|
>>> vis.add_datasample('demo', video)
|
|
|
>>> assert Path('outputs/vis_data/demo/frames_0/1.png').exists()
|
|
|
>>> assert Path('outputs/vis_data/demo/frames_0/2.png').exists()
|
|
|
>>> # Save another visualization result with the same name.
|
|
|
>>> vis.add_datasample('demo', video, step=1)
|
|
|
>>> assert Path('outputs/vis_data/demo/frames_1/2.png').exists()
|
|
|
"""
|
|
|
|
|
|
def __init__(
|
|
|
self,
|
|
|
name='visualizer',
|
|
|
vis_backends: Optional[List[Dict]] = None,
|
|
|
save_dir: Optional[str] = None,
|
|
|
fig_save_cfg=dict(frameon=False),
|
|
|
fig_show_cfg=dict(frameon=False)
|
|
|
) -> None:
|
|
|
super().__init__(
|
|
|
name=name,
|
|
|
image=None,
|
|
|
vis_backends=vis_backends,
|
|
|
save_dir=save_dir,
|
|
|
fig_save_cfg=fig_save_cfg,
|
|
|
fig_show_cfg=fig_show_cfg)
|
|
|
|
|
|
def _load_video(self,
|
|
|
video: Union[np.ndarray, Sequence[np.ndarray], str],
|
|
|
target_resolution: Optional[Tuple[int]] = None):
|
|
|
"""Load video from multiple source and convert to target resolution.
|
|
|
|
|
|
Args:
|
|
|
video (np.ndarray, str): The video to draw.
|
|
|
target_resolution (Tuple[int], optional): Set to
|
|
|
(desired_width desired_height) to have resized frames. If
|
|
|
either dimension is None, the frames are resized by keeping
|
|
|
the existing aspect ratio. Defaults to None.
|
|
|
"""
|
|
|
if isinstance(video, np.ndarray) or isinstance(video, list):
|
|
|
frames = video
|
|
|
elif isinstance(video, str):
|
|
|
|
|
|
if isfile(video):
|
|
|
try:
|
|
|
import decord
|
|
|
except ImportError:
|
|
|
raise ImportError(
|
|
|
'Please install decord to load video file.')
|
|
|
video = decord.VideoReader(video)
|
|
|
frames = [x.asnumpy()[..., ::-1] for x in video]
|
|
|
|
|
|
elif isdir(video):
|
|
|
frame_list = sorted(list_dir_or_file(video, list_dir=False))
|
|
|
frames = [mmcv.imread(join_path(video, x)) for x in frame_list]
|
|
|
else:
|
|
|
raise TypeError(f'type of video {type(video)} not supported')
|
|
|
|
|
|
if target_resolution is not None:
|
|
|
w, h = target_resolution
|
|
|
frame_h, frame_w, _ = frames[0].shape
|
|
|
if w == -1:
|
|
|
w = int(h / frame_h * frame_w)
|
|
|
if h == -1:
|
|
|
h = int(w / frame_w * frame_h)
|
|
|
frames = [mmcv.imresize(f, (w, h)) for f in frames]
|
|
|
|
|
|
return frames
|
|
|
|
|
|
@master_only
|
|
|
def add_datasample(self,
|
|
|
name: str,
|
|
|
video: Union[np.ndarray, Sequence[np.ndarray], str],
|
|
|
data_sample: Optional[ActionDataSample] = None,
|
|
|
draw_gt: bool = True,
|
|
|
draw_pred: bool = True,
|
|
|
draw_score: bool = True,
|
|
|
rescale_factor: Optional[float] = None,
|
|
|
show_frames: bool = False,
|
|
|
text_cfg: dict = dict(),
|
|
|
wait_time: float = 0.1,
|
|
|
out_path: Optional[str] = None,
|
|
|
out_type: str = 'img',
|
|
|
target_resolution: Optional[Tuple[int]] = None,
|
|
|
step: int = 0,
|
|
|
fps: int = 4) -> None:
|
|
|
"""Draw datasample and save to all backends.
|
|
|
|
|
|
- If ``out_path`` is specified, all storage backends are ignored
|
|
|
and save the videos to the ``out_path``.
|
|
|
- If ``show_frames`` is True, plot the frames in a window sequentially,
|
|
|
please confirm you are able to access the graphical interface.
|
|
|
|
|
|
Args:
|
|
|
name (str): The frame identifier.
|
|
|
video (np.ndarray, str): The video to draw. supports decoded
|
|
|
np.ndarray, video file path, rawframes folder path.
|
|
|
data_sample (:obj:`ActionDataSample`, optional): The annotation of
|
|
|
the frame. Defaults to None.
|
|
|
draw_gt (bool): Whether to draw ground truth labels.
|
|
|
Defaults to True.
|
|
|
draw_pred (bool): Whether to draw prediction labels.
|
|
|
Defaults to True.
|
|
|
draw_score (bool): Whether to draw the prediction scores
|
|
|
of prediction categories. Defaults to True.
|
|
|
rescale_factor (float, optional): Rescale the frame by the rescale
|
|
|
factor before visualization. Defaults to None.
|
|
|
show_frames (bool): Whether to display the frames of the video.
|
|
|
Defaults to False.
|
|
|
text_cfg (dict): Extra text setting, which accepts
|
|
|
arguments of :attr:`mmengine.Visualizer.draw_texts`.
|
|
|
Defaults to an empty dict.
|
|
|
wait_time (float): Delay in seconds. 0 is the special
|
|
|
value that means "forever". Defaults to 0.1.
|
|
|
out_path (str, optional): Extra folder to save the visualization
|
|
|
result. If specified, the visualizer will only save the result
|
|
|
frame to the out_path and ignore its storage backends.
|
|
|
Defaults to None.
|
|
|
out_type (str): Output format type, choose from 'img', 'gif',
|
|
|
'video'. Defaults to ``'img'``.
|
|
|
target_resolution (Tuple[int], optional): Set to
|
|
|
(desired_width desired_height) to have resized frames. If
|
|
|
either dimension is None, the frames are resized by keeping
|
|
|
the existing aspect ratio. Defaults to None.
|
|
|
step (int): Global step value to record. Defaults to 0.
|
|
|
fps (int): Frames per second for saving video. Defaults to 4.
|
|
|
"""
|
|
|
classes = None
|
|
|
video = self._load_video(video, target_resolution)
|
|
|
tol_video = len(video)
|
|
|
|
|
|
if self.dataset_meta is not None:
|
|
|
classes = self.dataset_meta.get('classes', None)
|
|
|
|
|
|
if data_sample is None:
|
|
|
data_sample = ActionDataSample()
|
|
|
|
|
|
resulted_video = []
|
|
|
for frame_idx, frame in enumerate(video):
|
|
|
frame_name = 'frame %d of %s' % (frame_idx + 1, name)
|
|
|
if rescale_factor is not None:
|
|
|
frame = mmcv.imrescale(frame, rescale_factor)
|
|
|
|
|
|
texts = ['Frame %d of total %d frames' % (frame_idx, tol_video)]
|
|
|
self.set_image(frame)
|
|
|
|
|
|
if draw_gt and 'gt_labels' in data_sample:
|
|
|
gt_labels = data_sample.gt_label
|
|
|
idx = gt_labels.tolist()
|
|
|
class_labels = [''] * len(idx)
|
|
|
if classes is not None:
|
|
|
class_labels = [f' ({classes[i]})' for i in idx]
|
|
|
labels = [
|
|
|
str(idx[i]) + class_labels[i] for i in range(len(idx))
|
|
|
]
|
|
|
prefix = 'Ground truth: '
|
|
|
texts.append(prefix + ('\n' + ' ' * len(prefix)).join(labels))
|
|
|
|
|
|
if draw_pred and 'pred_labels' in data_sample:
|
|
|
pred_labels = data_sample.pred_labels
|
|
|
idx = pred_labels.item.tolist()
|
|
|
score_labels = [''] * len(idx)
|
|
|
class_labels = [''] * len(idx)
|
|
|
if draw_score and 'score' in pred_labels:
|
|
|
score_labels = [
|
|
|
f', {pred_labels.score[i].item():.2f}' for i in idx
|
|
|
]
|
|
|
|
|
|
if classes is not None:
|
|
|
class_labels = [f' ({classes[i]})' for i in idx]
|
|
|
|
|
|
labels = [
|
|
|
str(idx[i]) + score_labels[i] + class_labels[i]
|
|
|
for i in range(len(idx))
|
|
|
]
|
|
|
prefix = 'Prediction: '
|
|
|
texts.append(prefix + ('\n' + ' ' * len(prefix)).join(labels))
|
|
|
|
|
|
img_scale = _get_adaptive_scale(frame.shape[:2])
|
|
|
_text_cfg = {
|
|
|
'positions':
|
|
|
np.array([(img_scale * 5, ) * 2]).astype(np.int32),
|
|
|
'font_sizes': int(img_scale * 7),
|
|
|
'font_families': 'monospace',
|
|
|
'colors': 'white',
|
|
|
'bboxes': dict(facecolor='black', alpha=0.5, boxstyle='Round'),
|
|
|
}
|
|
|
_text_cfg.update(text_cfg)
|
|
|
self.draw_texts('\n'.join(texts), **_text_cfg)
|
|
|
drawn_img = self.get_image()
|
|
|
resulted_video.append(drawn_img)
|
|
|
|
|
|
if show_frames:
|
|
|
frame_wait_time = 1. / fps
|
|
|
for frame_idx, drawn_img in enumerate(resulted_video):
|
|
|
frame_name = 'frame %d of %s' % (frame_idx + 1, name)
|
|
|
if frame_idx < len(resulted_video) - 1:
|
|
|
wait_time = frame_wait_time
|
|
|
else:
|
|
|
wait_time = wait_time
|
|
|
self.show(
|
|
|
drawn_img[:, :, ::-1],
|
|
|
win_name=frame_name,
|
|
|
wait_time=wait_time)
|
|
|
|
|
|
resulted_video = np.array(resulted_video)
|
|
|
if out_path is not None:
|
|
|
save_dir, save_name = osp.split(out_path)
|
|
|
vis_backend_cfg = dict(type='LocalVisBackend', save_dir=save_dir)
|
|
|
tmp_local_vis_backend = VISBACKENDS.build(vis_backend_cfg)
|
|
|
tmp_local_vis_backend.add_video(
|
|
|
save_name,
|
|
|
resulted_video,
|
|
|
step=step,
|
|
|
fps=fps,
|
|
|
out_type=out_type)
|
|
|
else:
|
|
|
self.add_video(
|
|
|
name, resulted_video, step=step, fps=fps, out_type=out_type)
|
|
|
return resulted_video
|
|
|
|
|
|
@master_only
|
|
|
def add_video(
|
|
|
self,
|
|
|
name: str,
|
|
|
image: np.ndarray,
|
|
|
step: int = 0,
|
|
|
fps: int = 4,
|
|
|
out_type: str = 'img',
|
|
|
) -> None:
|
|
|
"""Record the image.
|
|
|
|
|
|
Args:
|
|
|
name (str): The image identifier.
|
|
|
image (np.ndarray, optional): The image to be saved. The format
|
|
|
should be RGB. Default to None.
|
|
|
step (int): Global step value to record. Default to 0.
|
|
|
fps (int): Frames per second for saving video. Defaults to 4.
|
|
|
out_type (str): Output format type, choose from 'img', 'gif',
|
|
|
'video'. Defaults to ``'img'``.
|
|
|
"""
|
|
|
for vis_backend in self._vis_backends.values():
|
|
|
vis_backend.add_video(
|
|
|
name, image, step=step, fps=fps,
|
|
|
out_type=out_type)
|
|
|
|
