supervision/tests/utils/test_image.py

import numpy as np
import pytest
from PIL import Image, ImageChops

from supervision.utils.image import (
    crop_image,
    get_image_resolution_wh,
    letterbox_image,
    resize_image,
)


def test_resize_image_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(
        image=image,
        resolution_wh=(1024, 1024),
        keep_aspect_ratio=True,
    )

    # then
    assert np.allclose(result, expected_result), (
        "Expected output shape to be (w, h): (1024, 768)"
    )


def test_resize_image_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(
        image=image,
        resolution_wh=(1024, 1024),
        keep_aspect_ratio=True,
    )

    # 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, resolution_wh=(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_grayscale_opencv_image() -> None:
    image = np.zeros((4, 6), dtype=np.uint8)
    expected_result = np.concatenate(
        [
            np.ones((2, 10), dtype=np.uint8) * 255,
            np.zeros((6, 10), dtype=np.uint8),
            np.ones((2, 10), dtype=np.uint8) * 255,
        ],
        axis=0,
    )

    result = letterbox_image(image=image, resolution_wh=(10, 10), color=(255, 255, 255))

    assert result.shape == (10, 10)
    assert np.array_equal(result, expected_result)


def test_letterbox_image_for_rgba_opencv_image() -> None:
    """RGBA input: padded alpha=0, interior alpha preserved, input array not mutated."""
    # given
    image = np.zeros((4, 6, 4), dtype=np.uint8)
    image[:, :, 3] = 128
    image_before = image.copy()

    # when
    result = letterbox_image(image=image, resolution_wh=(10, 10), color=(0, 0, 0))

    # then
    assert result.shape == (10, 10, 4)
    assert np.all(result[:2, :, 3] == 0), "padded top rows must have alpha=0"
    assert np.all(result[8:, :, 3] == 0), "padded bottom rows must have alpha=0"
    assert np.all(result[2:8, :, 3] == 128), "interior rows must preserve alpha"
    assert np.array_equal(image, image_before), "input must not be mutated"


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, resolution_wh=(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"
    )


@pytest.mark.parametrize(
    ("image", "xyxy", "expected_size"),
    [
        # NumPy RGB
        (
            np.zeros((4, 6, 3), dtype=np.uint8),
            (2, 1, 5, 3),
            (3, 2),  # width = 5-2, height = 3-1
        ),
        # NumPy grayscale
        (
            np.zeros((5, 5), dtype=np.uint8),
            (1, 1, 4, 4),
            (3, 3),
        ),
        # Pillow RGB
        (
            Image.new("RGB", (6, 4), color=0),
            (2, 1, 5, 3),
            (3, 2),
        ),
        # Pillow grayscale
        (
            Image.new("L", (5, 5), color=0),
            (1, 1, 4, 4),
            (3, 3),
        ),
    ],
)
def test_crop_image(image, xyxy, expected_size):
    cropped = crop_image(image=image, xyxy=xyxy)
    if isinstance(image, np.ndarray):
        assert isinstance(cropped, np.ndarray)
        assert cropped.shape[1] == expected_size[0]  # width
        assert cropped.shape[0] == expected_size[1]  # height
    else:
        assert isinstance(cropped, Image.Image)
        assert cropped.size == expected_size


@pytest.mark.parametrize(
    ("image", "expected"),
    [
        # NumPy RGB
        (np.zeros((4, 6, 3), dtype=np.uint8), (6, 4)),
        # NumPy grayscale
        (np.zeros((10, 20), dtype=np.uint8), (20, 10)),
        # Pillow RGB
        (Image.new("RGB", (6, 4), color=0), (6, 4)),
        # Pillow grayscale
        (Image.new("L", (20, 10), color=0), (20, 10)),
    ],
)
def test_get_image_resolution_wh(image, expected):
    resolution = get_image_resolution_wh(image)
    assert resolution == expected