|
|
|
"""Unit test for visualizer.""" |
|
|
|
import os |
|
import skvideo.datasets |
|
|
|
from ..image_utils import save_image |
|
from . import GridVisualizer |
|
from . import HtmlVisualizer |
|
from . import HtmlReader |
|
from . import GifVisualizer |
|
from . import VideoVisualizer |
|
from . import VideoReader |
|
|
|
__all__ = ['test_visualizer'] |
|
|
|
_TEST_DIR = 'visualizer_test' |
|
|
|
|
|
def test_visualizer(test_dir=_TEST_DIR): |
|
"""Tests visualizers.""" |
|
print('========== Start Visualizer Test ==========') |
|
|
|
frame_dir = os.path.join(test_dir, 'test_frames') |
|
os.makedirs(frame_dir, exist_ok=True) |
|
|
|
print('===== Testing `VideoReader` =====') |
|
|
|
video_reader = VideoReader(skvideo.datasets.bigbuckbunny()) |
|
frame_height = video_reader.frame_height |
|
frame_width = video_reader.frame_width |
|
frame_size = (frame_height, frame_width) |
|
half_size = (frame_height // 2, frame_width // 2) |
|
|
|
for idx in range(80): |
|
frame = video_reader.read() |
|
save_image(os.path.join(frame_dir, f'{idx:02d}.png'), frame) |
|
|
|
print('===== Testing `GirdVisualizer` =====') |
|
grid_visualizer = GridVisualizer() |
|
grid_visualizer.set_row_spacing(30) |
|
grid_visualizer.set_col_spacing(30) |
|
grid_visualizer.set_background(use_black=True) |
|
path = os.path.join(test_dir, 'portrait_row_major_ori_space30_black.png') |
|
grid_visualizer.visualize_directory(frame_dir, path, |
|
is_portrait=True, is_row_major=True) |
|
path = os.path.join( |
|
test_dir, 'landscape_col_major_downsample_space15_white.png') |
|
grid_visualizer.set_image_size(half_size) |
|
grid_visualizer.set_row_spacing(15) |
|
grid_visualizer.set_col_spacing(15) |
|
grid_visualizer.set_background(use_black=False) |
|
grid_visualizer.visualize_directory(frame_dir, path, |
|
is_portrait=False, is_row_major=False) |
|
|
|
print('===== Testing `HtmlVisualizer` =====') |
|
html_visualizer = HtmlVisualizer() |
|
path = os.path.join(test_dir, 'portrait_col_major_ori.html') |
|
html_visualizer.visualize_directory(frame_dir, path, |
|
is_portrait=True, is_row_major=False) |
|
path = os.path.join(test_dir, 'landscape_row_major_downsample.html') |
|
html_visualizer.set_image_size(half_size) |
|
html_visualizer.visualize_directory(frame_dir, path, |
|
is_portrait=False, is_row_major=True) |
|
|
|
print('===== Testing `HtmlReader` =====') |
|
path = os.path.join(test_dir, 'landscape_row_major_downsample.html') |
|
html_reader = HtmlReader(path) |
|
for j in range(html_reader.num_cols): |
|
assert html_reader.get_header(j) == '' |
|
parsed_dir = os.path.join(test_dir, 'parsed_frames') |
|
os.makedirs(parsed_dir, exist_ok=True) |
|
for i in range(html_reader.num_rows): |
|
for j in range(html_reader.num_cols): |
|
idx = i * html_reader.num_cols + j |
|
assert html_reader.get_text(i, j).endswith(f'(index {idx:03d})') |
|
image = html_reader.get_image(i, j, image_size=frame_size) |
|
assert image.shape[0:2] == frame_size |
|
save_image(os.path.join(parsed_dir, f'{idx:02d}.png'), image) |
|
|
|
print('===== Testing `GifVisualizer` =====') |
|
gif_visualizer = GifVisualizer() |
|
path = os.path.join(test_dir, 'gif_ori.gif') |
|
gif_visualizer.visualize_directory(frame_dir, path) |
|
gif_visualizer.set_image_size(half_size) |
|
path = os.path.join(test_dir, 'gif_downsample.gif') |
|
gif_visualizer.visualize_directory(frame_dir, path) |
|
|
|
print('===== Testing `VideoVisualizer` =====') |
|
video_visualizer = VideoVisualizer() |
|
path = os.path.join(test_dir, 'video_ori.mp4') |
|
video_visualizer.visualize_directory(frame_dir, path) |
|
path = os.path.join(test_dir, 'video_downsample.mp4') |
|
video_visualizer.set_frame_size(half_size) |
|
video_visualizer.visualize_directory(frame_dir, path) |
|
|
|
print('========== Finish Visualizer Test ==========') |
|
|
