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