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