Add first draft of implementation for function creating image tiles
This commit is contained in:
Paweł Pęczek
parent
6208013dcf
commit
7eba338243
|
|
@ -27,7 +27,6 @@ repos:
|
|||
- id: mixed-line-ending
|
||||
|
||||
|
||||
|
||||
- repo: https://github.com/PyCQA/bandit
|
||||
rev: '1.7.8'
|
||||
hooks:
|
||||
|
|
|
|||
|
|
@ -2,7 +2,6 @@ from enum import Enum
|
|||
from functools import wraps
|
||||
from typing import Optional, Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
|
@ -10,6 +9,7 @@ from supervision.annotators.base import ImageType
|
|||
from supervision.detection.core import Detections
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
from supervision.geometry.core import Position
|
||||
from supervision.utils.image import cv2_to_pillow, pillow_to_cv2
|
||||
|
||||
|
||||
class ColorLookup(Enum):
|
||||
|
|
@ -125,12 +125,6 @@ class Trace:
|
|||
return self.xy[self.tracker_id == tracker_id]
|
||||
|
||||
|
||||
def pillow_to_cv2(image: Image.Image) -> np.ndarray:
|
||||
scene = np.array(image)
|
||||
scene = cv2.cvtColor(scene, cv2.COLOR_RGB2BGR)
|
||||
return scene
|
||||
|
||||
|
||||
def scene_to_annotator_img_type(annotate_func):
|
||||
"""
|
||||
Decorates `BaseAnnotator.annotate` implementations, converts scene to
|
||||
|
|
@ -146,9 +140,7 @@ def scene_to_annotator_img_type(annotate_func):
|
|||
if isinstance(scene, Image.Image):
|
||||
scene = pillow_to_cv2(scene)
|
||||
annotated = annotate_func(self, scene, *args, **kwargs)
|
||||
annotated = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
|
||||
annotated = Image.fromarray(annotated)
|
||||
return annotated
|
||||
return cv2_to_pillow(image=annotated)
|
||||
|
||||
raise ValueError(f"Unsupported image type: {type(scene)}")
|
||||
|
||||
|
|
|
|||
|
|
@ -1,11 +1,43 @@
|
|||
import itertools
|
||||
import math
|
||||
import os
|
||||
import shutil
|
||||
from typing import Optional, Tuple
|
||||
from functools import partial, wraps
|
||||
from typing import Callable, List, Literal, Optional, Tuple, Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from supervision.annotators.base import ImageType
|
||||
from supervision.draw.color import Color
|
||||
from supervision.utils.iterables import create_batches
|
||||
|
||||
MAX_COLUMNS_FOR_SINGLE_ROW_GRID = 3
|
||||
|
||||
|
||||
def adjust_image_to_cv2_processing(image_processing_fun):
|
||||
"""
|
||||
Decorates image processing functions that accept np.ndarray, converting `image` to
|
||||
np.ndarray, converts back when processing is complete.
|
||||
"""
|
||||
|
||||
@wraps(image_processing_fun)
|
||||
def wrapper(image: ImageType, *args, **kwargs):
|
||||
if isinstance(image, np.ndarray):
|
||||
return image_processing_fun(image, *args, **kwargs)
|
||||
|
||||
if isinstance(image, Image.Image):
|
||||
scene = pillow_to_cv2(image)
|
||||
annotated = image_processing_fun(scene, *args, **kwargs)
|
||||
return cv2_to_pillow(image=annotated)
|
||||
|
||||
raise ValueError(f"Unsupported image type: {type(image)}")
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
@adjust_image_to_cv2_processing
|
||||
def crop_image(image: np.ndarray, xyxy: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Crops the given image based on the given bounding box.
|
||||
|
|
@ -35,6 +67,7 @@ def crop_image(image: np.ndarray, xyxy: np.ndarray) -> np.ndarray:
|
|||
return image[y1:y2, x1:x2]
|
||||
|
||||
|
||||
@adjust_image_to_cv2_processing
|
||||
def resize_image(image: np.ndarray, scale_factor: float) -> np.ndarray:
|
||||
"""
|
||||
Resizes an image by a given scale factor using cv2.INTER_LINEAR interpolation.
|
||||
|
|
@ -167,3 +200,314 @@ class ImageSink:
|
|||
|
||||
def __exit__(self, exc_type, exc_value, exc_traceback):
|
||||
pass
|
||||
|
||||
|
||||
def create_tiles(
|
||||
images: List[ImageType],
|
||||
grid_size: Optional[Tuple[Optional[int], Optional[int]]] = None,
|
||||
single_tile_size: Optional[Tuple[int, int]] = None,
|
||||
tile_scaling: Literal["min", "max", "avg"] = "avg",
|
||||
tile_padding_color: Union[Tuple[int, int, int], Color] = Color.WHITE,
|
||||
tile_margin: int = 15,
|
||||
tile_margin_color: Union[Tuple[int, int, int], Color] = Color.BLACK,
|
||||
return_type: Literal["auto", "cv2", "pillow"] = "auto",
|
||||
) -> ImageType:
|
||||
"""
|
||||
Creates tiles mosaic from input images, automating grid placement and
|
||||
converting images to common resolution maintaining aspect ratio.
|
||||
|
||||
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.
|
||||
|
||||
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 = _color_to_bgr(color=tile_padding_color)
|
||||
tile_margin_color = _color_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, desired_size=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}."
|
||||
)
|
||||
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,
|
||||
)
|
||||
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: Optional[Tuple[Optional[int], Optional[int]]]
|
||||
) -> 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]
|
||||
return grid_size[0], math.ceil(len(images) / grid_size[0])
|
||||
|
||||
|
||||
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],
|
||||
) -> np.ndarray:
|
||||
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 _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
|
||||
|
||||
|
||||
@adjust_image_to_cv2_processing
|
||||
def letterbox_image(
|
||||
image: np.ndarray,
|
||||
desired_size: Tuple[int, int],
|
||||
color: Union[Tuple[int, int, int], Color] = (0, 0, 0),
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Resize and pad image to fit the desired size, preserving its aspect
|
||||
ratio, adding padding of given color if needed to maintain aspect ratio.
|
||||
|
||||
Parameters:
|
||||
- image (np.ndarray): Input image (type will be adjusted by decorator,
|
||||
you can provide PIL.Image)
|
||||
- desired_size (Tuple[int, int]): image size (width, height) representing
|
||||
the target dimensions.
|
||||
- color (Union[Tuple[int, int, int], Color]): the color to pad with - If
|
||||
tuple provided - should be BGR.
|
||||
|
||||
Returns:
|
||||
np.ndarray: letterboxed image (type may be adjusted to PIL.Image by
|
||||
decorator if function was called with PIL.Image)
|
||||
"""
|
||||
color = _color_to_bgr(color=color)
|
||||
resized_img = resize_image_keeping_aspect_ratio(
|
||||
image=image,
|
||||
desired_size=desired_size,
|
||||
)
|
||||
new_height, new_width = resized_img.shape[:2]
|
||||
top_padding = (desired_size[1] - new_height) // 2
|
||||
bottom_padding = desired_size[1] - new_height - top_padding
|
||||
left_padding = (desired_size[0] - new_width) // 2
|
||||
right_padding = desired_size[0] - new_width - left_padding
|
||||
return cv2.copyMakeBorder(
|
||||
resized_img,
|
||||
top_padding,
|
||||
bottom_padding,
|
||||
left_padding,
|
||||
right_padding,
|
||||
cv2.BORDER_CONSTANT,
|
||||
value=color,
|
||||
)
|
||||
|
||||
|
||||
@adjust_image_to_cv2_processing
|
||||
def resize_image_keeping_aspect_ratio(
|
||||
image: np.ndarray,
|
||||
desired_size: Tuple[int, int],
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Resize and pad image preserving its aspect ratio.
|
||||
|
||||
Parameters:
|
||||
- image (np.ndarray): Input image (type will be adjusted by decorator,
|
||||
you can provide PIL.Image)
|
||||
- desired_size (Tuple[int, int]): image size (width, height) representing the
|
||||
target dimensions. Parameter will be used to dictate maximum size of
|
||||
output image. Output size may be smaller - to preserve aspect ratio of original
|
||||
image.
|
||||
|
||||
Returns:
|
||||
np.ndarray: resized image (type may be adjusted to PIL.Image by decorator
|
||||
if function was called with PIL.Image)
|
||||
"""
|
||||
img_ratio = image.shape[1] / image.shape[0]
|
||||
desired_ratio = desired_size[0] / desired_size[1]
|
||||
if img_ratio >= desired_ratio:
|
||||
new_width = desired_size[0]
|
||||
new_height = int(desired_size[0] / img_ratio)
|
||||
else:
|
||||
new_height = desired_size[1]
|
||||
new_width = int(desired_size[1] * img_ratio)
|
||||
return cv2.resize(image, (new_width, new_height))
|
||||
|
||||
|
||||
def _color_to_bgr(color: Union[Tuple[int, int, int], Color]) -> Tuple[int, int, int]:
|
||||
if issubclass(type(color), Color):
|
||||
return color.as_bgr()
|
||||
return color
|
||||
|
||||
|
||||
def images_to_cv2(images: List[ImageType]) -> List[np.ndarray]:
|
||||
result = []
|
||||
for image in images:
|
||||
if issubclass(type(image), Image.Image):
|
||||
image = pillow_to_cv2(image=image)
|
||||
result.append(image)
|
||||
return result
|
||||
|
||||
|
||||
def pillow_to_cv2(image: Image.Image) -> np.ndarray:
|
||||
scene = np.array(image)
|
||||
scene = cv2.cvtColor(scene, cv2.COLOR_RGB2BGR)
|
||||
return scene
|
||||
|
||||
|
||||
def cv2_to_pillow(image: np.ndarray) -> Image.Image:
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
return Image.fromarray(image)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,17 @@
|
|||
from typing import Generator, Iterable, List, TypeVar
|
||||
|
||||
SequenceElement = TypeVar("SequenceElement")
|
||||
|
||||
|
||||
def create_batches(
|
||||
sequence: Iterable[SequenceElement], batch_size: int
|
||||
) -> Generator[List[SequenceElement], None, None]:
|
||||
batch_size = max(batch_size, 1)
|
||||
current_batch = []
|
||||
for element in sequence:
|
||||
if len(current_batch) == batch_size:
|
||||
yield current_batch
|
||||
current_batch = []
|
||||
current_batch.append(element)
|
||||
if len(current_batch) > 0:
|
||||
yield current_batch
|
||||
|
|
@ -5,7 +5,7 @@ import matplotlib.pyplot as plt
|
|||
from PIL import Image
|
||||
|
||||
from supervision.annotators.base import ImageType
|
||||
from supervision.annotators.utils import pillow_to_cv2
|
||||
from supervision.utils.image import pillow_to_cv2
|
||||
|
||||
|
||||
def plot_image(
|
||||
|
|
|
|||
Loading…
Reference in New Issue