Source code for epic_kitchens.gulp.visualisation

import PIL.Image
import numpy as np
from abc import ABC
from moviepy.editor import ImageSequenceClip
from moviepy.video.io.html_tools import ipython_display
from typing import List, Union, Tuple
from IPython.display import HTML

from epic_kitchens.dataset.epic_dataset import EpicVideoDataset


def _grey_to_rgb(frames: np.ndarray) -> np.ndarray:
    """
    Convert frame(s) from gray (2D array) to RGB (3D array)

    Args:
        frames: A single frame or set of frames
    """
    if not (2 <= frames.ndim <= 3):
        raise ValueError(
            "Expected frame(s) to be a 2D array (single frame) "
            "or 3D array, but was {}D".format(frames.ndim)
        )
    n_channels = 3
    return np.repeat(np.array(frames)[..., np.newaxis], n_channels, axis=-1)


[docs]def clipify_rgb( frames: List[PIL.Image.Image], *, fps: float = 60. ) -> ImageSequenceClip: """ Args: frames: A list of frames fps: FPS of clip Returns ------- moviepy.editor.ImageSequenceClip Frames concatenated into clip """ if frames == 0: raise ValueError("Expected at least one frame, but received none") frames = [np.array(frame) for frame in frames] segment_clip = ImageSequenceClip(frames, fps=fps) return segment_clip
[docs]def clipify_flow( frames: List[PIL.Image.Image], *, fps: float = 30. ) -> ImageSequenceClip: """ Destack flow frames, join them side by side and then create a clip for display Args: frames: A list of alternating :math:`u`, :math:`v` flow frames to join into a video. Even indices should be :math:`u` flow frames, and odd indices, :math:`v` flow frames. fps: float, optional FPS of generated :py:class:`moviepy.editor.ImageSequenceClip` """ if frames == 0: raise ValueError("Expected at least one frame, but received none") frames = _grey_to_rgb(combine_flow_uv_frames(frames)) return ImageSequenceClip(list(frames), fps=fps)
[docs]def combine_flow_uv_frames( uv_frames: Union[List[PIL.Image.Image], np.ndarray], *, method="hstack", width_axis=2 ) -> np.ndarray: """Destack (u, v) frames and concatenate them side by side for display purposes""" if isinstance(uv_frames[0], PIL.Image.Image): uv_frames = list(map(np.array, uv_frames)) u_frames = np.array(uv_frames[::2]) v_frames = np.array(uv_frames[1::2]) return hstack_frames(u_frames, v_frames, width_axis=width_axis)
[docs]def hstack_frames(*frame_sequences: np.ndarray, width_axis: int = 2) -> np.ndarray: print([f.shape for f in frame_sequences]) return np.concatenate(frame_sequences, axis=width_axis)
[docs]class Visualiser(ABC): def __init__(self, dataset: EpicVideoDataset) -> None: self.dataset = dataset
[docs] def show(self, uid: Union[int, str], **kwargs) -> HTML: """ Show the given video corresponding to ``uid`` in a HTML5 video element. Args: uid: UID of video segment fps (float, optional): FPS of video sequence """ frames = self.dataset.load_frames(self.dataset[uid]) clip = self._clipify_frames(frames, **kwargs) return ipython_display(clip)
def _clipify_frames( self, frames: List[PIL.Image.Image], fps: float = 60. ) -> ImageSequenceClip: raise NotImplementedError()
[docs]class RgbVisualiser(Visualiser): """Visualiser for video dataset containing RGB frames""" def _clipify_frames(self, frames: List[PIL.Image.Image], fps: float = 60.): return clipify_rgb(frames, fps=fps)
[docs]class FlowVisualiser(Visualiser): """Visualiser for video dataset containing optical flow :math:`(u, v)` frames""" def _clipify_frames(self, frames: List[PIL.Image.Image], fps: float = 30): return clipify_flow(frames, fps=fps)