Feat/blend images

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

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

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)

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(