|
|
|
"""Contains the visualizer to visualize images as a video. |
|
|
|
This file relies on `FFmpeg`. Use `sudo apt-get install ffmpeg` and |
|
`brew install ffmpeg` to install on Ubuntu and MacOS respectively. |
|
""" |
|
|
|
import os.path |
|
from skvideo.io import FFmpegWriter |
|
from skvideo.io import FFmpegReader |
|
|
|
from ..image_utils import parse_image_size |
|
from ..image_utils import load_image |
|
from ..image_utils import resize_image |
|
from ..image_utils import list_images_from_dir |
|
|
|
__all__ = ['VideoVisualizer', 'VideoReader'] |
|
|
|
|
|
class VideoVisualizer(object): |
|
"""Defines the video visualizer that presents images as a video.""" |
|
|
|
def __init__(self, |
|
path=None, |
|
frame_size=None, |
|
fps=25.0, |
|
codec='libx264', |
|
pix_fmt='yuv420p', |
|
crf=1): |
|
"""Initializes the video visualizer. |
|
|
|
Args: |
|
path: Path to write the video. (default: None) |
|
frame_size: Frame size, i.e., (height, width). (default: None) |
|
fps: Frames per second. (default: 24) |
|
codec: Codec. (default: `libx264`) |
|
pix_fmt: Pixel format. (default: `yuv420p`) |
|
crf: Constant rate factor, which controls the compression. The |
|
larger this field is, the higher compression and lower quality. |
|
`0` means no compression and consequently the highest quality. |
|
To enable QuickTime playing (requires YUV to be 4:2:0, but |
|
`crf = 0` results YUV to be 4:4:4), please set this field as |
|
at least 1. (default: 1) |
|
""" |
|
self.set_path(path) |
|
self.set_frame_size(frame_size) |
|
self.set_fps(fps) |
|
self.set_codec(codec) |
|
self.set_pix_fmt(pix_fmt) |
|
self.set_crf(crf) |
|
self.video = None |
|
|
|
def set_path(self, path=None): |
|
"""Sets the path to save the video.""" |
|
self.path = path |
|
|
|
def set_frame_size(self, frame_size=None): |
|
"""Sets the video frame size.""" |
|
height, width = parse_image_size(frame_size) |
|
self.frame_height = height |
|
self.frame_width = width |
|
|
|
def set_fps(self, fps=25.0): |
|
"""Sets the FPS (frame per second) of the video.""" |
|
self.fps = fps |
|
|
|
def set_codec(self, codec='libx264'): |
|
"""Sets the video codec.""" |
|
self.codec = codec |
|
|
|
def set_pix_fmt(self, pix_fmt='yuv420p'): |
|
"""Sets the video pixel format.""" |
|
self.pix_fmt = pix_fmt |
|
|
|
def set_crf(self, crf=1): |
|
"""Sets the CRF (constant rate factor) of the video.""" |
|
self.crf = crf |
|
|
|
def init_video(self): |
|
"""Initializes an empty video with expected settings.""" |
|
assert not os.path.exists(self.path), f'Video `{self.path}` existed!' |
|
assert self.frame_height > 0 |
|
assert self.frame_width > 0 |
|
|
|
video_setting = { |
|
'-r': f'{self.fps:.2f}', |
|
'-s': f'{self.frame_width}x{self.frame_height}', |
|
'-vcodec': f'{self.codec}', |
|
'-crf': f'{self.crf}', |
|
'-pix_fmt': f'{self.pix_fmt}', |
|
} |
|
self.video = FFmpegWriter(self.path, outputdict=video_setting) |
|
|
|
def add(self, frame): |
|
"""Adds a frame into the video visualizer. |
|
|
|
NOTE: The input frame is assumed to be with `RGB` channel order. |
|
""" |
|
if self.video is None: |
|
height, width = frame.shape[0:2] |
|
height = self.frame_height or height |
|
width = self.frame_width or width |
|
self.set_frame_size((height, width)) |
|
self.init_video() |
|
if frame.shape[0:2] != (self.frame_height, self.frame_width): |
|
frame = resize_image(frame, (self.frame_width, self.frame_height)) |
|
self.video.writeFrame(frame) |
|
|
|
def visualize_collection(self, images, save_path=None): |
|
"""Visualizes a collection of images one by one.""" |
|
if save_path is not None and save_path != self.path: |
|
self.save() |
|
self.set_path(save_path) |
|
for image in images: |
|
self.add(image) |
|
self.save() |
|
|
|
def visualize_list(self, image_list, save_path=None): |
|
"""Visualizes a list of image files.""" |
|
if save_path is not None and save_path != self.path: |
|
self.save() |
|
self.set_path(save_path) |
|
for filename in image_list: |
|
image = load_image(filename) |
|
self.add(image) |
|
self.save() |
|
|
|
def visualize_directory(self, directory, save_path=None): |
|
"""Visualizes all images under a directory.""" |
|
image_list = list_images_from_dir(directory) |
|
self.visualize_list(image_list, save_path) |
|
|
|
def save(self): |
|
"""Saves the video by closing the file.""" |
|
if self.video is not None: |
|
self.video.close() |
|
self.video = None |
|
self.set_path(None) |
|
|
|
|
|
class VideoReader(object): |
|
"""Defines the video reader. |
|
|
|
This class can be used to read frames from a given video. |
|
|
|
NOTE: Each frame can be read only once. |
|
TODO: Fix this? |
|
""" |
|
|
|
def __init__(self, path, inputdict=None): |
|
"""Initializes the video reader by loading the video from disk.""" |
|
self.path = path |
|
self.video = FFmpegReader(path, inputdict=inputdict) |
|
|
|
self.length = self.video.inputframenum |
|
self.frame_height = self.video.inputheight |
|
self.frame_width = self.video.inputwidth |
|
self.fps = self.video.inputfps |
|
self.pix_fmt = self.video.pix_fmt |
|
|
|
def __del__(self): |
|
"""Releases the opened video.""" |
|
self.video.close() |
|
|
|
def read(self, image_size=None): |
|
"""Reads the next frame.""" |
|
frame = next(self.video.nextFrame()) |
|
height, width = parse_image_size(image_size) |
|
height = height or frame.shape[0] |
|
width = width or frame.shape[1] |
|
if frame.shape[0:2] != (height, width): |
|
frame = resize_image(frame, (width, height)) |
|
return frame |
|
|
