Revive `create_tiles` implementation and all related functions
This commit is contained in:
jirka
Jirka Borovec
parent
bc521bc2d0
commit
650ca0ae2f
|
|
@ -1,8 +1,12 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
import math
|
||||
import os
|
||||
import shutil
|
||||
from typing import Any, Literal
|
||||
from collections.abc import Callable
|
||||
from functools import partial
|
||||
from typing import Literal
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
|
@ -12,10 +16,18 @@ from PIL import Image
|
|||
|
||||
from supervision.annotators.base import ImageType
|
||||
from supervision.draw.color import Color, unify_to_bgr
|
||||
from supervision.draw.utils import calculate_optimal_text_scale, draw_text
|
||||
from supervision.geometry.core import Point
|
||||
from supervision.utils.conversion import (
|
||||
cv2_to_pillow,
|
||||
ensure_cv2_image_for_standalone_function,
|
||||
images_to_cv2,
|
||||
)
|
||||
from supervision.utils.iterables import create_batches, fill
|
||||
|
||||
RelativePosition = Literal["top", "bottom"]
|
||||
|
||||
MAX_COLUMNS_FOR_SINGLE_ROW_GRID = 3
|
||||
|
||||
|
||||
def crop_image(
|
||||
|
|
@ -654,6 +666,336 @@ def create_tiles(
|
|||
titles_padding: int = 10,
|
||||
titles_text_font: int = cv2.FONT_HERSHEY_SIMPLEX,
|
||||
titles_background_color: tuple[int, int, int] | Color = Color.from_hex("#D9D9D9"),
|
||||
default_title_placement: Any = "top",
|
||||
default_title_placement: RelativePosition = "top",
|
||||
) -> ImageType:
|
||||
return None
|
||||
"""
|
||||
Creates tiles mosaic from input images, automating grid placement and
|
||||
converting images to common resolution maintaining aspect ratio. It is
|
||||
also possible to render text titles on tiles, using optional set of
|
||||
parameters specifying text drawing (see parameters description).
|
||||
|
||||
Automated grid placement will try to maintain square shape of grid
|
||||
(with size being the nearest integer square root of #images), up to two exceptions:
|
||||
* if there are up to 3 images - images will be displayed in single row
|
||||
* if square-grid placement causes last row to be empty - number of rows is trimmed
|
||||
until last row has at least one image
|
||||
|
||||
Args:
|
||||
images (List[ImageType]): Images to create tiles. Elements can be either
|
||||
np.ndarray or PIL.Image, common representation will be agreed by the
|
||||
function.
|
||||
grid_size (Optional[Tuple[Optional[int], Optional[int]]]): Expected grid
|
||||
size in format (n_rows, n_cols). If not given - automated grid placement
|
||||
will be applied. One may also provide only one out of two elements of the
|
||||
tuple - then grid will be created with either n_rows or n_cols fixed,
|
||||
leaving the other dimension to be adjusted by the number of images
|
||||
single_tile_size (Optional[Tuple[int, int]]): sizeof a single tile element
|
||||
provided in (width, height) format. If not given - size of tile will be
|
||||
automatically calculated based on `tile_scaling` parameter.
|
||||
tile_scaling (Literal["min", "max", "avg"]): If `single_tile_size` is not
|
||||
given - parameter will be used to calculate tile size - using
|
||||
min / max / avg size of image provided in `images` list.
|
||||
tile_padding_color (Union[Tuple[int, int, int], sv.Color]): Color to be used in
|
||||
images letterbox procedure (while standardising tiles sizes) as a padding.
|
||||
If tuple provided - should be BGR.
|
||||
tile_margin (int): size of margin between tiles (in pixels)
|
||||
tile_margin_color (Union[Tuple[int, int, int], sv.Color]): Color of tile margin.
|
||||
If tuple provided - should be BGR.
|
||||
return_type (Literal["auto", "cv2", "pillow"]): Parameter dictates the format of
|
||||
return image. One may choose specific type ("cv2" or "pillow") to enforce
|
||||
conversion. "auto" mode takes a majority vote between types of elements in
|
||||
`images` list - resolving draws in favour of OpenCV format. "auto" can be
|
||||
safely used when all input images are of the same type.
|
||||
titles (Optional[List[Optional[str]]]): Optional titles to be added to tiles.
|
||||
Elements of that list may be empty - then specific tile (in order presented
|
||||
in `images` parameter) will not be filled with title. It is possible to
|
||||
provide list of titles shorter than `images` - then remaining titles will
|
||||
be assumed empty.
|
||||
titles_anchors (Optional[Union[Point, List[Optional[Point]]]]): Parameter to
|
||||
specify anchor points for titles. It is possible to specify anchor either
|
||||
globally or for specific tiles (following order of `images`).
|
||||
If not given (either globally, or for specific element of the list),
|
||||
it will be calculated automatically based on `default_title_placement`.
|
||||
titles_color (Union[Tuple[int, int, int], Color]): Color of titles text.
|
||||
If tuple provided - should be BGR.
|
||||
titles_scale (Optional[float]): Scale of titles. If not provided - value will
|
||||
be calculated using `calculate_optimal_text_scale(...)`.
|
||||
titles_thickness (int): Thickness of titles text.
|
||||
titles_padding (int): Size of titles padding.
|
||||
titles_text_font (int): Font to be used to render titles. Must be integer
|
||||
constant representing OpenCV font.
|
||||
(See docs: https://docs.opencv.org/4.x/d6/d6e/group__imgproc__draw.html)
|
||||
titles_background_color (Union[Tuple[int, int, int], Color]): Color of title
|
||||
text padding.
|
||||
default_title_placement (Literal["top", "bottom"]): Parameter specifies title
|
||||
anchor placement in case if explicit anchor is not provided.
|
||||
|
||||
Returns:
|
||||
ImageType: Image with all input images located in tails grid. The output type is
|
||||
determined by `return_type` parameter.
|
||||
|
||||
Raises:
|
||||
ValueError: In case when input images list is empty, provided `grid_size` is too
|
||||
small to fit all images, `tile_scaling` mode is invalid.
|
||||
"""
|
||||
if len(images) == 0:
|
||||
raise ValueError("Could not create image tiles from empty list of images.")
|
||||
if return_type == "auto":
|
||||
return_type = _negotiate_tiles_format(images=images)
|
||||
tile_padding_color = unify_to_bgr(color=tile_padding_color)
|
||||
tile_margin_color = unify_to_bgr(color=tile_margin_color)
|
||||
images = images_to_cv2(images=images)
|
||||
if single_tile_size is None:
|
||||
single_tile_size = _aggregate_images_shape(images=images, mode=tile_scaling)
|
||||
resized_images = [
|
||||
letterbox_image(
|
||||
image=i, resolution_wh=single_tile_size, color=tile_padding_color
|
||||
)
|
||||
for i in images
|
||||
]
|
||||
grid_size = _establish_grid_size(images=images, grid_size=grid_size)
|
||||
if len(images) > grid_size[0] * grid_size[1]:
|
||||
raise ValueError(
|
||||
f"Could not place {len(images)} in grid with size: {grid_size}."
|
||||
)
|
||||
if titles is not None:
|
||||
titles = fill(sequence=titles, desired_size=len(images), content=None)
|
||||
titles_anchors = (
|
||||
[titles_anchors]
|
||||
if not issubclass(type(titles_anchors), list)
|
||||
else titles_anchors
|
||||
)
|
||||
titles_anchors = fill(
|
||||
sequence=titles_anchors, desired_size=len(images), content=None
|
||||
)
|
||||
titles_color = unify_to_bgr(color=titles_color)
|
||||
titles_background_color = unify_to_bgr(color=titles_background_color)
|
||||
tiles = _generate_tiles(
|
||||
images=resized_images,
|
||||
grid_size=grid_size,
|
||||
single_tile_size=single_tile_size,
|
||||
tile_padding_color=tile_padding_color,
|
||||
tile_margin=tile_margin,
|
||||
tile_margin_color=tile_margin_color,
|
||||
titles=titles,
|
||||
titles_anchors=titles_anchors,
|
||||
titles_color=titles_color,
|
||||
titles_scale=titles_scale,
|
||||
titles_thickness=titles_thickness,
|
||||
titles_padding=titles_padding,
|
||||
titles_text_font=titles_text_font,
|
||||
titles_background_color=titles_background_color,
|
||||
default_title_placement=default_title_placement,
|
||||
)
|
||||
if return_type == "pillow":
|
||||
tiles = cv2_to_pillow(image=tiles)
|
||||
return tiles
|
||||
|
||||
|
||||
def _negotiate_tiles_format(images: list[ImageType]) -> Literal["cv2", "pillow"]:
|
||||
number_of_np_arrays = sum(issubclass(type(i), np.ndarray) for i in images)
|
||||
if number_of_np_arrays >= (len(images) // 2):
|
||||
return "cv2"
|
||||
return "pillow"
|
||||
|
||||
|
||||
def _calculate_aggregated_images_shape(
|
||||
images: list[np.ndarray], aggregator: Callable[[list[int]], float]
|
||||
) -> tuple[int, int]:
|
||||
height = round(aggregator([i.shape[0] for i in images]))
|
||||
width = round(aggregator([i.shape[1] for i in images]))
|
||||
return width, height
|
||||
|
||||
|
||||
SHAPE_AGGREGATION_FUN = {
|
||||
"min": partial(_calculate_aggregated_images_shape, aggregator=np.min),
|
||||
"max": partial(_calculate_aggregated_images_shape, aggregator=np.max),
|
||||
"avg": partial(_calculate_aggregated_images_shape, aggregator=np.average),
|
||||
}
|
||||
|
||||
|
||||
def _aggregate_images_shape(
|
||||
images: list[np.ndarray], mode: Literal["min", "max", "avg"]
|
||||
) -> tuple[int, int]:
|
||||
if mode not in SHAPE_AGGREGATION_FUN:
|
||||
raise ValueError(
|
||||
f"Could not aggregate images shape - provided unknown mode: {mode}. "
|
||||
f"Supported modes: {list(SHAPE_AGGREGATION_FUN.keys())}."
|
||||
)
|
||||
return SHAPE_AGGREGATION_FUN[mode](images)
|
||||
|
||||
|
||||
def _establish_grid_size(
|
||||
images: list[np.ndarray], grid_size: tuple[int | None, int | None] | None
|
||||
) -> tuple[int, int]:
|
||||
if grid_size is None or all(e is None for e in grid_size):
|
||||
return _negotiate_grid_size(images=images)
|
||||
if grid_size[0] is None:
|
||||
return math.ceil(len(images) / grid_size[1]), grid_size[1]
|
||||
if grid_size[1] is None:
|
||||
return grid_size[0], math.ceil(len(images) / grid_size[0])
|
||||
return grid_size
|
||||
|
||||
|
||||
def _negotiate_grid_size(images: list[np.ndarray]) -> tuple[int, int]:
|
||||
if len(images) <= MAX_COLUMNS_FOR_SINGLE_ROW_GRID:
|
||||
return 1, len(images)
|
||||
nearest_sqrt = math.ceil(np.sqrt(len(images)))
|
||||
proposed_columns = nearest_sqrt
|
||||
proposed_rows = nearest_sqrt
|
||||
while proposed_columns * (proposed_rows - 1) >= len(images):
|
||||
proposed_rows -= 1
|
||||
return proposed_rows, proposed_columns
|
||||
|
||||
|
||||
def _generate_tiles(
|
||||
images: list[np.ndarray],
|
||||
grid_size: tuple[int, int],
|
||||
single_tile_size: tuple[int, int],
|
||||
tile_padding_color: tuple[int, int, int],
|
||||
tile_margin: int,
|
||||
tile_margin_color: tuple[int, int, int],
|
||||
titles: list[str | None] | None,
|
||||
titles_anchors: list[Point | None],
|
||||
titles_color: tuple[int, int, int],
|
||||
titles_scale: float | None,
|
||||
titles_thickness: int,
|
||||
titles_padding: int,
|
||||
titles_text_font: int,
|
||||
titles_background_color: tuple[int, int, int],
|
||||
default_title_placement: RelativePosition,
|
||||
) -> np.ndarray:
|
||||
images = _draw_texts(
|
||||
images=images,
|
||||
titles=titles,
|
||||
titles_anchors=titles_anchors,
|
||||
titles_color=titles_color,
|
||||
titles_scale=titles_scale,
|
||||
titles_thickness=titles_thickness,
|
||||
titles_padding=titles_padding,
|
||||
titles_text_font=titles_text_font,
|
||||
titles_background_color=titles_background_color,
|
||||
default_title_placement=default_title_placement,
|
||||
)
|
||||
rows, columns = grid_size
|
||||
tiles_elements = list(create_batches(sequence=images, batch_size=columns))
|
||||
while len(tiles_elements[-1]) < columns:
|
||||
tiles_elements[-1].append(
|
||||
_generate_color_image(shape=single_tile_size, color=tile_padding_color)
|
||||
)
|
||||
while len(tiles_elements) < rows:
|
||||
tiles_elements.append(
|
||||
[_generate_color_image(shape=single_tile_size, color=tile_padding_color)]
|
||||
* columns
|
||||
)
|
||||
return _merge_tiles_elements(
|
||||
tiles_elements=tiles_elements,
|
||||
grid_size=grid_size,
|
||||
single_tile_size=single_tile_size,
|
||||
tile_margin=tile_margin,
|
||||
tile_margin_color=tile_margin_color,
|
||||
)
|
||||
|
||||
|
||||
def _draw_texts(
|
||||
images: list[np.ndarray],
|
||||
titles: list[str | None] | None,
|
||||
titles_anchors: list[Point | None],
|
||||
titles_color: tuple[int, int, int],
|
||||
titles_scale: float | None,
|
||||
titles_thickness: int,
|
||||
titles_padding: int,
|
||||
titles_text_font: int,
|
||||
titles_background_color: tuple[int, int, int],
|
||||
default_title_placement: RelativePosition,
|
||||
) -> list[np.ndarray]:
|
||||
if titles is None:
|
||||
return images
|
||||
titles_anchors = _prepare_default_titles_anchors(
|
||||
images=images,
|
||||
titles_anchors=titles_anchors,
|
||||
default_title_placement=default_title_placement,
|
||||
)
|
||||
if titles_scale is None:
|
||||
image_height, image_width = images[0].shape[:2]
|
||||
titles_scale = calculate_optimal_text_scale(
|
||||
resolution_wh=(image_width, image_height)
|
||||
)
|
||||
result = []
|
||||
for image, text, anchor in zip(images, titles, titles_anchors):
|
||||
if text is None:
|
||||
result.append(image)
|
||||
continue
|
||||
processed_image = draw_text(
|
||||
scene=image,
|
||||
text=text,
|
||||
text_anchor=anchor,
|
||||
text_color=Color.from_bgr_tuple(titles_color),
|
||||
text_scale=titles_scale,
|
||||
text_thickness=titles_thickness,
|
||||
text_padding=titles_padding,
|
||||
text_font=titles_text_font,
|
||||
background_color=Color.from_bgr_tuple(titles_background_color),
|
||||
)
|
||||
result.append(processed_image)
|
||||
return result
|
||||
|
||||
|
||||
def _prepare_default_titles_anchors(
|
||||
images: list[np.ndarray],
|
||||
titles_anchors: list[Point | None],
|
||||
default_title_placement: RelativePosition,
|
||||
) -> list[Point]:
|
||||
result = []
|
||||
for image, anchor in zip(images, titles_anchors):
|
||||
if anchor is not None:
|
||||
result.append(anchor)
|
||||
continue
|
||||
image_height, image_width = image.shape[:2]
|
||||
if default_title_placement == "top":
|
||||
default_anchor = Point(x=image_width / 2, y=image_height * 0.1)
|
||||
else:
|
||||
default_anchor = Point(x=image_width / 2, y=image_height * 0.9)
|
||||
result.append(default_anchor)
|
||||
return result
|
||||
|
||||
|
||||
def _merge_tiles_elements(
|
||||
tiles_elements: list[list[np.ndarray]],
|
||||
grid_size: tuple[int, int],
|
||||
single_tile_size: tuple[int, int],
|
||||
tile_margin: int,
|
||||
tile_margin_color: tuple[int, int, int],
|
||||
) -> np.ndarray:
|
||||
vertical_padding = (
|
||||
np.ones((single_tile_size[1], tile_margin, 3)) * tile_margin_color
|
||||
)
|
||||
merged_rows = [
|
||||
np.concatenate(
|
||||
list(
|
||||
itertools.chain.from_iterable(
|
||||
zip(row, [vertical_padding] * grid_size[1])
|
||||
)
|
||||
)[:-1],
|
||||
axis=1,
|
||||
)
|
||||
for row in tiles_elements
|
||||
]
|
||||
row_width = merged_rows[0].shape[1]
|
||||
horizontal_padding = (
|
||||
np.ones((tile_margin, row_width, 3), dtype=np.uint8) * tile_margin_color
|
||||
)
|
||||
rows_with_paddings = []
|
||||
for row in merged_rows:
|
||||
rows_with_paddings.append(row)
|
||||
rows_with_paddings.append(horizontal_padding)
|
||||
return np.concatenate(
|
||||
rows_with_paddings[:-1],
|
||||
axis=0,
|
||||
).astype(np.uint8)
|
||||
|
||||
|
||||
def _generate_color_image(
|
||||
shape: tuple[int, int], color: tuple[int, int, int]
|
||||
) -> np.ndarray:
|
||||
return np.ones((*shape[::-1], 3), dtype=np.uint8) * color
|
||||
|
|
|
|||
Loading…
Reference in New Issue