Repository: albumentations-team/albumentations

Issue:
Wrong type of parameters to `PixelDropout`, disallowing per-channel values
The parameters `drop_value` and `mask_drop_value` of `PixelDropout` have the wrong type `ScaleFloatType = Union[float, tuple[float, float]]`, see [here](https://github.com/albumentations-team/albumentations/blob/6ef2d979d7dcafcc3f8eed57d2184a33ec1c1af4/albumentations/augmentations/transforms.py#L4345). This makes it impossible to pass a triplet of values, i.e. one for each channel, like it was and still is supposed to be possible (according to docs and judging from the `prepare_drop_values` function).

E.g. this code:
```
import albumentations as A
A.PixelDropout(drop_value=(0, 0, 0))
```

fails with the error:
```
Traceback (most recent call last):
 File "<stdin>", line 1, in <module>
 File "[...]/albumentations/core/validation.py", line 35, in custom_init
   config = dct["InitSchema"](**full_kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 File "[...]/pydantic/main.py", line 214, in __init__
   validated_self = self.__pydantic_validator__.validate_python(data, self_instance=self)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
pydantic_core._pydantic_core.ValidationError: 2 validation errors for InitSchema
drop_value.float
 Input should be a valid number [type=float_type, input_value=(0, 0, 0), input_type=tuple]
   For further information visit https://errors.pydantic.dev/2.10/v/float_type
drop_value.tuple[float, float]
 Tuple should have at most 2 items after validation, not 3 [type=too_long, input_value=(0, 0, 0), input_type=tuple]
   For further information visit https://errors.pydantic.dev/2.10/v/too_long
```

Hints from issue comments:
@ternaus: @jleuschn Thank you. Fixed.

Solution patch:
diff --git a/albumentations/augmentations/functional.py b/albumentations/augmentations/functional.py
index a663469..d138e3b 100644
--- a/albumentations/augmentations/functional.py
+++ b/albumentations/augmentations/functional.py
@@ -2911,11 +2911,11 @@ def prepare_drop_values(
     else:
         values = np.array(value, dtype=array.dtype).reshape(-1)
 
-    # For 2D input, return single value
+    # For monochannel input, return single value
     if array.ndim == 2:
         return np.full(array.shape, values[0], dtype=array.dtype)
 
-    # For 3D input, broadcast values to full shape
+    # For multichannel input, broadcast values to full shape
     return np.full(array.shape[:2] + (len(values),), values, dtype=array.dtype)
 
 
diff --git a/albumentations/augmentations/transforms.py b/albumentations/augmentations/transforms.py
index 000a01a..b3ddf1b 100644
--- a/albumentations/augmentations/transforms.py
+++ b/albumentations/augmentations/transforms.py
@@ -4307,8 +4307,7 @@ class PixelDropout(DualTransform):
         mask_drop_value (float | Sequence[float] | None): Value to assign to dropped pixels in the mask.
             If None, the mask will remain unchanged.
             If a single number, that value will be used for all dropped pixels in the mask.
-            If a sequence, it should contain one value per channel of the mask.
-            Note: Only applicable when per_channel=False.
+            If a sequence, it should contain one value per channel.
             Default: None
 
         p (float): Probability of applying the transform. Should be in the range [0, 1].
@@ -4325,9 +4324,6 @@ class PixelDropout(DualTransform):
           if all pixels within the box are dropped. Such boxes will be removed.
         - When applied to keypoints, keypoints that fall on dropped pixels will be removed if
           the keypoint processor is configured to remove invisible keypoints.
-        - The 'per_channel' option is not supported for mask dropout. If you need to drop pixels
-          in a multi-channel mask independently, consider applying this transform multiple times
-          with per_channel=False.
 
     Example:
         >>> import numpy as np
@@ -4342,15 +4338,8 @@ class PixelDropout(DualTransform):
     class InitSchema(BaseTransformInitSchema):
         dropout_prob: ProbabilityType
         per_channel: bool
-        drop_value: ScaleFloatType | None
-        mask_drop_value: ScaleFloatType | None
-
-        @model_validator(mode="after")
-        def validate_mask_drop_value(self) -> Self:
-            if self.mask_drop_value is not None and self.per_channel:
-                msg = "PixelDropout supports mask only with per_channel=False."
-                raise ValueError(msg)
-            return self
+        drop_value: ColorType | None
+        mask_drop_value: ColorType | None
 
     _targets = ALL_TARGETS
 
@@ -4358,8 +4347,8 @@ class PixelDropout(DualTransform):
         self,
         dropout_prob: float = 0.01,
         per_channel: bool = False,
-        drop_value: ScaleFloatType | None = 0,
-        mask_drop_value: ScaleFloatType | None = None,
+        drop_value: ColorType | None = 0,
+        mask_drop_value: ColorType | None = None,
         p: float = 0.5,
     ):
         super().__init__(p=p)
@@ -4372,7 +4361,7 @@ class PixelDropout(DualTransform):
         self,
         img: np.ndarray,
         drop_mask: np.ndarray,
-        drop_values: float | Sequence[float],
+        drop_values: np.ndarray,
         **params: Any,
     ) -> np.ndarray:
         return fmain.pixel_dropout(img, drop_mask, drop_values)


Test patch:
diff --git a/tests/functional/test_functional.py b/tests/functional/test_functional.py
index c78f3a4..03cdc86 100644
--- a/tests/functional/test_functional.py
+++ b/tests/functional/test_functional.py
@@ -1500,6 +1500,34 @@ def test_get_drop_mask_reproducibility():
     assert not np.array_equal(mask1, mask3)
 
 
+def test_pixel_dropout_sequence_per_channel():
+    """Test pixel dropout with sequence values and per_channel=True"""
+    # Setup
+    image = np.ones((10, 10, 3), dtype=np.uint8) * 255
+    drop_values = (1, 2, 3)
+
+    # With dropout_prob=1.0, drop_mask will be all True
+    drop_mask = np.ones((10, 10, 3), dtype=bool)
+
+    # Verify drop_mask is all True
+    assert np.all(drop_mask)
+    assert drop_mask.shape == image.shape
+
+    # Prepare drop values
+    prepared_values = fmain.prepare_drop_values(image, drop_values, np.random.default_rng(42))
+
+    # Verify prepared_values shape matches image
+    assert prepared_values.shape == image.shape
+
+    # Apply dropout
+    result = fmain.pixel_dropout(image, drop_mask, prepared_values)
+
+    # Each channel should be entirely filled with its corresponding value
+    for channel_idx, expected_value in enumerate(drop_values):
+        assert np.all(result[:, :, channel_idx] == expected_value), \
+            f"Channel {channel_idx} should be filled with value {expected_value}"
+
+
 def test_prepare_drop_values_random_two_channels():
     """Test random value generation for 2-channel images."""
     rng = np.random.default_rng(42)
diff --git a/tests/test_augmentations.py b/tests/test_augmentations.py
index 004adbb..9e00d9a 100644
--- a/tests/test_augmentations.py
+++ b/tests/test_augmentations.py
@@ -1165,3 +1165,62 @@ def test_constrained_coarse_dropout_with_bboxes(bbox_labels, bboxes, expected_nu
                 overlaps_any = True
                 break
         assert overlaps_any, f"Hole {hole} doesn't overlap with any target box"
+
+
+
+@pytest.mark.parametrize(
+    ["drop_value", "expected_values"],
+    [
+        (None, None),  # Random values will be generated
+        (0, 0),  # Single value
+        ((1, 2, 3), np.array([1, 2, 3])),  # Sequence of values
+    ],
+)
+def test_pixel_dropout_drop_values(drop_value, expected_values):
+    image = np.ones((10, 10, 3), dtype=np.uint8) * 255
+    transform = A.PixelDropout(dropout_prob=1.0, drop_value=drop_value, p=1.0)
+
+    result = transform(image=image)["image"]
+
+    if drop_value is None:
+        # For None, we just verify values are within valid range
+        assert result.dtype == np.uint8
+        assert np.all((result >= 0) & (result <= 255))
+    else:
+        if isinstance(drop_value, (int, float)):
+            # For single value, all channels should have same value
+            assert np.all(result == expected_values)
+        else:
+            # For sequence, each channel should have corresponding value
+            for channel_idx, expected_value in enumerate(expected_values):
+                assert np.all(result[:, :, channel_idx] == expected_value)
+
+
+def test_pixel_dropout_per_channel():
+    """Test that per_channel=True works correctly with different drop_values"""
+    image = np.ones((10, 10, 3), dtype=np.uint8) * 255
+
+    # Test with single value
+    transform = A.PixelDropout(
+        dropout_prob=0.5,
+        drop_value=0,
+        per_channel=True,
+        p=1.0
+    )
+    result = transform(image=image)["image"]
+    assert np.any(result == 0)  # Should have some dropped pixels
+
+    # Test with sequence
+    transform = A.PixelDropout(
+        dropout_prob=0.5,
+        drop_value=(1, 2, 3),
+        per_channel=True,
+        p=1.0
+    )
+    result = transform(image=image)["image"]
+    # Each channel should only contain original values or its drop value
+    for channel_idx, drop_val in enumerate((1, 2, 3)):
+        unique_values = np.unique(result[:, :, channel_idx])
+        assert len(unique_values) == 2  # Should only have original value and drop value
+        assert drop_val in unique_values
+        assert 255 in unique_values