From 734bdc16dcc9449be6f15401a4d54e3dd9ba1ce4 Mon Sep 17 00:00:00 2001 From: TechyDaniel <117300186+TechyDaniel@users.noreply.github.com> Date: Wed, 6 Dec 2023 19:26:57 +0100 Subject: [PATCH] feat: enhance blend_images tests - Add tests for non-uniform mask values in blend_images. - Include RGB images blending tests. - Handle edge cases and exception scenarios. - Cover different image sizes and empty masks. --- src/readyplayerme/meshops/draw/rasterize.py | 8 +- src/readyplayerme/meshops/image.py | 59 +++++++++++- .../readyplayerme/meshops/unit/test_image.py | 92 ++++++++++++++++++- 3 files changed, 151 insertions(+), 8 deletions(-) diff --git a/src/readyplayerme/meshops/draw/rasterize.py b/src/readyplayerme/meshops/draw/rasterize.py index 9d236b8..a256be7 100644 --- a/src/readyplayerme/meshops/draw/rasterize.py +++ b/src/readyplayerme/meshops/draw/rasterize.py @@ -3,6 +3,7 @@ import numpy as np import skimage +from readyplayerme.meshops.image import get_color_array_color_mode, get_image_color_mode from readyplayerme.meshops.types import Color, ColorMode, Edges, Image, PixelCoord @@ -211,11 +212,12 @@ def rasterize( return clean_image(image, inplace=inplace) # Check if the image and colors are compatible (both grayscale or both color) - is_image_grayscale = image.ndim == 2 # noqa: PLR2004 - is_colors_grayscale = (colors.ndim == 2 and colors.shape[-1] == 1) or (colors.ndim == 1) # noqa: PLR2004 - if is_image_grayscale != is_colors_grayscale: + image_mode = get_image_color_mode(image) + colors_mode = get_color_array_color_mode(colors) + if image_mode != colors_mode: msg = "Color mode of 'image' and 'colors' must match (both grayscale or both color)." raise ValueError(msg) + try: unique_indices = np.unique(edges.flatten()) # Failing early before proceeding with the code because draw line loops over the indices diff --git a/src/readyplayerme/meshops/image.py b/src/readyplayerme/meshops/image.py index 1fd19c2..0c29aa0 100644 --- a/src/readyplayerme/meshops/image.py +++ b/src/readyplayerme/meshops/image.py @@ -1,6 +1,57 @@ +from readyplayerme.meshops.types import Color, ColorMode from readyplayerme.meshops.types import Image as IMG_type +def get_image_color_mode(image: IMG_type) -> ColorMode: + """ + Determine the color mode of an image. + + :param image: An image array. + :return ColorMode:: Enum indicating the color mode of the image (GRAYSCALE, RGB, or RGBA). + """ + match image.ndim: + case 2: + return ColorMode.GRAYSCALE + case 3: + match image.shape[-1]: + case 3: + return ColorMode.RGB + case 4: + return ColorMode.RGBA + case _: + msg = "Invalid channel count for a color image." + raise ValueError(msg) + case _: + msg = "Invalid dimensions for an image." + raise ValueError(msg) + + +def get_color_array_color_mode(color_array: Color) -> ColorMode: + """ + Determine the color mode of a color array. + + :param color_array: An array representing colors. + :return ColorMode: Enum indicating the color mode of the color array (GRAYSCALE, RGB, or RGBA). + """ + match color_array.ndim: + case 1: + return ColorMode.GRAYSCALE + case 2: + match color_array.shape[-1]: + case 1: + return ColorMode.GRAYSCALE + case 3: + return ColorMode.RGB + case 4: + return ColorMode.RGBA + case _: + msg = "Invalid format for a color array." + raise ValueError(msg) + case _: + msg = "Invalid dimensions for a color array." + raise ValueError(msg) + + def blend_images(image1: IMG_type, image2: IMG_type, mask: IMG_type) -> IMG_type: """ Blend two images using a mask. @@ -8,17 +59,17 @@ def blend_images(image1: IMG_type, image2: IMG_type, mask: IMG_type) -> IMG_type This function performs a blending operation on two images using a mask. The mask determines the blending ratio at each pixel. The blending is done via a vectorized operation for efficiency. - :param image_1: The first image to blend, as a NumPy array. - :param image_2: The second image to blend, as a NumPy array. Must be the same shape as np_image. + :param image1: The first image to blend, as a NumPy array. + :param image2: The second image to blend, as a NumPy array. Must be the same shape as the first image. :param mask: The blending mask, as a NumPy array. Must be the same shape as the images. - :return: The blended image, as a NumPy array. + :return: The blended image. """ try: # Check if the error is due to shape mismatch if not (image1.shape == image2.shape): msg = "image1 and image2 must have the same shape." raise ValueError(msg) - # Define a constant variable for the expected number of dimensions + # Mask is reshaped to be able to perform the blending expected_dimensions = 2 if mask.ndim == expected_dimensions and image1.ndim == expected_dimensions + 1: mask = mask[:, :, None] diff --git a/tests/readyplayerme/meshops/unit/test_image.py b/tests/readyplayerme/meshops/unit/test_image.py index 7314b40..d01bf4d 100644 --- a/tests/readyplayerme/meshops/unit/test_image.py +++ b/tests/readyplayerme/meshops/unit/test_image.py @@ -4,6 +4,7 @@ import pytest import readyplayerme.meshops.image as img +from readyplayerme.meshops.types import ColorMode @pytest.mark.parametrize( @@ -37,6 +38,20 @@ np.array([[1, 1], [1, 1]]), np.array([[1, 1], [1, 1]]), ), + # Non uniform mask values: + ( + np.array([[0, 0], [0, 0]]), + np.array([[1, 1], [1, 1]]), + np.array([[0, 1], [1, 0]]), + np.array([[0, 1], [1, 0]]), + ), + # Full RGB Image + ( + np.zeros((2, 2, 3)), # Black RGB image + np.ones((2, 2, 3)), # White RGB image + np.full((2, 2), 0.5), # Uniform grayscale mask + np.full((2, 2, 3), 0.5), # Expected output: gray RGB image + ), ], ) def test_blend_images(image_1, image_2, mask, expected_output): @@ -56,10 +71,85 @@ def test_blend_images(image_1, image_2, mask, expected_output): ValueError, ), # Invalid input type - ("0, 0, 0", np.array([[1, 1], [1, 1]]), np.array([[0.5, 0.5], [0.5, 0.5]]), TypeError), + ("0, 0, 0", np.array([[1, 1], [1, 1]]), np.array([[0.5, 0.5], [0.5, 0.5]]), AttributeError), + # Empty mask + ( + np.array([[0, 0], [0, 0]]), + np.array([[1, 1], [1, 1]]), + np.array([]), + ValueError, + ), ], ) def test_blend_images_should_fail(image_1, image_2, mask, expected_exception): """Test the blend_images function with invalid input scenarios.""" with pytest.raises(expected_exception): img.blend_images(image_1, image_2, mask) + + +@pytest.mark.parametrize( + "image, expected_mode", + [ + # Grayscale image (2D array) + (np.array([[0, 1], [1, 0]], dtype=np.float32), ColorMode.GRAYSCALE), + # RGB image (3D array with shape (h, w, 3)) + (np.random.rand(10, 10, 3).astype(np.float32), ColorMode.RGB), + # RGBA image (3D array with shape (h, w, 4)) + (np.random.rand(10, 10, 4).astype(np.float32), ColorMode.RGBA), + ], +) +def test_get_image_color_mode(image, expected_mode): + """Test the get_image_color_mode function with valid inputs.""" + assert img.get_image_color_mode(image) == expected_mode + + +@pytest.mark.parametrize( + "image", + [ + # Invalid image: 2D array with incorrect channel count + np.random.rand(10, 10, 5).astype(np.float32), + # Invalid image: 1D array + np.array([1, 2, 3], dtype=np.float32), + # Invalid image: 4D array + np.random.rand(10, 10, 10, 3).astype(np.float32), + ], +) +def test_get_image_color_mode_should_fail(image): + """Test the get_image_color_mode function with invalid inputs.""" + with pytest.raises(ValueError): + img.get_image_color_mode(image) + + +@pytest.mark.parametrize( + "color_array, expected_mode", + [ + # Grayscale color array (1D array) + (np.array([128, 255, 100], dtype=np.uint8), ColorMode.GRAYSCALE), + # Grayscale color array (2D array with single channel) + (np.array([[128], [255], [100]], dtype=np.uint8), ColorMode.GRAYSCALE), + # RGB color array (2D array with shape (n, 3)) + (np.array([[255, 0, 0], [0, 255, 0], [0, 0, 255]], dtype=np.uint8), ColorMode.RGB), + # RGBA color array (2D array with shape (n, 4)) + (np.array([[255, 0, 0, 255], [0, 255, 0, 255], [0, 0, 255, 255]], dtype=np.uint8), ColorMode.RGBA), + ], +) +def test_get_color_array_color_mode(color_array, expected_mode): + """Test the get_color_array_color_mode function with valid inputs.""" + assert img.get_color_array_color_mode(color_array) == expected_mode + + +@pytest.mark.parametrize( + "color_array", + [ + # Invalid color array: 2D array with incorrect channel count + np.array([[1, 2, 3, 4, 5]], dtype=np.uint8), + # Invalid color array: 3D array + np.random.rand(10, 10, 3).astype(np.uint8), + # Invalid color array: 0-dimensional array + np.array(128, dtype=np.uint8), + ], +) +def test_get_color_array_color_mode_should_fail(color_array): + """Test the get_color_array_color_mode function with invalid inputs.""" + with pytest.raises(ValueError): + img.get_color_array_color_mode(color_array)