#!/usr/bin/env python3
from typing import List, Optional, Tuple, Union
import cv2
import numpy as np
__author__ = "Kevin Allen"
[docs]class ImageMux:
"""
Utility to create a customizable square grid of images, with labels and borders,
which will generate a single image from the grid at any time. Useful for
combining several debug images into one.
See the ``mil_vision_tools/image_mnux.py`` file for a usage example.
.. container:: operations
.. describe:: x[key] = img
Sets the ``key`` index in the panes of the grid to have ``img`` as the
source image for that pane. Equivalent to using :meth:`.set_image`.
.. describe:: x()
Returns the grid image, with decorations. Equivalent to calling :meth:`.get_image`.
Attributes:
size (np.ndarray): Tuple representing the size of the grid image, in pixels.
shape (np.ndarray): Tuple containing number of ``(rows, cols)``
of smaller images in the grid.
labels (List[str]): List of strings or ``None`` of length ``shape[0] * shape[1]``.
Each label corresponds to one image in the grid.
pane_size (np.ndarray): A numpy array of ``(height, width)`` representing the size
of each pane, in pixels.
keep_ratio (bool): If True, do not stretch image to insert into grid pane
border_color (Tuple[int, int, int]): The color of the border to use, in
BGR format. Defaults to ``[255, 255, 255]``, or white.
border_thickness (int): The thickness of the border between the images.
Defaults to 1.
text_color (Tuple[int, int, int]): The color of the text to use for the image
labels. Defaults to ``[255, 255, 255]``, or white.
text_font (int): An OpenCV font to use for the image labels.
text_scale (int): Scaling factor for label text. Defaults to 1.
text_thickness (int): Thickness of the label text. Defaults to 2.
"""
def __init__(
self,
size: Tuple[int, int] = (480, 640),
shape: Tuple[int, int] = (2, 2),
labels: Optional[List[str]] = None,
keep_ratio: bool = True,
border_color: Tuple[int, int, int] = (255, 255, 255),
border_thickness: int = 1,
text_color: Tuple[int, int, int] = (255, 255, 255),
text_font: int = cv2.FONT_HERSHEY_COMPLEX_SMALL,
text_scale: int = 1,
text_thickness: int = 2,
):
"""
Args:
size (Tuple[int, int]): Tuple representing the size of the grid image,
in pixels.
shape (Tuple[int, int]): Tuple containing number of ``(rows, cols)``
of smaller images in the grid.
labels (Optional[List[str]]): Optional list of strings of length ``shape[0] * shape[1]``.
Each label corresponds to one image in the grid.
keep_ratio (bool): If True, do not stretch image to insert into grid pane
border_color (Tuple[int, int, int]): The color of the border to use, in
BGR format. Defaults to ``[255, 255, 255]``, or white.
border_thickness (int): The thickness of the border between the images.
Defaults to 1.
text_color (Tuple[int, int, int]): The color of the text to use for the image
labels. Defaults to ``[255, 255, 255]``, or white.
text_font (int): An OpenCV font to use for the image labels.
text_scale (int): Scaling factor for label text. Defaults to 1.
text_thickness (int): Thickness of the label text. Defaults to 2.
"""
self.size = np.array(size, dtype=np.uint)
self.shape = np.array(shape, dtype=np.uint)
self.keep_ratio = keep_ratio
self.pane_size = np.array(self.size / self.shape, dtype=int)
self.border_color = border_color
self.border_thickness = border_thickness
self.text_color = text_color
self.text_font = text_font
self.text_scale = text_scale
self.text_thickness = text_thickness
# If labels not specified, fill a list with None's
if labels is None:
self.labels = [None for _ in range(self.shape[0] * self.shape[1])]
else:
assert len(labels) == self.shape[0] * self.shape[1], "not enough labels"
self.labels = labels
self._image = np.zeros((size[0], size[1], 3), dtype=np.uint8)
def _index_to_tuple(self, index: int) -> Tuple[int, int]:
"""
Internal helper function, returns row, col index
from a single index integer
"""
return (int(index / self.shape[1]), int(index % self.shape[1]))
def _apply_decorations(self) -> None:
"""
Internal helper function, adds border lines and label text to internal image.
"""
# Add border if thickness > 0
if self.border_thickness > 0:
for row in range(1, self.shape[0]): # Add horizontal line for rows 1 - m
y = int(self.pane_size[0] * row)
cv2.line(
self._image,
(0, y),
(self.size[1], y),
self.border_color,
self.border_thickness,
)
for col in range(1, self.shape[1]): # Add vertical line for rows 1 - n
x = int(self.pane_size[1] * col)
cv2.line(
self._image,
(x, 0),
(x, self.size[0]),
self.border_color,
self.border_thickness,
)
# Add label text for each pane if it is not None
for i, label in enumerate(self.labels):
if label is None:
continue
tup = self._index_to_tuple(i)
(text_width, text_height), _ = cv2.getTextSize(
label,
self.text_font,
self.text_scale,
self.text_thickness,
)
x = int(self.pane_size[1] * tup[1])
y = int(self.pane_size[0] * tup[0] + text_height)
# Adjust text position to not overlap border
if tup[0] != 0:
y += self.border_thickness
if tup[1] != 0:
x += self.border_thickness
cv2.putText(
self._image,
label,
(x, y),
self.text_font,
self.text_scale,
self.text_color,
self.text_thickness,
)
[docs] def set_image(self, key: Union[int, Tuple[int, int]], img: np.ndarray):
"""
Sets the content of one pane in the image grid.
Args:
key (Union[int, Tuple[int, int]]): The index of the pane to set to
the update. If an integer, the pane at that index is updated (counting
left to right, then top to bottom). If a tuple, then the pane at ``(row, col)``
is updated.
img (np.ndarray): Array with shape ``(m, n, 3)`` or ``(m, n, 1)`` representing
the image to insert in the pane specified in key. If a one-channel
image, first convert grayscale to BGR. If :attr:`.keep_ratio`
was True in constructor, will add black bars as necessary to fill pane.
Otherwise, use standard cv2.resize to fit img into pane.
Raises:
AssertionError: If key is wrong type of out of bounds.
"""
assert isinstance(img, np.ndarray), "img must be numpy array"
# If image is grayscale, convert to 3 channel
if len(img.shape) == 2 or img.shape[2] == 1:
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
if isinstance(key, int): # Accept a single index, ex 5 -> (2, 1)
key = self._index_to_tuple(key)
assert isinstance(key, tuple), "must be tuple"
assert len(key) == 2, "index best be 2D"
assert key[0] < self.shape[0] and key[1] < self.shape[1], "out of bounds"
rows = slice(key[0] * self.pane_size[0], (key[0] + 1) * self.pane_size[0])
cols = slice(key[1] * self.pane_size[1], (key[1] + 1) * self.pane_size[1])
if self.keep_ratio:
row_count = rows.stop - rows.start
col_count = cols.stop - cols.start
ratio = np.array([img.shape[0] / row_count, img.shape[1] / col_count])
scale = 1 / np.max(ratio)
size = (int(img.shape[1] * scale), int(img.shape[0] * scale))
v_border = int((row_count - size[1]) / 2)
h_border = int((col_count - size[0]) / 2)
rows = slice(rows.start + v_border, rows.start + size[1] + v_border)
cols = slice(cols.start + h_border, cols.start + size[0] + h_border)
self._image[rows, cols] = cv2.resize(img, size)
else:
size = (self.pane_size[1], self.pane_size[0])
self._image[rows, cols] = cv2.resize(img, size)
__setitem__ = set_image # Overload index [] operator to set image
[docs] def get_image(self) -> np.ndarray:
"""
Returns the image grid, with labels and borders.
Returns:
np.ndarray: The grid image, with decorations.
"""
self._apply_decorations()
return self._image
@property
def image(
self,
) -> np.ndarray:
"""
Returns the decorated grid image. Equivalent to calling :meth:`.get_image`.
Returns:
np.ndarray: The grid image, with decorations.
"""
return self.get_image()
__call__ = get_image # Overload () operator to access grid image
if __name__ == "__main__":
"""
ImageMux is intended to be used as a class, not an executable.
The following is an example of how to use it in a python program.
Creates a 2x2 grid of Raccoon images with labels, using some custom parameters.
To run this yourself, download some images, put them in $HOME/Pictures/[1.jpg, 2.jpg, 3.jpg, 4.jpg]
"""
import os
labels = [
"Chubby Raccoon",
"Kiddo Raccoons",
"wide",
"tall",
"big wide",
"big tall",
] # Create strings for labels
images = [
cv2.imread(os.path.join(os.environ["HOME"], "Pictures", str(i + 1) + ".jpg"))
for i in range(2)
]
# Add strange ratio white blocks to test keep_ratio flag
images.append(255 * np.ones((20, 201, 3), dtype=np.uint8)) # A small, wide image
images.append(255 * np.ones((200, 20, 3), dtype=np.uint8)) # A small, tall image
images.append(255 * np.ones((200, 2000, 3), dtype=np.uint8)) # A large, wide image
images.append(255 * np.ones((2000, 200, 3), dtype=np.uint8)) # A large, tall image
t = ImageMux(
size=(500, 900),
border_color=(0, 0, 255),
border_thickness=3,
shape=(3, 2),
labels=labels,
text_scale=1,
keep_ratio=True,
)
for i in range(len(images)):
t[i] = np.array(images[i])
cv2.imshow("Grid", t.image)
cv2.imshow("Grid2", t())
print("Press any key in GUI window to exit")
cv2.waitKey(0)