File size: 5,043 Bytes
d670799 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
# Copyright (c) OpenMMLab. All rights reserved.
import os
import os.path as osp
from typing import Optional
import cv2
import numpy as np
from mmengine.visualization import (LocalVisBackend, TensorboardVisBackend,
WandbVisBackend)
from mmengine.visualization.vis_backend import force_init_env
from mmaction.registry import VISBACKENDS
try:
import wandb
except ImportError:
pass
@VISBACKENDS.register_module()
class LocalVisBackend(LocalVisBackend):
"""Local visualization backend class with video support.
See mmengine.visualization.LocalVisBackend for more details.
"""
@force_init_env
def add_video(self,
name: str,
frames: np.ndarray,
step: int = 0,
fps: Optional[int] = 4,
out_type: Optional[int] = 'img',
**kwargs) -> None:
"""Record the frames of a video to disk.
Args:
name (str): The video identifier (frame folder).
frames (np.ndarray): The frames to be saved. The format
should be RGB. The shape should be (T, H, W, C).
step (int): Global step value to record. Defaults to 0.
out_type (str): Output format type, choose from 'img', 'gif',
'video'. Defaults to ``'img'``.
fps (int): Frames per second for saving video. Defaults to 4.
"""
assert frames.dtype == np.uint8
if out_type == 'img':
frames_dir = osp.join(self._save_dir, name, f'frames_{step}')
os.makedirs(frames_dir, exist_ok=True)
for idx, frame in enumerate(frames):
drawn_image = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
save_file_name = f'{idx}.png'
cv2.imwrite(osp.join(frames_dir, save_file_name), drawn_image)
else:
try:
from moviepy.editor import ImageSequenceClip
except ImportError:
raise ImportError('Please install moviepy to enable '
'output file.')
frames = [x[..., ::-1] for x in frames]
video_clips = ImageSequenceClip(frames, fps=fps)
name = osp.splitext(name)[0]
if out_type == 'gif':
out_path = osp.join(self._save_dir, name + '.gif')
video_clips.write_gif(out_path, logger=None)
elif out_type == 'video':
out_path = osp.join(self._save_dir, name + '.mp4')
video_clips.write_videofile(
out_path, remove_temp=True, logger=None)
@VISBACKENDS.register_module()
class WandbVisBackend(WandbVisBackend):
"""Wandb visualization backend class with video support. See
mmengine.visualization.WandbVisBackend for more details.
Note that this requires the ``wandb`` and ``moviepy`` package. A wandb
account login is also required at ``https://wandb.ai/authorize``.
"""
@force_init_env
def add_video(self,
name: str,
frames: np.ndarray,
fps: int = 4,
**kwargs) -> None:
"""Record the frames of a video to wandb.
Note that this requires the ``moviepy`` package.
Args:
name (str): The video identifier (frame folder).
frames (np.ndarray): The frames to be saved. The format
should be RGB. The shape should be (T, H, W, C).
step is a useless parameter that Wandb does not need.
fps (int): Frames per second. Defaults to 4.
"""
frames = frames.transpose(0, 3, 1, 2)
self._wandb.log({'video': wandb.Video(frames, fps=fps, format='gif')})
@VISBACKENDS.register_module()
class TensorboardVisBackend(TensorboardVisBackend):
"""Tensorboard visualization backend class with video support. See
mmengine.visualization.TensorboardVisBackend for more details.
Note that this requires the ``future`` and ``tensorboard`` package.
"""
@force_init_env
def add_video(self,
name: str,
frames: np.ndarray,
step: int = 0,
fps: int = 4,
**kwargs) -> None:
"""Record the frames of a video to tensorboard.
Note that this requires the ``moviepy`` package.
Args:
name (str): The video identifier (frame folder).
frames (np.ndarray): The frames to be saved. The format
should be RGB. The shape should be (T, H, W, C).
step (int): Global step value to record. Defaults to 0.
fps (int): Frames per second. Defaults to 4.
"""
frames = frames.transpose(0, 3, 1, 2)
frames = frames.reshape(1, *frames.shape)
self._tensorboard.add_video(name, frames, global_step=step, fps=fps)
|