diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0d53b4f30..7569e6b4b 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/docs/utils/image.md b/docs/utils/image.md index 8a6768d74..087aebd73 100644 --- a/docs/utils/image.md +++ b/docs/utils/image.md @@ -17,6 +17,12 @@ status: new :::supervision.utils.image.crop_image +
+

letterbox_image

+
+ +:::supervision.utils.image.letterbox_image +

resize_image

diff --git a/docs/utils/iterables.md b/docs/utils/iterables.md new file mode 100644 index 000000000..3ccf98e59 --- /dev/null +++ b/docs/utils/iterables.md @@ -0,0 +1,18 @@ +--- +comments: true +status: new +--- + +# Iterables Utils + +
+

create_batches

+
+ +:::supervision.utils.iterables.create_batches + +
+

fill

+
+ +:::supervision.utils.iterables.fill diff --git a/mkdocs.yml b/mkdocs.yml index 5444290b9..2afadfa55 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -65,6 +65,7 @@ nav: - Utils: - Video: utils/video.md - Image: utils/image.md + - Iterables: utils/iterables.md - Notebook: utils/notebook.md - File: utils/file.md - Assets: assets.md diff --git a/supervision/__init__.py b/supervision/__init__.py index 35f525b04..8c3acc265 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -72,7 +72,15 @@ from supervision.metrics.detection import ConfusionMatrix, MeanAveragePrecision from supervision.tracker.byte_tracker.core import ByteTrack from supervision.utils.file import list_files_with_extensions -from supervision.utils.image import ImageSink, crop_image, place_image, resize_image +from supervision.utils.image import ( + ImageSink, + create_tiles, + crop_image, + letterbox_image, + place_image, + resize_image, + resize_image_keeping_aspect_ratio, +) from supervision.utils.notebook import plot_image, plot_images_grid from supervision.utils.video import ( FPSMonitor, diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 2be2562ee..bdf7e9b2f 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -5,18 +5,14 @@ import numpy as np from supervision.annotators.base import BaseAnnotator, ImageType -from supervision.annotators.utils import ( - ColorLookup, - Trace, - resolve_color, - scene_to_annotator_img_type, -) +from supervision.annotators.utils import ColorLookup, Trace, resolve_color from supervision.config import CLASS_NAME_DATA_FIELD, ORIENTED_BOX_COORDINATES from supervision.detection.core import Detections from supervision.detection.utils import clip_boxes, mask_to_polygons from supervision.draw.color import Color, ColorPalette from supervision.draw.utils import draw_polygon from supervision.geometry.core import Position +from supervision.utils.conversion import convert_for_annotation_method from supervision.utils.image import crop_image, place_image, resize_image @@ -43,7 +39,7 @@ def __init__( self.thickness: int = thickness self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -124,7 +120,7 @@ def __init__( self.thickness: int = thickness self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -212,7 +208,7 @@ def __init__( self.opacity = opacity self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -299,7 +295,7 @@ def __init__( self.thickness: int = thickness self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -385,7 +381,7 @@ def __init__( self.color_lookup: ColorLookup = color_lookup self.opacity = opacity - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -479,7 +475,7 @@ def __init__( self.color_lookup: ColorLookup = color_lookup self.kernel_size: int = kernel_size - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -577,7 +573,7 @@ def __init__( self.end_angle: int = end_angle self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -668,7 +664,7 @@ def __init__( self.corner_length: int = corner_length self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -756,7 +752,7 @@ def __init__( self.thickness: int = thickness self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -846,7 +842,7 @@ def __init__( self.position: Position = position self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -989,7 +985,7 @@ def resolve_text_background_xyxy( center_y + text_h // 2, ) - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -1112,7 +1108,7 @@ def __init__(self, kernel_size: int = 15): """ self.kernel_size: int = kernel_size - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -1197,7 +1193,7 @@ def __init__( self.thickness = thickness self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -1304,7 +1300,7 @@ def __init__( self.top_hue = top_hue self.low_hue = low_hue - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate(self, scene: ImageType, detections: Detections) -> ImageType: """ Annotates the scene with a heatmap based on the provided detections. @@ -1380,7 +1376,7 @@ def __init__(self, pixel_size: int = 20): """ self.pixel_size: int = pixel_size - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -1468,7 +1464,7 @@ def __init__( self.position: Position = position self.color_lookup: ColorLookup = color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -1564,7 +1560,7 @@ def __init__( raise ValueError("roundness attribute must be float between (0, 1.0]") self.roundness: float = roundness - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -1701,7 +1697,7 @@ def __init__( if border_thickness is None: self.border_thickness = int(0.15 * self.height) - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, @@ -1874,7 +1870,7 @@ def __init__( self.border_thickness: int = border_thickness self.border_color_lookup: ColorLookup = border_color_lookup - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, diff --git a/supervision/annotators/utils.py b/supervision/annotators/utils.py index 6f9cd9db2..e206c8cbb 100644 --- a/supervision/annotators/utils.py +++ b/supervision/annotators/utils.py @@ -1,12 +1,8 @@ from enum import Enum -from functools import wraps from typing import Optional, Union -import cv2 import numpy as np -from PIL import Image -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 @@ -123,33 +119,3 @@ def put(self, detections: Detections) -> None: def get(self, tracker_id: int) -> np.ndarray: 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 - an image type used internally by the annotators, converts back when annotation - is complete. - """ - - @wraps(annotate_func) - def wrapper(self, scene: ImageType, *args, **kwargs): - if isinstance(scene, np.ndarray): - return annotate_func(self, scene, *args, **kwargs) - - 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 - - raise ValueError(f"Unsupported image type: {type(scene)}") - - return wrapper diff --git a/supervision/detection/annotate.py b/supervision/detection/annotate.py index a8df0a573..f32c73440 100644 --- a/supervision/detection/annotate.py +++ b/supervision/detection/annotate.py @@ -3,9 +3,9 @@ import cv2 from supervision.annotators.base import ImageType -from supervision.annotators.utils import scene_to_annotator_img_type from supervision.detection.core import Detections from supervision.draw.color import Color, ColorPalette +from supervision.utils.conversion import convert_for_annotation_method from supervision.utils.internal import deprecated @@ -46,7 +46,7 @@ def __init__( self.text_thickness: int = text_thickness self.text_padding: int = text_padding - @scene_to_annotator_img_type + @convert_for_annotation_method def annotate( self, scene: ImageType, diff --git a/supervision/draw/color.py b/supervision/draw/color.py index 149908d8c..debb46f3b 100644 --- a/supervision/draw/color.py +++ b/supervision/draw/color.py @@ -104,10 +104,13 @@ def from_hex(cls, color_hex: str) -> Color: Create a Color instance from a hex string. Args: - color_hex (str): Hex string of the color. + color_hex (str): The hex string representing the color. This string can + start with '#' followed by either 3 or 6 hexadecimal characters. In + case of 3 characters, each character is repeated to form the full + 6-character hex code. Returns: - Color: Instance representing the color. + Color: An instance representing the color. Example: ```python @@ -115,6 +118,9 @@ def from_hex(cls, color_hex: str) -> Color: sv.Color.from_hex('#ff00ff') # Color(r=255, g=0, b=255) + + sv.Color.from_hex('#f0f') + # Color(r=255, g=0, b=255) ``` """ _validate_color_hex(color_hex) @@ -124,6 +130,52 @@ def from_hex(cls, color_hex: str) -> Color: r, g, b = (int(color_hex[i : i + 2], 16) for i in range(0, 6, 2)) return cls(r, g, b) + @classmethod + def from_rgb_tuple(cls, color_tuple: Tuple[int, int, int]) -> Color: + """ + Create a Color instance from an RGB tuple. + + Args: + color_tuple (Tuple[int, int, int]): A tuple representing the color in RGB + format, where each element is an integer in the range 0-255. + + Returns: + Color: An instance representing the color. + + Example: + ```python + import supervision as sv + + sv.Color.from_rgb_tuple((255, 255, 0)) + # Color(r=255, g=255, b=0) + ``` + """ + r, g, b = color_tuple + return cls(r=r, g=g, b=b) + + @classmethod + def from_bgr_tuple(cls, color_tuple: Tuple[int, int, int]) -> Color: + """ + Create a Color instance from a BGR tuple. + + Args: + color_tuple (Tuple[int, int, int]): A tuple representing the color in BGR + format, where each element is an integer in the range 0-255. + + Returns: + Color: An instance representing the color. + + Example: + ```python + import supervision as sv + + sv.Color.from_bgr_tuple((0, 255, 255)) + # Color(r=255, g=255, b=0) + ``` + """ + b, g, r = color_tuple + return cls(r=r, g=g, b=b) + def as_hex(self) -> str: """ Converts the Color instance to a hex string. diff --git a/supervision/utils/conversion.py b/supervision/utils/conversion.py new file mode 100644 index 000000000..608104bcd --- /dev/null +++ b/supervision/utils/conversion.py @@ -0,0 +1,103 @@ +from functools import wraps +from typing import List + +import cv2 +import numpy as np +from PIL import Image + +from supervision.annotators.base import ImageType + + +def convert_for_annotation_method(annotate_func): + """ + Decorates `BaseAnnotator.annotate` implementations, converts scene to + an image type used internally by the annotators, converts back when annotation + is complete. + """ + + @wraps(annotate_func) + def wrapper(self, scene: ImageType, *args, **kwargs): + if isinstance(scene, np.ndarray): + return annotate_func(self, scene, *args, **kwargs) + + if isinstance(scene, Image.Image): + scene = pillow_to_cv2(scene) + annotated = annotate_func(self, scene, *args, **kwargs) + return cv2_to_pillow(image=annotated) + + raise ValueError(f"Unsupported image type: {type(scene)}") + + return wrapper + + +def convert_for_image_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 + + +def images_to_cv2(images: List[ImageType]) -> List[np.ndarray]: + """ + Converts images provided either as Pillow images or OpenCV + images into OpenCV format. + + Args: + images (List[ImageType]): Images to be converted + + Returns: + List[np.ndarray]: List of input images in OpenCV format + (with order preserved). + + """ + 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: + """ + Converts Pillow image into OpenCV image, handling RGB -> BGR + conversion. + + Args: + image (Image.Image): Pillow image (in RGB format). + + Returns: + np.ndarray: Input image converted to OpenCV format. + """ + scene = np.array(image) + scene = cv2.cvtColor(scene, cv2.COLOR_RGB2BGR) + return scene + + +def cv2_to_pillow(image: np.ndarray) -> Image.Image: + """ + Converts OpenCV image into Pillow image, handling BGR -> RGB + conversion. + + Args: + image (np.ndarray): OpenCV image (in BGR format). + + Returns: + Image.Image: Input image converted to Pillow format. + """ + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + return Image.fromarray(image) diff --git a/supervision/utils/image.py b/supervision/utils/image.py index 4b9372ee2..d79da8d6f 100644 --- a/supervision/utils/image.py +++ b/supervision/utils/image.py @@ -1,11 +1,30 @@ +import itertools +import math import os import shutil -from typing import Optional, Tuple +from functools import partial +from typing import Callable, List, Literal, Optional, Tuple, Union import cv2 import numpy as np +from supervision.annotators.base import ImageType +from supervision.draw.color import Color +from supervision.draw.utils import calculate_optimal_text_scale, draw_text +from supervision.geometry.core import Point +from supervision.utils.conversion import ( + convert_for_image_processing, + cv2_to_pillow, + images_to_cv2, +) +from supervision.utils.iterables import create_batches, fill +RelativePosition = Literal["top", "bottom"] + +MAX_COLUMNS_FOR_SINGLE_ROW_GRID = 3 + + +@convert_for_image_processing def crop_image(image: np.ndarray, xyxy: np.ndarray) -> np.ndarray: """ Crops the given image based on the given bounding box. @@ -35,14 +54,15 @@ def crop_image(image: np.ndarray, xyxy: np.ndarray) -> np.ndarray: return image[y1:y2, x1:x2] +@convert_for_image_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. Args: image (np.ndarray): The input image to be resized. - scale_factor (float): The factor by which the image will be scaled. Scale factor - > 1.0 zooms in, < 1.0 zooms out. + scale_factor (float): The factor by which the image will be scaled. Scale + factor > 1.0 zooms in, < 1.0 zooms out. Returns: np.ndarray: The resized image. @@ -167,3 +187,446 @@ def save_image(self, image: np.ndarray, image_name: Optional[str] = None): 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.from_hex("#D9D9D9"), + tile_margin: int = 10, + tile_margin_color: Union[Tuple[int, int, int], Color] = Color.from_hex("#BFBEBD"), + return_type: Literal["auto", "cv2", "pillow"] = "auto", + titles: Optional[List[Optional[str]]] = None, + titles_anchors: Optional[Union[Point, List[Optional[Point]]]] = None, + titles_color: Union[Tuple[int, int, int], Color] = Color.from_hex("#262523"), + titles_scale: Optional[float] = None, + titles_thickness: int = 1, + titles_padding: int = 10, + titles_text_font: int = cv2.FONT_HERSHEY_SIMPLEX, + titles_background_color: Union[Tuple[int, int, int], Color] = Color.from_hex( + "#D9D9D9" + ), + default_title_placement: RelativePosition = "top", +) -> ImageType: + """ + 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 = _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}." + ) + 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 = _color_to_bgr(color=titles_color) + titles_background_color = _color_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: 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] + 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: Optional[List[Optional[str]]], + titles_anchors: List[Optional[Point]], + titles_color: Tuple[int, int, int], + titles_scale: Optional[float], + 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: Optional[List[Optional[str]]], + titles_anchors: List[Optional[Point]], + titles_color: Tuple[int, int, int], + titles_scale: Optional[float], + 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[Optional[Point]], + 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 + + +@convert_for_image_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. + + Args: + 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, + ) + + +@convert_for_image_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. + + For example: input image is (640, 480) and we want to resize into + (1024, 1024). If this rectangular image is just resized naively + to square-shape output - aspect ratio would be altered. If we do not + want this to happen - we may resize bigger dimension (640) to 1024. + Ratio of change is 1.6. This ratio is later on used to calculate scaling + in the other dimension. As a result we have (1024, 768) image. + + 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) + """ + if image.shape[:2] == desired_size[::-1]: + return 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 diff --git a/supervision/utils/iterables.py b/supervision/utils/iterables.py new file mode 100644 index 000000000..ad570379c --- /dev/null +++ b/supervision/utils/iterables.py @@ -0,0 +1,70 @@ +from typing import Generator, Iterable, List, TypeVar + +V = TypeVar("V") + + +def create_batches( + sequence: Iterable[V], batch_size: int +) -> Generator[List[V], None, None]: + """ + Provides a generator that yields chunks of the input sequence + of the size specified by the `batch_size` parameter. The last + chunk may be a smaller batch. + + Args: + sequence (Iterable[V]): The sequence to be split into batches. + batch_size (int): The expected size of a batch. + + Returns: + Generator[List[V], None, None]: A generator that yields chunks + of `sequence` of size `batch_size`, up to the length of + the input `sequence`. + + Examples: + ```python + list(create_batches([1, 2, 3, 4, 5], 2)) + # [[1, 2], [3, 4], [5]] + + list(create_batches("abcde", 3)) + # [['a', 'b', 'c'], ['d', 'e']] + ``` + """ + 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 current_batch: + yield current_batch + + +def fill(sequence: List[V], desired_size: int, content: V) -> List[V]: + """ + Fill the sequence with padding elements until the sequence reaches + the desired size. + + Args: + sequence (List[V]): The input sequence. + desired_size (int): The expected size of the output list. The + difference between this value and the actual length of `sequence` + (if positive) dictates how many elements will be added as padding. + content (V): The element to be placed at the end of the input + `sequence` as padding. + + Returns: + List[V]: A padded version of the input `sequence` (if needed). + + Examples: + ```python + fill([1, 2], 4, 0) + # [1, 2, 0, 0] + + fill(['a', 'b'], 3, 'c') + # ['a', 'b', 'c'] + ``` + """ + missing_size = max(0, desired_size - len(sequence)) + sequence.extend([content] * missing_size) + return sequence diff --git a/supervision/utils/notebook.py b/supervision/utils/notebook.py index 159f3c17c..19f5eaed5 100644 --- a/supervision/utils/notebook.py +++ b/supervision/utils/notebook.py @@ -5,7 +5,7 @@ from PIL import Image from supervision.annotators.base import ImageType -from supervision.annotators.utils import pillow_to_cv2 +from supervision.utils.conversion import pillow_to_cv2 def plot_image( diff --git a/test/utils/assets/1.jpg b/test/utils/assets/1.jpg new file mode 100644 index 000000000..ed88f941f Binary files /dev/null and b/test/utils/assets/1.jpg differ diff --git a/test/utils/assets/2.jpg b/test/utils/assets/2.jpg new file mode 100644 index 000000000..cbe01e983 Binary files /dev/null and b/test/utils/assets/2.jpg differ diff --git a/test/utils/assets/3.jpg b/test/utils/assets/3.jpg new file mode 100644 index 000000000..8fc100b97 Binary files /dev/null and b/test/utils/assets/3.jpg differ diff --git a/test/utils/assets/4.jpg b/test/utils/assets/4.jpg new file mode 100644 index 000000000..16467d977 Binary files /dev/null and b/test/utils/assets/4.jpg differ diff --git a/test/utils/assets/5.jpg b/test/utils/assets/5.jpg new file mode 100644 index 000000000..e58fd9e1a Binary files /dev/null and b/test/utils/assets/5.jpg differ diff --git a/test/utils/assets/all_images_tile.png b/test/utils/assets/all_images_tile.png new file mode 100644 index 000000000..ef5066a73 Binary files /dev/null and b/test/utils/assets/all_images_tile.png differ diff --git a/test/utils/assets/all_images_tile_and_custom_colors.png b/test/utils/assets/all_images_tile_and_custom_colors.png new file mode 100644 index 000000000..db8ed13bf Binary files /dev/null and b/test/utils/assets/all_images_tile_and_custom_colors.png differ diff --git a/test/utils/assets/all_images_tile_and_custom_colors_and_titles.png b/test/utils/assets/all_images_tile_and_custom_colors_and_titles.png new file mode 100644 index 000000000..8a058f979 Binary files /dev/null and b/test/utils/assets/all_images_tile_and_custom_colors_and_titles.png differ diff --git a/test/utils/assets/all_images_tile_and_custom_grid.png b/test/utils/assets/all_images_tile_and_custom_grid.png new file mode 100644 index 000000000..1407a97a0 Binary files /dev/null and b/test/utils/assets/all_images_tile_and_custom_grid.png differ diff --git a/test/utils/assets/all_images_tile_and_titles_with_custom_configs.png b/test/utils/assets/all_images_tile_and_titles_with_custom_configs.png new file mode 100644 index 000000000..0da1d2b80 Binary files /dev/null and b/test/utils/assets/all_images_tile_and_titles_with_custom_configs.png differ diff --git a/test/utils/assets/four_images_tile.png b/test/utils/assets/four_images_tile.png new file mode 100644 index 000000000..220df4465 Binary files /dev/null and b/test/utils/assets/four_images_tile.png differ diff --git a/test/utils/assets/single_image_tile.png b/test/utils/assets/single_image_tile.png new file mode 100644 index 000000000..434deb002 Binary files /dev/null and b/test/utils/assets/single_image_tile.png differ diff --git a/test/utils/assets/single_image_tile_enforced_grid.png b/test/utils/assets/single_image_tile_enforced_grid.png new file mode 100644 index 000000000..0f8b5ce40 Binary files /dev/null and b/test/utils/assets/single_image_tile_enforced_grid.png differ diff --git a/test/utils/assets/three_images_tile.png b/test/utils/assets/three_images_tile.png new file mode 100644 index 000000000..104bba8f9 Binary files /dev/null and b/test/utils/assets/three_images_tile.png differ diff --git a/test/utils/assets/two_images_tile.png b/test/utils/assets/two_images_tile.png new file mode 100644 index 000000000..9922c68e7 Binary files /dev/null and b/test/utils/assets/two_images_tile.png differ diff --git a/test/utils/conftest.py b/test/utils/conftest.py new file mode 100644 index 000000000..cc134f8ce --- /dev/null +++ b/test/utils/conftest.py @@ -0,0 +1,99 @@ +import os +from typing import List + +import cv2 +import numpy as np +from _pytest.fixtures import fixture +from PIL import Image + +ASSETS_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "assets")) +ALL_IMAGES_LIST = [os.path.join(ASSETS_DIR, f"{i}.jpg") for i in range(1, 6)] + + +@fixture(scope="function") +def empty_opencv_image() -> np.ndarray: + return np.zeros((128, 128, 3), dtype=np.uint8) + + +@fixture(scope="function") +def empty_pillow_image() -> Image.Image: + return Image.new(mode="RGB", size=(128, 128), color=(0, 0, 0)) + + +@fixture(scope="function") +def all_images() -> List[np.ndarray]: + return [cv2.imread(path) for path in ALL_IMAGES_LIST] + + +@fixture(scope="function") +def one_image() -> np.ndarray: + return cv2.imread(ALL_IMAGES_LIST[0]) + + +@fixture(scope="function") +def two_images() -> List[np.ndarray]: + return [cv2.imread(path) for path in ALL_IMAGES_LIST[:2]] + + +@fixture(scope="function") +def three_images() -> List[np.ndarray]: + return [cv2.imread(path) for path in ALL_IMAGES_LIST[:3]] + + +@fixture(scope="function") +def four_images() -> List[np.ndarray]: + return [cv2.imread(path) for path in ALL_IMAGES_LIST[:4]] + + +@fixture(scope="function") +def all_images_tile() -> np.ndarray: + return cv2.imread(os.path.join(ASSETS_DIR, "all_images_tile.png")) + + +@fixture(scope="function") +def all_images_tile_and_custom_colors() -> np.ndarray: + return cv2.imread(os.path.join(ASSETS_DIR, "all_images_tile_and_custom_colors.png")) + + +@fixture(scope="function") +def all_images_tile_and_custom_grid() -> np.ndarray: + return cv2.imread(os.path.join(ASSETS_DIR, "all_images_tile_and_custom_grid.png")) + + +@fixture(scope="function") +def four_images_tile() -> np.ndarray: + return cv2.imread(os.path.join(ASSETS_DIR, "four_images_tile.png")) + + +@fixture(scope="function") +def single_image_tile() -> np.ndarray: + return cv2.imread(os.path.join(ASSETS_DIR, "single_image_tile.png")) + + +@fixture(scope="function") +def single_image_tile_enforced_grid() -> np.ndarray: + return cv2.imread(os.path.join(ASSETS_DIR, "single_image_tile_enforced_grid.png")) + + +@fixture(scope="function") +def three_images_tile() -> np.ndarray: + return cv2.imread(os.path.join(ASSETS_DIR, "three_images_tile.png")) + + +@fixture(scope="function") +def two_images_tile() -> np.ndarray: + return cv2.imread(os.path.join(ASSETS_DIR, "two_images_tile.png")) + + +@fixture(scope="function") +def all_images_tile_and_custom_colors_and_titles() -> np.ndarray: + return cv2.imread( + os.path.join(ASSETS_DIR, "all_images_tile_and_custom_colors_and_titles.png") + ) + + +@fixture(scope="function") +def all_images_tile_and_titles_with_custom_configs() -> np.ndarray: + return cv2.imread( + os.path.join(ASSETS_DIR, "all_images_tile_and_titles_with_custom_configs.png") + ) diff --git a/test/utils/test_conversion.py b/test/utils/test_conversion.py new file mode 100644 index 000000000..fb3d8faff --- /dev/null +++ b/test/utils/test_conversion.py @@ -0,0 +1,173 @@ +import numpy as np +from PIL import Image, ImageChops + +from supervision.utils.conversion import ( + convert_for_image_processing, + cv2_to_pillow, + images_to_cv2, + pillow_to_cv2, +) + + +def test_convert_for_image_processing_when_pillow_image_submitted( + empty_opencv_image: np.ndarray, empty_pillow_image: Image.Image +) -> None: + # given + param_a_value = 3 + param_b_value = "some" + + @convert_for_image_processing + def my_custom_processing_function( + image: np.ndarray, + param_a: int, + param_b: str, + ) -> np.ndarray: + assert np.allclose( + image, empty_opencv_image + ), "Expected conversion to OpenCV image to happen" + assert ( + param_a == param_a_value + ), f"Parameter a expected to be {param_a_value} in target function" + assert ( + param_b == param_b_value + ), f"Parameter b expected to be {param_b_value} in target function" + return image + + # when + result = my_custom_processing_function( + empty_pillow_image, + param_a_value, + param_b=param_b_value, + ) + + # then + difference = ImageChops.difference(result, empty_pillow_image) + assert difference.getbbox() is None, ( + "Wrapper is expected to convert-back the OpenCV image " + "into Pillow format without changes to content" + ) + + +def test_convert_for_image_processing_when_opencv_image_submitted( + empty_opencv_image: np.ndarray, +) -> None: + # given + param_a_value = 3 + param_b_value = "some" + + @convert_for_image_processing + def my_custom_processing_function( + image: np.ndarray, + param_a: int, + param_b: str, + ) -> np.ndarray: + assert np.allclose( + image, empty_opencv_image + ), "Expected conversion to OpenCV image to happen" + assert ( + param_a == param_a_value + ), f"Parameter a expected to be {param_a_value} in target function" + assert ( + param_b == param_b_value + ), f"Parameter b expected to be {param_b_value} in target function" + return image + + # when + result = my_custom_processing_function( + empty_opencv_image, + param_a_value, + param_b=param_b_value, + ) + + # then + assert ( + result is empty_opencv_image + ), "Expected to return OpenCV image without changes" + + +def test_cv2_to_pillow( + empty_opencv_image: np.ndarray, empty_pillow_image: Image.Image +) -> None: + # when + result = cv2_to_pillow(image=empty_opencv_image) + + # then + difference = ImageChops.difference(result, empty_pillow_image) + assert ( + difference.getbbox() is None + ), "Conversion to PIL.Image expected not to change the content of image" + + +def test_pillow_to_cv2( + empty_opencv_image: np.ndarray, empty_pillow_image: Image.Image +) -> None: + # when + result = pillow_to_cv2(image=empty_pillow_image) + + # then + assert np.allclose( + result, empty_opencv_image + ), "Conversion to OpenCV image expected not to change the content of image" + + +def test_images_to_cv2_when_empty_input_provided() -> None: + # when + result = images_to_cv2(images=[]) + + # then + assert result == [], "Expected empty output when empty input provided" + + +def test_images_to_cv2_when_only_cv2_images_provided( + empty_opencv_image: np.ndarray, +) -> None: + # given + images = [empty_opencv_image] * 5 + + # when + result = images_to_cv2(images=images) + + # then + assert len(result) == 5, "Expected the same number of output element as input ones" + for result_element in result: + assert ( + result_element is empty_opencv_image + ), "Expected CV images not to be touched by conversion" + + +def test_images_to_cv2_when_only_pillow_images_provided( + empty_pillow_image: Image.Image, + empty_opencv_image: np.ndarray, +) -> None: + # given + images = [empty_pillow_image] * 5 + + # when + result = images_to_cv2(images=images) + + # then + assert len(result) == 5, "Expected the same number of output element as input ones" + for result_element in result: + assert np.allclose( + result_element, empty_opencv_image + ), "Output images expected to be equal to empty OpenCV image" + + +def test_images_to_cv2_when_mixed_input_provided( + empty_pillow_image: Image.Image, + empty_opencv_image: np.ndarray, +) -> None: + # given + images = [empty_pillow_image, empty_opencv_image] + + # when + result = images_to_cv2(images=images) + + # then + assert len(result) == 2, "Expected the same number of output element as input ones" + assert np.allclose( + result[0], empty_opencv_image + ), "PIL image should be converted to OpenCV one, equal to example empty image" + assert ( + result[1] is empty_opencv_image + ), "Expected CV images not to be touched by conversion" diff --git a/test/utils/test_image.py b/test/utils/test_image.py new file mode 100644 index 000000000..e50f2e574 --- /dev/null +++ b/test/utils/test_image.py @@ -0,0 +1,246 @@ +from typing import List + +import numpy as np +import pytest +from PIL import Image, ImageChops + +from supervision import Color, Point +from supervision.utils.image import ( + create_tiles, + letterbox_image, + resize_image_keeping_aspect_ratio, +) + + +def test_resize_image_keeping_aspect_ratio_for_opencv_image() -> None: + # given + image = np.zeros((480, 640, 3), dtype=np.uint8) + expected_result = np.zeros((768, 1024, 3), dtype=np.uint8) + + # when + result = resize_image_keeping_aspect_ratio( + image=image, + desired_size=(1024, 1024), + ) + + # then + assert np.allclose( + result, expected_result + ), "Expected output shape to be (w, h): (1024, 768)" + + +def test_resize_image_keeping_aspect_ratio_for_pillow_image() -> None: + # given + image = Image.new(mode="RGB", size=(640, 480), color=(0, 0, 0)) + expected_result = Image.new(mode="RGB", size=(1024, 768), color=(0, 0, 0)) + + # when + result = resize_image_keeping_aspect_ratio( + image=image, + desired_size=(1024, 1024), + ) + + # then + assert result.size == (1024, 768), "Expected output shape to be (w, h): (1024, 768)" + difference = ImageChops.difference(result, expected_result) + assert ( + difference.getbbox() is None + ), "Expected no difference in resized image content as the image is all zeros" + + +def test_letterbox_image_for_opencv_image() -> None: + # given + image = np.zeros((480, 640, 3), dtype=np.uint8) + expected_result = np.concatenate( + [ + np.ones((128, 1024, 3), dtype=np.uint8) * 255, + np.zeros((768, 1024, 3), dtype=np.uint8), + np.ones((128, 1024, 3), dtype=np.uint8) * 255, + ], + axis=0, + ) + + # when + result = letterbox_image( + image=image, desired_size=(1024, 1024), color=(255, 255, 255) + ) + + # then + assert np.allclose(result, expected_result), ( + "Expected output shape to be (w, h): " + "(1024, 1024) with padding added top and bottom" + ) + + +def test_letterbox_image_for_pillow_image() -> None: + # given + image = Image.new(mode="RGB", size=(640, 480), color=(0, 0, 0)) + expected_result = Image.fromarray( + np.concatenate( + [ + np.ones((128, 1024, 3), dtype=np.uint8) * 255, + np.zeros((768, 1024, 3), dtype=np.uint8), + np.ones((128, 1024, 3), dtype=np.uint8) * 255, + ], + axis=0, + ) + ) + + # when + result = letterbox_image( + image=image, desired_size=(1024, 1024), color=(255, 255, 255) + ) + + # then + assert result.size == ( + 1024, + 1024, + ), "Expected output shape to be (w, h): (1024, 1024)" + difference = ImageChops.difference(result, expected_result) + assert ( + difference.getbbox() is None + ), "Expected padding to be added top and bottom with padding added top and bottom" + + +def test_create_tiles_with_one_image( + one_image: np.ndarray, single_image_tile: np.ndarray +) -> None: + # when + result = create_tiles(images=[one_image], single_tile_size=(240, 240)) + + # # then + assert np.allclose(result, single_image_tile, atol=5.0) + + +def test_create_tiles_with_one_image_and_enforced_grid( + one_image: np.ndarray, single_image_tile_enforced_grid: np.ndarray +) -> None: + # when + result = create_tiles( + images=[one_image], + grid_size=(None, 3), + single_tile_size=(240, 240), + ) + + # then + assert np.allclose(result, single_image_tile_enforced_grid, atol=5.0) + + +def test_create_tiles_with_two_images( + two_images: List[np.ndarray], two_images_tile: np.ndarray +) -> None: + # when + result = create_tiles(images=two_images, single_tile_size=(240, 240)) + + # then + assert np.allclose(result, two_images_tile, atol=5.0) + + +def test_create_tiles_with_three_images( + three_images: List[np.ndarray], three_images_tile: np.ndarray +) -> None: + # when + result = create_tiles(images=three_images, single_tile_size=(240, 240)) + + # then + assert np.allclose(result, three_images_tile, atol=5.0) + + +def test_create_tiles_with_four_images( + four_images: List[np.ndarray], + four_images_tile: np.ndarray, +) -> None: + # when + result = create_tiles(images=four_images, single_tile_size=(240, 240)) + + # then + assert np.allclose(result, four_images_tile, atol=5.0) + + +def test_create_tiles_with_all_images( + all_images: List[np.ndarray], + all_images_tile: np.ndarray, +) -> None: + # when + result = create_tiles(images=all_images, single_tile_size=(240, 240)) + + # then + assert np.allclose(result, all_images_tile, atol=5.0) + + +def test_create_tiles_with_all_images_and_custom_grid( + all_images: List[np.ndarray], all_images_tile_and_custom_grid: np.ndarray +) -> None: + # when + result = create_tiles( + images=all_images, + grid_size=(3, 3), + single_tile_size=(240, 240), + ) + + # then + assert np.allclose(result, all_images_tile_and_custom_grid, atol=5.0) + + +def test_create_tiles_with_all_images_and_custom_colors( + all_images: List[np.ndarray], all_images_tile_and_custom_colors: np.ndarray +) -> None: + # when + result = create_tiles( + images=all_images, + tile_margin_color=(127, 127, 127), + tile_padding_color=(224, 224, 224), + single_tile_size=(240, 240), + ) + + # then + assert np.allclose(result, all_images_tile_and_custom_colors, atol=5.0) + + +def test_create_tiles_with_all_images_and_titles( + all_images: List[np.ndarray], + all_images_tile_and_custom_colors_and_titles: np.ndarray, +) -> None: + # when + result = create_tiles( + images=all_images, + titles=["Image 1", None, "Image 3", "Image 4"], + single_tile_size=(240, 240), + ) + + # then + assert np.allclose(result, all_images_tile_and_custom_colors_and_titles, atol=5.0) + + +def test_create_tiles_with_all_images_and_titles_with_custom_configs( + all_images: List[np.ndarray], + all_images_tile_and_titles_with_custom_configs: np.ndarray, +) -> None: + # when + result = create_tiles( + images=all_images, + titles=["Image 1", None, "Image 3", "Image 4"], + single_tile_size=(240, 240), + titles_anchors=[ + Point(x=200, y=300), + Point(x=300, y=400), + None, + Point(x=300, y=400), + ], + titles_color=Color.RED, + titles_scale=1.5, + titles_thickness=3, + titles_padding=20, + titles_background_color=Color.BLACK, + default_title_placement="bottom", + ) + + # then + assert np.allclose(result, all_images_tile_and_titles_with_custom_configs, atol=5.0) + + +def test_create_tiles_with_all_images_and_custom_grid_to_small_to_fit_images( + all_images: List[np.ndarray], +) -> None: + with pytest.raises(ValueError): + _ = create_tiles(images=all_images, grid_size=(2, 2)) diff --git a/test/utils/test_iterables.py b/test/utils/test_iterables.py new file mode 100644 index 000000000..2d34605c0 --- /dev/null +++ b/test/utils/test_iterables.py @@ -0,0 +1,43 @@ +import pytest + +from supervision.utils.iterables import create_batches, fill + + +@pytest.mark.parametrize( + "sequence, batch_size, expected", + [ + # Empty sequence, non-zero batch size. Expect empty list. + ([], 4, []), + # Non-zero size sequence, batch size of 0. Each item is its own batch. + ([1, 2, 3], 0, [[1], [2], [3]]), + # Batch size larger than sequence. All items in a single batch. + ([1, 2], 4, [[1, 2]]), + # Batch size evenly divides the sequence. Equal size batches. + ([1, 2, 3, 4], 2, [[1, 2], [3, 4]]), + # Batch size doesn't evenly divide sequence. Last batch smaller. + ([1, 2, 3, 4], 3, [[1, 2, 3], [4]]), + ], +) +def test_create_batches(sequence, batch_size, expected) -> None: + result = list(create_batches(sequence=sequence, batch_size=batch_size)) + assert result == expected + + +@pytest.mark.parametrize( + "sequence, desired_size, content, expected", + [ + # Empty sequence, desired size 0. Expect empty list. + ([], 0, 1, []), + # Empty sequence, non-zero desired size. Filled with padding. + ([], 3, 1, [1, 1, 1]), + # Sequence at desired size. No changes. + ([2, 2, 2], 3, 1, [2, 2, 2]), + # Sequence exceeds desired size. No changes. + ([2, 2, 2, 2], 3, 1, [2, 2, 2, 2]), + # Non-empty sequence, shorter than desired. Padding added. + ([2], 3, 1, [2, 1, 1]), + ], +) +def test_fill(sequence, desired_size, content, expected) -> None: + result = fill(sequence=sequence, desired_size=desired_size, content=content) + assert result == expected