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

@@ -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

src/readyplayerme/meshops/image.py

@@ -1,24 +1,75 @@
+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.
 
     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]

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)