Source code for mmaction.visualization.video_backend

# 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)