flags support for get_num_channels (#55)

* flags support for get_num_channels

* flags support for get_num_channels

* refresh tests

* update version

Author: trylinka

albucore/utils.py

@@ -122,12 +122,126 @@ def wrapped_function(img: np.ndarray, *args: P.args, **kwargs: P.kwargs) -> np.n
     return wrapped_function
 
 
-def get_num_channels(image: np.ndarray) -> int:
-    return image.shape[-1] if image.ndim >= NUM_MULTI_CHANNEL_DIMENSIONS else 1
+def get_num_channels(image: np.ndarray, has_batch_dim: bool = False, has_depth_dim: bool = False) -> int:
+    """Get the number of channels in an image array.
 
+    This function determines the number of channels in an image array by analyzing its shape
+    and accounting for optional batch and depth dimensions. The function assumes that the
+    last dimension represents channels when the array has more than 2 spatial dimensions
+    after accounting for any batch or depth dimensions.
 
-def is_grayscale_image(image: np.ndarray) -> bool:
-    return get_num_channels(image) == 1
+    Args:
+        image: Input image array. Can have various shapes:
+            - HW: (height, width) - grayscale image
+            - HWC: (height, width, channels) - multi-channel image
+            - NHW: (batch, height, width) - batch of grayscale images
+            - NHWC: (batch, height, width, channels) - batch of multi-channel images
+            - DHW: (depth, height, width) - 3D grayscale volume
+            - DHWC: (depth, height, width, channels) - 3D multi-channel volume
+            - DNHW: (depth, batch, height, width) - batch of 3D grayscale volumes
+            - DNHWC: (depth, batch, height, width, channels) - batch of 3D multi-channel volumes
+        has_batch_dim: If True, the first dimension is treated as a batch dimension (N).
+        has_depth_dim: If True, the first dimension (or second if has_batch_dim is True)
+                       is treated as a depth dimension (D).
+
+    Returns:
+        int: Number of channels in the image. Returns 1 for grayscale images and the size
+             of the last dimension for multi-channel images.
+
+    Examples:
+        >>> # 2D grayscale image
+        >>> img = np.zeros((100, 200))
+        >>> get_num_channels(img)
+        1
+
+        >>> # RGB image
+        >>> img = np.zeros((100, 200, 3))
+        >>> get_num_channels(img)
+        3
+
+        >>> # Batch of grayscale images
+        >>> img = np.zeros((10, 100, 200))
+        >>> get_num_channels(img, has_batch_dim=True)
+        1
+
+        >>> # Batch of RGB images
+        >>> img = np.zeros((10, 100, 200, 3))
+        >>> get_num_channels(img, has_batch_dim=True)
+        3
+
+        >>> # 3D volume
+        >>> img = np.zeros((5, 100, 200))
+        >>> get_num_channels(img, has_depth_dim=True)
+        1
+
+        >>> # Batch of 3D volumes with RGB
+        >>> img = np.zeros((5, 10, 100, 200, 3))
+        >>> get_num_channels(img, has_batch_dim=True, has_depth_dim=True)
+        3
+
+    Note:
+        The function assumes that after accounting for batch and depth dimensions,
+        the remaining dimensions follow the pattern HW (grayscale) or HWC (multi-channel).
+    """
+    # Calculate how many dimensions to skip from the beginning
+    dims_to_skip = int(has_depth_dim) + int(has_batch_dim)
+
+    # After skipping D and/or N dimensions, we should have HW or HWC
+    remaining_dims = image.ndim - dims_to_skip
+
+    # If we have more than 2 spatial dimensions (H, W), the last one is channels
+    # Otherwise, it's single channel
+    return image.shape[-1] if remaining_dims > 2 else 1
+
+
+def is_grayscale_image(image: np.ndarray, has_batch_dim: bool = False, has_depth_dim: bool = False) -> bool:
+    """Check if an image array represents a grayscale (single-channel) image.
+
+    This function determines whether an image has only one channel by calling get_num_channels
+    and checking if the result equals 1. It properly handles various array shapes including
+    batched images and 3D volumes.
+
+    Args:
+        image: Input image array. Can have various shapes as described in get_num_channels.
+        has_batch_dim: If True, the first dimension is treated as a batch dimension (N).
+        has_depth_dim: If True, the first dimension (or second if has_batch_dim is True)
+                       is treated as a depth dimension (D).
+
+    Returns:
+        bool: True if the image has only 1 channel (grayscale), False otherwise.
+
+    Examples:
+        >>> # 2D grayscale image
+        >>> img = np.zeros((100, 200))
+        >>> is_grayscale_image(img)
+        True
+
+        >>> # RGB image
+        >>> img = np.zeros((100, 200, 3))
+        >>> is_grayscale_image(img)
+        False
+
+        >>> # Single channel image with explicit channel dimension
+        >>> img = np.zeros((100, 200, 1))
+        >>> is_grayscale_image(img)
+        True
+
+        >>> # Batch of grayscale images
+        >>> img = np.zeros((10, 100, 200))
+        >>> is_grayscale_image(img, has_batch_dim=True)
+        True
+
+        >>> # Batch of RGB images
+        >>> img = np.zeros((10, 100, 200, 3))
+        >>> is_grayscale_image(img, has_batch_dim=True)
+        False
+
+    See Also:
+        get_num_channels: For getting the exact number of channels.
+        is_rgb_image: For checking if an image has exactly 3 channels (RGB).
+        is_multispectral_image: For checking if an image has channels other than 1 or 3.
+    """
+    return get_num_channels(image, has_batch_dim=has_batch_dim, has_depth_dim=has_depth_dim) == 1
 
 
 def get_opencv_dtype_from_numpy(value: np.ndarray | int | np.dtype | object) -> int:

pyproject.toml

@@ -5,7 +5,7 @@ requires = [ "setuptools>=45", "wheel" ]
 
 [project]
 name = "albucore"
-version = "0.0.25"
+version = "0.0.26"
 
 description = "High-performance image processing functions for deep learning and computer vision."
 readme = "README.md"

tests/test_utils.py

@@ -3,7 +3,7 @@
 import cv2
 from albucore.decorators import contiguous
 from albucore.functions import float32_io, from_float, to_float, uint8_io
-from albucore.utils import NPDTYPE_TO_OPENCV_DTYPE, clip, convert_value, get_opencv_dtype_from_numpy, get_num_channels
+from albucore.utils import NPDTYPE_TO_OPENCV_DTYPE, clip, convert_value, get_opencv_dtype_from_numpy, get_num_channels, is_grayscale_image
 
 
 @pytest.mark.parametrize("input_img, dtype, expected", [
@@ -257,3 +257,114 @@ def test_get_num_channels(shape, expected_channels, description):
     """Test get_num_channels for various array dimensions."""
     image = np.zeros(shape)
     assert get_num_channels(image) == expected_channels, f"Failed for {description} with shape {shape}"
+
+
+@pytest.mark.parametrize("shape, has_batch_dim, has_depth_dim, expected_channels, description", [
+    # HW: shape=(100, 200) → channels=1
+    ((100, 200), False, False, 1, "HW: grayscale image"),
+
+    # HWC: shape=(100, 200, 3) → channels=3
+    ((100, 200, 3), False, False, 3, "HWC: RGB image"),
+    ((100, 200, 1), False, False, 1, "HWC: single channel image"),
+    ((100, 200, 4), False, False, 4, "HWC: RGBA image"),
+
+    # NHW: shape=(10, 100, 200) → channels=1 (batch of grayscale)
+    ((10, 100, 200), True, False, 1, "NHW: batch of grayscale images"),
+
+    # NHWC: shape=(10, 100, 200, 3) → channels=3 (batch of RGB)
+    ((10, 100, 200, 3), True, False, 3, "NHWC: batch of RGB images"),
+    ((10, 100, 200, 1), True, False, 1, "NHWC: batch of single channel images"),
+
+    # DHW: shape=(5, 100, 200) → channels=1 (3D volume)
+    ((5, 100, 200), False, True, 1, "DHW: 3D volume"),
+
+    # DHWC: shape=(5, 100, 200, 3) → channels=3 (3D volume with RGB slices)
+    ((5, 100, 200, 3), False, True, 3, "DHWC: 3D volume with RGB slices"),
+    ((5, 100, 200, 1), False, True, 1, "DHWC: 3D volume with single channel"),
+
+    # DNHW: shape=(5, 10, 100, 200) → channels=1 (batch of 3D volumes)
+    ((5, 10, 100, 200), True, True, 1, "DNHW: batch of 3D volumes"),
+
+    # DNHWC: shape=(5, 10, 100, 200, 3) → channels=3 (batch of 3D volumes with RGB)
+    ((5, 10, 100, 200, 3), True, True, 3, "DNHWC: batch of 3D volumes with RGB"),
+    ((5, 10, 100, 200, 1), True, True, 1, "DNHWC: batch of 3D volumes with single channel"),
+
+    # Additional edge cases
+    ((32, 32), False, False, 1, "HW: square grayscale"),
+    ((224, 224, 3), False, False, 3, "HWC: standard RGB image size"),
+    ((1, 512, 512), True, False, 1, "NHW: single image in batch"),
+    ((1, 512, 512, 3), True, False, 3, "NHWC: single RGB image in batch"),
+])
+def test_get_num_channels_with_dimension_flags(shape, has_batch_dim, has_depth_dim, expected_channels, description):
+    """Test get_num_channels with batch and depth dimension flags."""
+    image = np.zeros(shape)
+    result = get_num_channels(image, has_batch_dim=has_batch_dim, has_depth_dim=has_depth_dim)
+    assert result == expected_channels, f"Failed for {description} with shape {shape}, has_batch_dim={has_batch_dim}, has_depth_dim={has_depth_dim}"
+
+
+@pytest.mark.parametrize("shape, has_batch_dim, has_depth_dim, expected_grayscale, description", [
+    # HW: shape=(100, 200) → grayscale=True
+    ((100, 200), False, False, True, "HW: grayscale image"),
+
+    # HWC: shape=(100, 200, 3) → grayscale=False
+    ((100, 200, 3), False, False, False, "HWC: RGB image"),
+    ((100, 200, 1), False, False, True, "HWC: single channel image"),
+    ((100, 200, 4), False, False, False, "HWC: RGBA image"),
+
+    # NHW: shape=(10, 100, 200) → grayscale=True
+    ((10, 100, 200), True, False, True, "NHW: batch of grayscale images"),
+
+    # NHWC: shape=(10, 100, 200, 3) → grayscale=False
+    ((10, 100, 200, 3), True, False, False, "NHWC: batch of RGB images"),
+    ((10, 100, 200, 1), True, False, True, "NHWC: batch of single channel images"),
+
+    # DHW: shape=(5, 100, 200) → grayscale=True
+    ((5, 100, 200), False, True, True, "DHW: 3D volume"),
+
+    # DHWC: shape=(5, 100, 200, 3) → grayscale=False
+    ((5, 100, 200, 3), False, True, False, "DHWC: 3D volume with RGB slices"),
+    ((5, 100, 200, 1), False, True, True, "DHWC: 3D volume with single channel"),
+
+    # DNHW: shape=(5, 10, 100, 200) → grayscale=True
+    ((5, 10, 100, 200), True, True, True, "DNHW: batch of 3D volumes"),
+
+    # DNHWC: shape=(5, 10, 100, 200, 3) → grayscale=False
+    ((5, 10, 100, 200, 3), True, True, False, "DNHWC: batch of 3D volumes with RGB"),
+    ((5, 10, 100, 200, 1), True, True, True, "DNHWC: batch of 3D volumes with single channel"),
+])
+def test_is_grayscale_image(shape, has_batch_dim, has_depth_dim, expected_grayscale, description):
+    """Test is_grayscale_image with various shape combinations."""
+    image = np.zeros(shape)
+    result = is_grayscale_image(image, has_batch_dim=has_batch_dim, has_depth_dim=has_depth_dim)
+    assert result == expected_grayscale, f"Failed for {description} with shape {shape}, has_batch_dim={has_batch_dim}, has_depth_dim={has_depth_dim}"
+
+
+@pytest.mark.parametrize("shape, has_batch_dim, has_depth_dim", [
+    # Basic 2D cases
+    ((100, 200), False, False),
+    ((100, 200, 1), False, False),
+    ((100, 200, 3), False, False),
+    # Batch cases (NHW/NHWC)
+    ((10, 100, 200), True, False),
+    ((10, 100, 200, 1), True, False),
+    ((10, 100, 200, 3), True, False),
+    # Depth cases (DHW/DHWC)
+    ((5, 100, 200), False, True),
+    ((5, 100, 200, 1), False, True),
+    ((5, 100, 200, 3), False, True),
+    # Batch and depth cases (DNHW/DNHWC)
+    ((5, 10, 100, 200), True, True),
+    ((5, 10, 100, 200, 1), True, True),
+    ((5, 10, 100, 200, 3), True, True),
+])
+def test_get_num_channels_and_is_grayscale_consistency(shape, has_batch_dim, has_depth_dim):
+    """Test that get_num_channels and is_grayscale_image are consistent."""
+    image = np.zeros(shape)
+    num_channels = get_num_channels(image, has_batch_dim=has_batch_dim, has_depth_dim=has_depth_dim)
+    is_grayscale = is_grayscale_image(image, has_batch_dim=has_batch_dim, has_depth_dim=has_depth_dim)
+
+    # is_grayscale should be True if and only if num_channels == 1
+    assert (num_channels == 1) == is_grayscale, (
+        f"Inconsistency for shape {shape}, has_batch_dim={has_batch_dim}, has_depth_dim={has_depth_dim}: "
+        f"num_channels={num_channels}, is_grayscale={is_grayscale}"
+    )