From 7eba3382432c1bc98a30b66b3d65ce72ebb603bd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20P=C4=99czek?= Date: Wed, 20 Mar 2024 23:08:39 +0100 Subject: [PATCH] Add first draft of implementation for function creating image tiles --- .pre-commit-config.yaml | 1 - supervision/annotators/utils.py | 12 +- supervision/utils/image.py | 346 +++++++++++++++++++++++++++++++- supervision/utils/iterables.py | 17 ++ supervision/utils/notebook.py | 2 +- 5 files changed, 365 insertions(+), 13 deletions(-) create mode 100644 supervision/utils/iterables.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0d53b4f3..7569e6b4 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -27,7 +27,6 @@ repos: - id: mixed-line-ending - - repo: https://github.com/PyCQA/bandit rev: '1.7.8' hooks: diff --git a/supervision/annotators/utils.py b/supervision/annotators/utils.py index 6f9cd9db..0a7586a4 100644 --- a/supervision/annotators/utils.py +++ b/supervision/annotators/utils.py @@ -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)}") diff --git a/supervision/utils/image.py b/supervision/utils/image.py index 4b9372ee..39787ca4 100644 --- a/supervision/utils/image.py +++ b/supervision/utils/image.py @@ -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) diff --git a/supervision/utils/iterables.py b/supervision/utils/iterables.py new file mode 100644 index 00000000..85684661 --- /dev/null +++ b/supervision/utils/iterables.py @@ -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 diff --git a/supervision/utils/notebook.py b/supervision/utils/notebook.py index 159f3c17..346064ac 100644 --- a/supervision/utils/notebook.py +++ b/supervision/utils/notebook.py @@ -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(