|
|
|
"""Contains the visualizer to visualize images by composing them as a gird.""" |
|
|
|
from ..image_utils import get_blank_image |
|
from ..image_utils import get_grid_shape |
|
from ..image_utils import parse_image_size |
|
from ..image_utils import load_image |
|
from ..image_utils import save_image |
|
from ..image_utils import resize_image |
|
from ..image_utils import list_images_from_dir |
|
|
|
__all__ = ['GridVisualizer'] |
|
|
|
|
|
class GridVisualizer(object): |
|
"""Defines the visualizer that visualizes images as a grid. |
|
|
|
Basically, given a collection of images, this visualizer stitches them one |
|
by one. Notably, this class also supports adding spaces between images, |
|
adding borders around images, and using white/black background. |
|
|
|
Example: |
|
|
|
grid = GridVisualizer(num_rows, num_cols) |
|
for i in range(num_rows): |
|
for j in range(num_cols): |
|
grid.add(i, j, image) |
|
grid.save('visualize.jpg') |
|
""" |
|
|
|
def __init__(self, |
|
grid_size=0, |
|
num_rows=0, |
|
num_cols=0, |
|
is_portrait=False, |
|
image_size=None, |
|
image_channels=0, |
|
row_spacing=0, |
|
col_spacing=0, |
|
border_left=0, |
|
border_right=0, |
|
border_top=0, |
|
border_bottom=0, |
|
use_black_background=True): |
|
"""Initializes the grid visualizer. |
|
|
|
Args: |
|
grid_size: Total number of cells, i.e., height * width. (default: 0) |
|
num_rows: Number of rows. (default: 0) |
|
num_cols: Number of columns. (default: 0) |
|
is_portrait: Whether the grid should be portrait or landscape. |
|
This is only used when it requires to compute `num_rows` and |
|
`num_cols` automatically. See function `get_grid_shape()` in |
|
file `./image_utils.py` for details. (default: False) |
|
image_size: Size to visualize each image. (default: 0) |
|
image_channels: Number of image channels. (default: 0) |
|
row_spacing: Spacing between rows. (default: 0) |
|
col_spacing: Spacing between columns. (default: 0) |
|
border_left: Width of left border. (default: 0) |
|
border_right: Width of right border. (default: 0) |
|
border_top: Width of top border. (default: 0) |
|
border_bottom: Width of bottom border. (default: 0) |
|
use_black_background: Whether to use black background. |
|
(default: True) |
|
""" |
|
self.reset(grid_size, num_rows, num_cols, is_portrait) |
|
self.set_image_size(image_size) |
|
self.set_image_channels(image_channels) |
|
self.set_row_spacing(row_spacing) |
|
self.set_col_spacing(col_spacing) |
|
self.set_border_left(border_left) |
|
self.set_border_right(border_right) |
|
self.set_border_top(border_top) |
|
self.set_border_bottom(border_bottom) |
|
self.set_background(use_black_background) |
|
self.grid = None |
|
|
|
def reset(self, |
|
grid_size=0, |
|
num_rows=0, |
|
num_cols=0, |
|
is_portrait=False): |
|
"""Resets the grid shape, i.e., number of rows/columns.""" |
|
if grid_size > 0: |
|
num_rows, num_cols = get_grid_shape(grid_size, |
|
height=num_rows, |
|
width=num_cols, |
|
is_portrait=is_portrait) |
|
self.grid_size = num_rows * num_cols |
|
self.num_rows = num_rows |
|
self.num_cols = num_cols |
|
self.grid = None |
|
|
|
def set_image_size(self, image_size=None): |
|
"""Sets the image size of each cell in the grid.""" |
|
height, width = parse_image_size(image_size) |
|
self.image_height = height |
|
self.image_width = width |
|
|
|
def set_image_channels(self, image_channels=0): |
|
"""Sets the number of channels of the grid.""" |
|
self.image_channels = image_channels |
|
|
|
def set_row_spacing(self, row_spacing=0): |
|
"""Sets the spacing between grid rows.""" |
|
self.row_spacing = row_spacing |
|
|
|
def set_col_spacing(self, col_spacing=0): |
|
"""Sets the spacing between grid columns.""" |
|
self.col_spacing = col_spacing |
|
|
|
def set_border_left(self, border_left=0): |
|
"""Sets the width of the left border of the grid.""" |
|
self.border_left = border_left |
|
|
|
def set_border_right(self, border_right=0): |
|
"""Sets the width of the right border of the grid.""" |
|
self.border_right = border_right |
|
|
|
def set_border_top(self, border_top=0): |
|
"""Sets the width of the top border of the grid.""" |
|
self.border_top = border_top |
|
|
|
def set_border_bottom(self, border_bottom=0): |
|
"""Sets the width of the bottom border of the grid.""" |
|
self.border_bottom = border_bottom |
|
|
|
def set_background(self, use_black=True): |
|
"""Sets the grid background.""" |
|
self.use_black_background = use_black |
|
|
|
def init_grid(self): |
|
"""Initializes the grid with a blank image.""" |
|
assert self.num_rows > 0 |
|
assert self.num_cols > 0 |
|
assert self.image_height > 0 |
|
assert self.image_width > 0 |
|
assert self.image_channels > 0 |
|
grid_height = (self.image_height * self.num_rows + |
|
self.row_spacing * (self.num_rows - 1) + |
|
self.border_top + self.border_bottom) |
|
grid_width = (self.image_width * self.num_cols + |
|
self.col_spacing * (self.num_cols - 1) + |
|
self.border_left + self.border_right) |
|
self.grid = get_blank_image(grid_height, grid_width, |
|
channels=self.image_channels, |
|
use_black=self.use_black_background) |
|
|
|
def add(self, i, j, image): |
|
"""Adds an image into the grid. |
|
|
|
NOTE: The input image is assumed to be with `RGB` channel order. |
|
""" |
|
channels = 1 if image.ndim == 2 else image.shape[2] |
|
if self.grid is None: |
|
height, width = image.shape[0:2] |
|
height = self.image_height or height |
|
width = self.image_width or width |
|
channels = self.image_channels or channels |
|
self.set_image_size((height, width)) |
|
self.set_image_channels(channels) |
|
self.init_grid() |
|
if image.shape[0:2] != (self.image_height, self.image_width): |
|
image = resize_image(image, (self.image_width, self.image_height)) |
|
y = self.border_top + i * (self.image_height + self.row_spacing) |
|
x = self.border_left + j * (self.image_width + self.col_spacing) |
|
self.grid[y:y + self.image_height, |
|
x:x + self.image_width, |
|
:channels] = image |
|
|
|
def visualize_collection(self, |
|
images, |
|
save_path=None, |
|
num_rows=0, |
|
num_cols=0, |
|
is_portrait=False, |
|
is_row_major=True): |
|
"""Visualizes a collection of images one by one.""" |
|
self.grid = None |
|
self.reset(grid_size=len(images), |
|
num_rows=num_rows, |
|
num_cols=num_cols, |
|
is_portrait=is_portrait) |
|
for idx, image in enumerate(images): |
|
if is_row_major: |
|
row_idx, col_idx = divmod(idx, self.num_cols) |
|
else: |
|
col_idx, row_idx = divmod(idx, self.num_rows) |
|
self.add(row_idx, col_idx, image) |
|
if save_path: |
|
self.save(save_path) |
|
|
|
def visualize_list(self, |
|
image_list, |
|
save_path=None, |
|
num_rows=0, |
|
num_cols=0, |
|
is_portrait=False, |
|
is_row_major=True): |
|
"""Visualizes a list of image files.""" |
|
self.grid = None |
|
self.reset(grid_size=len(image_list), |
|
num_rows=num_rows, |
|
num_cols=num_cols, |
|
is_portrait=is_portrait) |
|
for idx, filename in enumerate(image_list): |
|
image = load_image(filename) |
|
if is_row_major: |
|
row_idx, col_idx = divmod(idx, self.num_cols) |
|
else: |
|
col_idx, row_idx = divmod(idx, self.num_rows) |
|
self.add(row_idx, col_idx, image) |
|
if save_path: |
|
self.save(save_path) |
|
|
|
def visualize_directory(self, |
|
directory, |
|
save_path=None, |
|
num_rows=0, |
|
num_cols=0, |
|
is_portrait=False, |
|
is_row_major=True): |
|
"""Visualizes all images under a directory.""" |
|
image_list = list_images_from_dir(directory) |
|
self.visualize_list(image_list=image_list, |
|
save_path=save_path, |
|
num_rows=num_rows, |
|
num_cols=num_cols, |
|
is_portrait=is_portrait, |
|
is_row_major=is_row_major) |
|
|
|
def save(self, path): |
|
"""Saves the grid.""" |
|
save_image(path, self.grid) |
|
|
