xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #1028 from roboflow/feature/add_tiles_generator 

Add util to create image tiles
This commit is contained in:
Piotr Skalski 2024-04-05 18:37:19 +02:00 committed by GitHub
parent 299c933c4c 0a919db9fd
commit 372e703a7a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
32 changed files with 1312 additions and 69 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.pre-commit-config.yaml
View File

@ -27,7 +27,6 @@ repos:
- id: mixed-line-ending
- repo: https://github.com/PyCQA/bandit
rev: '1.7.8'
hooks:

6
docs/utils/image.md
View File

@ -17,6 +17,12 @@ status: new
:::supervision.utils.image.crop_image
<div class="md-typeset">
<h2>letterbox_image</h2>
</div>
:::supervision.utils.image.letterbox_image
<div class="md-typeset">
<h2>resize_image</h2>
</div>

18
docs/utils/iterables.md Normal file
View File

@ -0,0 +1,18 @@
---
comments: true
status: new
---
# Iterables Utils
<div class="md-typeset">
<h2>create_batches</h2>
</div>
:::supervision.utils.iterables.create_batches
<div class="md-typeset">
<h2>fill</h2>
</div>
:::supervision.utils.iterables.fill

1
mkdocs.yml
View File

@ -66,6 +66,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

10
supervision/__init__.py
View File

@ -72,7 +72,15 @@ from supervision.geometry.utils import get_polygon_center
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,

46
supervision/annotators/core.py
View File

@ -5,18 +5,14 @@ import cv2
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 @@ class BoundingBoxAnnotator(BaseAnnotator):
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 @@ class OrientedBoxAnnotator(BaseAnnotator):
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 @@ class MaskAnnotator(BaseAnnotator):
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 @@ class PolygonAnnotator(BaseAnnotator):
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 @@ class ColorAnnotator(BaseAnnotator):
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 @@ class HaloAnnotator(BaseAnnotator):
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 @@ class EllipseAnnotator(BaseAnnotator):
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 @@ class BoxCornerAnnotator(BaseAnnotator):
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 @@ class CircleAnnotator(BaseAnnotator):
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 @@ class DotAnnotator(BaseAnnotator):
self.position: Position = position
self.color_lookup: ColorLookup = color_lookup
@scene_to_annotator_img_type
@convert_for_annotation_method
def annotate(
self,
scene: ImageType,
@ -993,7 +989,7 @@ class LabelAnnotator:
center_y + text_h // 2,
)
@scene_to_annotator_img_type
@convert_for_annotation_method
def annotate(
self,
scene: ImageType,
@ -1164,7 +1160,7 @@ class BlurAnnotator(BaseAnnotator):
"""
self.kernel_size: int = kernel_size
@scene_to_annotator_img_type
@convert_for_annotation_method
def annotate(
self,
scene: ImageType,
@ -1249,7 +1245,7 @@ class TraceAnnotator:
self.thickness = thickness
self.color_lookup: ColorLookup = color_lookup
@scene_to_annotator_img_type
@convert_for_annotation_method
def annotate(
self,
scene: ImageType,
@ -1356,7 +1352,7 @@ class HeatMapAnnotator:
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.
@ -1432,7 +1428,7 @@ class PixelateAnnotator(BaseAnnotator):
"""
self.pixel_size: int = pixel_size
@scene_to_annotator_img_type
@convert_for_annotation_method
def annotate(
self,
scene: ImageType,
@ -1520,7 +1516,7 @@ class TriangleAnnotator(BaseAnnotator):
self.position: Position = position
self.color_lookup: ColorLookup = color_lookup
@scene_to_annotator_img_type
@convert_for_annotation_method
def annotate(
self,
scene: ImageType,
@ -1616,7 +1612,7 @@ class RoundBoxAnnotator(BaseAnnotator):
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,
@ -1753,7 +1749,7 @@ class PercentageBarAnnotator(BaseAnnotator):
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,
@ -1926,7 +1922,7 @@ class CropAnnotator(BaseAnnotator):
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,

34
supervision/annotators/utils.py
View File

@ -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 @@ class Trace:
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

4
supervision/detection/annotate.py
View File

@ -3,9 +3,9 @@ from typing import List, Optional, Union
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 @@ class BoxAnnotator:
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,

56
supervision/draw/color.py
View File

@ -104,10 +104,13 @@ class 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 @@ class 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 @@ class 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.

103
supervision/utils/conversion.py Normal file
View File

@ -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)

469
supervision/utils/image.py
View File

@ -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.
@ -168,3 +188,446 @@ 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.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

70
supervision/utils/iterables.py Normal file
View File

@ -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

2
supervision/utils/notebook.py
View File

@ -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.conversion import pillow_to_cv2
def plot_image(

BIN
test/utils/assets/1.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 67 KiB

BIN
test/utils/assets/2.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 37 KiB

BIN
test/utils/assets/3.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 46 KiB

BIN
test/utils/assets/4.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 40 KiB

BIN
test/utils/assets/5.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 45 KiB

BIN
test/utils/assets/all_images_tile.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 329 KiB

BIN
test/utils/assets/all_images_tile_and_custom_colors.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 328 KiB

BIN
test/utils/assets/all_images_tile_and_custom_colors_and_titles.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 332 KiB

BIN
test/utils/assets/all_images_tile_and_custom_grid.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 334 KiB

BIN
test/utils/assets/all_images_tile_and_titles_with_custom_configs.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 323 KiB

BIN
test/utils/assets/four_images_tile.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 266 KiB

BIN
test/utils/assets/single_image_tile.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 79 KiB

BIN
test/utils/assets/single_image_tile_enforced_grid.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 85 KiB

BIN
test/utils/assets/three_images_tile.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 213 KiB

BIN
test/utils/assets/two_images_tile.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 141 KiB

99
test/utils/conftest.py Normal file
View File

@ -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")
)

173
test/utils/test_conversion.py Normal file
View File

@ -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"

246
test/utils/test_image.py Normal file
View File

@ -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))

43
test/utils/test_iterables.py Normal file
View File

@ -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