Add tests and docstrings for utils

supervision/utils/conversion.py

@@ -8,6 +8,18 @@ from supervision.annotators.base import ImageType
 
 
 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):
@@ -17,11 +29,31 @@ def images_to_cv2(images: List[ImageType]) -> List[np.ndarray]:
 
 
 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)

supervision/utils/image.py

@@ -317,14 +317,14 @@ def create_tiles(
             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), padding=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), padding=None
+        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)

supervision/utils/iterables.py

@@ -6,6 +6,21 @@ SequenceElement = TypeVar("SequenceElement")
 def create_batches(
     sequence: Iterable[SequenceElement], batch_size: int
 ) -> Generator[List[SequenceElement], None, None]:
+    """
+    Provides a generator that yields chunks of input sequence
+    of size specified by `batch_size` parameter. Last
+    chunk may be smaller batch.
+
+    Args:
+        sequence (Iterable[SequenceElement]): Sequence to be
+            split into batches.
+        batch_size (int): Expected size of a batch
+
+    Returns:
+        Generator[List[SequenceElement], None, None]: Generator
+            to yield chinks of `sequence` of size `batch_size`,
+            up to the length of input `sequence`.
+    """
     batch_size = max(batch_size, 1)
     current_batch = []
     for element in sequence:
@@ -20,9 +35,24 @@ def create_batches(
 def fill(
     sequence: List[SequenceElement],
     desired_size: int,
-    padding: SequenceElement,
+    content: SequenceElement,
 ) -> List[SequenceElement]:
+    """
+    Fill the sequence with padding elements until sequence reaches
+    desired size.
+
+    Args:
+        sequence (List[SequenceElement]): Input sequence.
+        desired_size (int): Expected size of output list - difference
+            between this value and actual `sequence` length (if positive)
+            dictates how many elements will be added as padding.
+        content (SequenceElement): Element to be placed at the end of
+            input `sequence` as padding.
+
+    Returns:
+        List[SequenceElement]: Padded version of input `sequence` (if needed)
+    """
     missing_size = max(0, desired_size - len(sequence))
-    required_padding = [padding] * missing_size
+    required_padding = [content] * missing_size
     sequence.extend(required_padding)
     return sequence

test/utils/conftest.py

@@ -0,0 +1,13 @@
+import numpy as np
+from _pytest.fixtures import fixture
+from PIL import Image
+
+
+@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))

test/utils/test_conversion.py

@@ -0,0 +1,92 @@
+import numpy as np
+from PIL import Image, ImageChops
+
+from supervision.utils.conversion import cv2_to_pillow, images_to_cv2, pillow_to_cv2
+
+
+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"

test/utils/test_iterables.py

@@ -0,0 +1,127 @@
+from supervision.utils.iterables import create_batches, fill
+
+
+def test_create_batches_when_empty_sequence_given() -> None:
+    # when
+    result = list(create_batches(sequence=[], batch_size=4))
+
+    # then
+    assert result == [], "Expected empty generator"
+
+
+def test_create_batches_when_not_allowed_batch_size_given() -> None:
+    # when
+    result = list(create_batches(sequence=[1, 2, 3], batch_size=0))
+
+    # then
+    assert result == [[1], [2], [3]], (
+        "Expected min_batch_size to be established and each element of input "
+        "list provided in separate batch"
+    )
+
+
+def test_create_batches_when_batch_size_larger_than_sequence() -> None:
+    # when
+    result = list(create_batches(sequence=[1, 2], batch_size=4))
+
+    # then
+    assert result == [[1, 2]], (
+        "Expected whole content to be returned in single batch as input sequence "
+        "is smaller than batch size"
+    )
+
+
+def test_create_batches_when_batch_size_multiplier_fits_sequence_length() -> None:
+    # when
+    result = list(create_batches(sequence=[1, 2, 3, 4], batch_size=2))
+
+    # then
+    assert result == [[1, 2], [3, 4]], (
+        "Expected input sequence to be returned in two chunks as batch size "
+        "is half of sequence length"
+    )
+
+
+def test_create_batches_when_batch_size_multiplier_does_not_fir_sequence_length() -> (
+    None
+):
+    # when
+    result = list(create_batches(sequence=[1, 2, 3, 4], batch_size=3))
+
+    # then
+    assert result == [[1, 2, 3], [4]], (
+        "Expected first batch to be of size 3 and last one to be not "
+        "full, with only one element"
+    )
+
+
+def test_fill_when_empty_sequence_given_and_padding_not_needed() -> None:
+    # given
+    sequence = []
+
+    # when
+    result = fill(sequence=sequence, desired_size=0, content=1)
+
+    # then
+    assert result == [], "Expected no elements to be added into sequence"
+
+
+def test_fill_when_empty_sequence_given_and_padding_needed() -> None:
+    # given
+    sequence = []
+
+    # when
+    result = fill(sequence=sequence, desired_size=3, content=1)
+
+    # then
+    assert result == [1, 1, 1], "Expected three padding element to be added"
+
+
+def test_fill_when_non_empty_sequence_given_and_sequence_equal_to_desired_size() -> (
+    None
+):
+    # given
+    sequence = [2, 2, 2]
+
+    # when
+    result = fill(sequence=sequence, desired_size=3, content=1)
+
+    # then
+    assert result == [
+        2,
+        2,
+        2,
+    ], "Expected nothing to be added to sequence, as it is already " "in desired size"
+
+
+def test_fill_when_non_empty_sequence_given_and_sequence_longer_then_desired_size() -> (
+    None
+):
+    # given
+    sequence = [2, 2, 2, 2]
+
+    # when
+    result = fill(sequence=sequence, desired_size=3, content=1)
+
+    # then
+    assert result == [
+        2,
+        2,
+        2,
+        2,
+    ], "Expected nothing to be added to sequence, as it already " "exceeds desired size"
+
+
+def test_fill_when_non_empty_sequence_given_and_padding_needed() -> None:
+    # given
+    sequence = [2]
+
+    # when
+    result = fill(sequence=sequence, desired_size=3, content=1)
+
+    # then
+    assert result == [
+        2,
+        1,
+        1,
+    ], "Expected 2 padding elements to be added to fit desired size"