utils/video_visualizer.py

import os.path
from skvideo.io import FFmpegWriter
# 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
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


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 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]
            if height & 1:
                height -= 1
            if width & 1:
                width -= 1
            # 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)


if __name__ == '__main__':
    from glob import glob
    import cv2
    video_visualizer = VideoVisualizer(path='/home/yehhh/DiffBIR/DAVIS_bear.mp4',
                                       frame_size=None,
                                       fps=25.0)
    img_folder = "/home/yehhh/DiffBIR/inputs/bear"
    imgs = sorted(glob(img_folder + '/*.png'))
    imgs = sorted(glob(img_folder + '/*.jpg'))
    for img in imgs:
        image = cv2.imread(img)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        video_visualizer.add(image)
    video_visualizer.save()