diff --git a/src/readyplayerme/meshops/draw/rasterize.py b/src/readyplayerme/meshops/draw/rasterize.py index ff87a87..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 @@ -27,7 +28,11 @@ def interpolate_values(start: Color, end: Color, num_steps: int) -> Color: raise ValueError(msg) t = np.arange(num_steps) / max(num_steps - 1, 1) - return start[None, :] + t[:, None] * (end - start) + + if start.size == 1 and start.ndim == 1: + return start + t * (end - start) + else: + return start[None, :] + t[:, None] * (end - start) def interpolate_segment(segment: Image) -> Image: @@ -206,6 +211,13 @@ def rasterize( if edges.size == 0 or image_coords.size == 0 or colors.size == 0: return clean_image(image, inplace=inplace) + # Check if the image and colors are compatible (both grayscale or both color) + 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 new file mode 100644 index 0000000..17af6ff --- /dev/null +++ b/src/readyplayerme/meshops/image.py @@ -0,0 +1,82 @@ +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). + """ + try: + n_channels = image.shape[-1] + except IndexError as error: + error_msg = "Image has invalid shape: zero dimensions." + raise ValueError(error_msg) from error + match image.ndim, n_channels: + case 2, _: + return ColorMode.GRAYSCALE + case 3, 3: + return ColorMode.RGB + case 3, 4: + return ColorMode.RGBA + case _: + error_msg = "Invalid color mode for an image." + raise ValueError(error_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). + """ + try: + n_channels = color_array.shape[-1] + except IndexError as error: + error_msg = "Color has invalid shape: zero dimensions." + raise ValueError(error_msg) from error + match color_array.ndim, n_channels: + case 1, _: + return ColorMode.GRAYSCALE + case 2, 1: + return ColorMode.GRAYSCALE + case 2, 3: + return ColorMode.RGB + case 2, 4: + return ColorMode.RGBA + 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. + + 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 image1: The first image to blend. + :param image2: The second image to blend. Must be the same shape as the first image. + :param mask: The blending mask. Must be the same shape as the images. + :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) + # 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] + # Perform the blending operation using vectorized NumPy operations + blended_image = (1 - mask) * image1 + mask * image2 + return blended_image + except AttributeError as error: + # Re-raise the original exception if it's not a shape mismatch + msg = "Could not blend the two images with the given mask" + raise AttributeError(msg) from error diff --git a/tests/readyplayerme/meshops/unit/test_image.py b/tests/readyplayerme/meshops/unit/test_image.py new file mode 100644 index 0000000..d01bf4d --- /dev/null +++ b/tests/readyplayerme/meshops/unit/test_image.py @@ -0,0 +1,155 @@ +"""Unit tests for the image module.""" + +import numpy as np +import pytest + +import readyplayerme.meshops.image as img +from readyplayerme.meshops.types import ColorMode + + +@pytest.mark.parametrize( + "image_1, image_2, mask, expected_output", + [ + # Basic blending with a uniform mask + ( + np.array([[0, 0], [0, 0]]), + np.array([[1, 1], [1, 1]]), + np.array([[0.5, 0.5], [0.5, 0.5]]), + np.array([[0.5, 0.5], [0.5, 0.5]]), + ), + # Blending with different mask values + ( + np.array([[0, 0], [0, 0]]), + np.array([[1, 1], [1, 1]]), + np.array([[0.2, 0.8], [0.3, 0.7]]), + np.array([[0.2, 0.8], [0.3, 0.7]]), + ), + # Mask with all zeros (full image_1) + ( + np.array([[0, 0], [0, 0]]), + np.array([[1, 1], [1, 1]]), + np.array([[0, 0], [0, 0]]), + np.array([[0, 0], [0, 0]]), + ), + # Mask with all ones (full image_2) + ( + np.array([[0, 0], [0, 0]]), + np.array([[1, 1], [1, 1]]), + 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): + """Test the blend_images function with various input scenarios.""" + output = img.blend_images(image_1, image_2, mask) + np.testing.assert_array_equal(output, expected_output) + + +@pytest.mark.parametrize( + "image_1, image_2, mask, expected_exception", + [ + # Shape mismatch + ( + np.array([[0, 0]]), # Shape (1, 2) + np.array([[1, 1], [1, 1], [1, 1]]), # Shape (3, 2) + np.array([[0.5, 0.5]]), # Shape (1, 2), + ValueError, + ), + # Invalid input type + ("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) diff --git a/tests/readyplayerme/meshops/unit/test_rasterizer.py b/tests/readyplayerme/meshops/unit/test_rasterizer.py index 2f8b033..c49605f 100644 --- a/tests/readyplayerme/meshops/unit/test_rasterizer.py +++ b/tests/readyplayerme/meshops/unit/test_rasterizer.py @@ -152,7 +152,7 @@ def test_draw_lines(image, edges, image_coords, colors, interpolate_func, expect np.array([0]), # Start color: Black in grayscale np.array([255]), # End color: White in grayscale 3, - np.array([[0], [127.5], [255]]), # Intermediate grayscale values + np.array([0, 127.5, 255]), # Intermediate grayscale values ), # Interpolation with more steps ( @@ -436,6 +436,26 @@ def test_rasterize( lambda img: np.nan_to_num(img).astype(np.uint8), IndexError, ), + # Mismatched Color Modes (Grayscale image with RGB colors) + ( + np.full((100, 100), np.nan, dtype=np.float32), # Grayscale image + np.array([[0, 1]]), + np.array([[10, 10], [20, 20]]), + np.array([[255, 0, 0], [0, 255, 0]]), # RGB colors + lambda color0, color1, steps: np.linspace(color0, color1, steps).astype(np.uint8), + lambda img: np.nan_to_num(img).astype(np.uint8), + ValueError, # Expecting a ValueError due to mismatched color modes + ), + # Mismatched Color Modes (RGB image with Grayscale colors) + ( + np.full((100, 100, 3), np.nan, dtype=np.float32), # RGB image + np.array([[0, 1]]), + np.array([[10, 10], [20, 20]]), + np.array([255, 0]), # Grayscale colors + lambda color0, color1, steps: np.linspace(color0, color1, steps).astype(np.uint8), + lambda img: np.nan_to_num(img).astype(np.uint8), + ValueError, # Expecting a ValueError due to mismatched color modes + ), ], ) def test_rasterize_should_fail(