supervision/test/utils/test_image.py

97 lines
2.7 KiB
Python

import numpy as np
from PIL import Image, ImageChops
from supervision.utils.image import 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_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"
)