Fix OPS typing and tests; add coverage for edge branches

Signed-off-by: Nicholas Audric Adriel <[email protected]>

Author: nicholasadriel

art/attacks/poisoning/one_pixel_shortcut_attack.py

@@ -19,10 +19,13 @@
 This module implements One Pixel Shortcut attacks on Deep Neural Networks.
 """
 
+from typing import Optional, Tuple
+
 import numpy as np
 
 from art.attacks.attack import PoisoningAttackBlackBox
 
+
 class OnePixelShortcutAttack(PoisoningAttackBlackBox):
     """
     One-Pixel Shortcut (OPS) poisoning attack.
@@ -31,6 +34,7 @@ class OnePixelShortcutAttack(PoisoningAttackBlackBox):
     images. The found pixel coordinate and color are applied to all images of the class
     (labels remain unchanged). Reference: Wu et al. (ICLR 2023).
     """
+
     attack_params: list = []  # No external parameters for this attack
     _estimator_requirements: tuple = ()
 
@@ -41,7 +45,7 @@ def _check_params(self):
         # No parameters to validate
         pass
 
-    def poison(self, x: np.ndarray, y: np.ndarray = None, **kwargs):
+    def poison(self, x: np.ndarray, y: Optional[np.ndarray] = None, **kwargs) -> Tuple[np.ndarray, np.ndarray]:
         """
         Generate an OPS-poisoned dataset from clean data.
 
@@ -81,7 +85,7 @@ def poison(self, x: np.ndarray, y: np.ndarray = None, **kwargs):
                 else:
                     x_orig = np.transpose(x_array, (0, 2, 3, 1)).astype(np.float32)
                     channels_first = True
-            grayscale = (x_orig.shape[3] == 1)
+            grayscale = x_orig.shape[3] == 1
         else:
             raise ValueError(f"Unsupported input tensor shape: {x_array.shape}")
 
@@ -113,19 +117,16 @@ def poison(self, x: np.ndarray, y: np.ndarray = None, **kwargs):
                     np.array([1.0], dtype=x_orig.dtype),
                 ]
             else:
-                target_options = [
-                    np.array(bits, dtype=x_orig.dtype)
-                    for bits in np.ndindex(*(2,) * c)
-                ]
+                target_options = [np.array(bits, dtype=x_orig.dtype) for bits in np.ndindex(*(2,) * c)]
             # Evaluate each candidate color
             for target_vec in target_options:
                 # Compute per-image average difference from target for all pixels
-                diffs = np.abs(imgs_c - target_vec)              # shape (n_c, H, W, C)
-                per_image_diff = diffs.mean(axis=3)             # shape (n_c, H, W), mean diff per image at each pixel
+                diffs = np.abs(imgs_c - target_vec)  # shape (n_c, H, W, C)
+                per_image_diff = diffs.mean(axis=3)  # shape (n_c, H, W), mean diff per image at each pixel
                 # Compute score = mean - var for each pixel position (vectorized over HxW)
-                mean_diff_map = per_image_diff.mean(axis=0)     # shape (H, W)
-                var_diff_map = per_image_diff.var(axis=0)       # shape (H, W)
-                score_map = mean_diff_map - var_diff_map        # shape (H, W)
+                mean_diff_map = per_image_diff.mean(axis=0)  # shape (H, W)
+                var_diff_map = per_image_diff.var(axis=0)  # shape (H, W)
+                score_map = mean_diff_map - var_diff_map  # shape (H, W)
                 # Find the pixel with maximum score for this target
                 max_idx_flat = np.argmax(score_map)
                 max_score = score_map.ravel()[max_idx_flat]

tests/attacks/poison/test_one_pixel_shortcut_attack.py

@@ -22,10 +22,12 @@
 import pytest
 
 from art.attacks.poisoning.one_pixel_shortcut_attack import OnePixelShortcutAttack
+from unittest.mock import patch
 from tests.utils import ARTTestException
 
 logger = logging.getLogger(__name__)
 
+
 def test_one_pixel_per_image_and_label_preservation():
     try:
         x = np.zeros((4, 3, 3))
@@ -46,6 +48,7 @@ def test_one_pixel_per_image_and_label_preservation():
         logger.warning("test_one_pixel_per_image_and_label_preservation failed: %s", e)
         raise ARTTestException("Pixel change or label consistency check failed") from e
 
+
 def test_missing_labels_raises_error():
     try:
         x = np.zeros((3, 5, 5))
@@ -55,6 +58,7 @@ def test_missing_labels_raises_error():
         logger.warning("test_missing_labels_raises_error failed: %s", e)
         raise ARTTestException("Expected error not raised for missing labels") from e
 
+
 def test_multi_channel_consistency():
     try:
         x = np.zeros((2, 2, 2, 3))
@@ -77,6 +81,7 @@ def test_multi_channel_consistency():
         logger.warning("test_multi_channel_consistency failed: %s", e)
         raise ARTTestException("Multi-channel image consistency check failed") from e
 
+
 def test_one_pixel_effect_with_pytorchclassifier():
     try:
         import torch
@@ -130,13 +135,114 @@ def test_one_pixel_effect_with_pytorchclassifier():
         acc_poisoned = np.mean(preds_poisoned.argmax(axis=1) == y_poison)
 
         # Adjusted assertions for robustness
-        assert acc_poisoned >= 1.0, (
-            f"Expected 100% poisoned accuracy, got {acc_poisoned:.3f}"
-        )
+        assert acc_poisoned >= 1.0, f"Expected 100% poisoned accuracy, got {acc_poisoned:.3f}"
         assert acc_clean < 0.95, f"Expected clean accuracy < 95%, got {acc_clean:.3f}"
 
     except Exception as e:
         logger.warning("test_one_pixel_effect_with_pytorchclassifier failed: %s", e)
-        raise ARTTestException(
-            "PyTorchClassifier integration with OPS attack failed"
-        ) from e
+        raise ARTTestException("PyTorchClassifier integration with OPS attack failed") from e
+
+
+def test_check_params_noop():
+    try:
+        attack = OnePixelShortcutAttack()
+        # _check_params should do nothing (no error raised)
+        result = attack._check_params()
+        assert result is None
+    except Exception as e:
+        logger.warning("test_check_params_noop failed: %s", e)
+        raise ARTTestException("Parameter check method failed unexpectedly") from e
+
+
+def test_ambiguous_layout_nhwc():
+    try:
+        # Shape (N=1, H=3, W=2, C=3) - ambiguous, both 3 could be channels
+        x = np.zeros((1, 3, 2, 3), dtype=np.float32)
+        y = np.array([0])  # single sample, class 0
+        attack = OnePixelShortcutAttack()
+        x_p, y_p = attack.poison(x.copy(), y.copy())
+        # Output shape should match input shape
+        assert x_p.shape == x.shape and x_p.dtype == x.dtype
+        assert np.array_equal(y_p, y)
+        # Verify exactly one pixel (position) was changed
+        diff_any = np.any(x_p != x, axis=3)  # collapse channel differences
+        changes = np.sum(diff_any, axis=(1, 2))
+        assert np.all(changes == 1)
+    except Exception as e:
+        logger.warning("test_ambiguous_layout_nhwc failed: %s", e)
+        raise ARTTestException("Ambiguous NHWC layout handling failed") from e
+
+
+def test_ambiguous_layout_nchw():
+    try:
+        # Shape (N=1, C=2, H=3, W=2) - ambiguous, no dim equals 1/3/4
+        x = np.zeros((1, 2, 3, 2), dtype=np.float32)
+        y = np.array([0])
+        attack = OnePixelShortcutAttack()
+        x_p, y_p = attack.poison(x.copy(), y.copy())
+        # Output shape unchanged
+        assert x_p.shape == x.shape and x_p.dtype == x.dtype
+        assert np.array_equal(y_p, y)
+        # Exactly one pixel changed (channels-first input)
+        diff_any = np.any(x_p != x, axis=1)  # collapse channels axis
+        changes = np.sum(diff_any, axis=(1, 2))
+        assert np.all(changes == 1)
+    except Exception as e:
+        logger.warning("test_ambiguous_layout_nchw failed: %s", e)
+        raise ARTTestException("Ambiguous NCHW layout handling failed") from e
+
+
+def test_unsupported_input_shape_raises_error():
+    try:
+        x = np.zeros((5, 5), dtype=np.float32)  # 2D input (unsupported)
+        y = np.array([0, 1, 2, 3, 4])  # Dummy labels (not actually used due to error)
+        with pytest.raises(ValueError):
+            OnePixelShortcutAttack().poison(x, y)
+    except Exception as e:
+        logger.warning("test_unsupported_input_shape_raises_error failed: %s", e)
+        raise ARTTestException("ValueError not raised for unsupported input shape") from e
+
+
+def test_one_hot_labels_preserve_format():
+    try:
+        # Two 2x2 grayscale images, with one-hot labels for classes 0 and 1
+        x = np.zeros((2, 2, 2), dtype=np.float32)  # shape (N=2, H=2, W=2)
+        y = np.array([[1, 0], [0, 1]], dtype=np.float32)  # shape (2, 2) one-hot labels
+        attack = OnePixelShortcutAttack()
+        x_p, y_p = attack.poison(x.copy(), y.copy())
+        # Labels should remain one-hot and unchanged
+        assert y_p.shape == y.shape
+        assert np.array_equal(y_p, y)
+        # Exactly one pixel changed per image
+        changes = np.sum(x_p != x, axis=(1, 2))
+        assert np.all(changes == 1)
+    except Exception as e:
+        logger.warning("test_one_hot_labels_preserve_format failed: %s", e)
+        raise ARTTestException("One-hot label handling failed") from e
+
+
+def test_class_skipping_when_no_samples():
+    try:
+        # Small dataset: 1 image 2x2, class 0
+        x = np.zeros((1, 2, 2), dtype=np.float32)
+        y = np.array([0])
+        attack = OnePixelShortcutAttack()
+        # Baseline output without any patch
+        x_ref, y_ref = attack.poison(x.copy(), y.copy())
+        # Monkey-patch np.unique in attack module to add a non-existent class (e.g., class 1)
+        dummy_class = np.array([1], dtype=int)
+        orig_unique = np.unique
+        with patch(
+            "art.attacks.poisoning.one_pixel_shortcut_attack.np.unique",
+            new=lambda arr: np.concatenate([orig_unique(arr), dummy_class]),
+        ):
+            x_p, y_p = attack.poison(x.copy(), y.copy())
+        # Output with dummy class skip should match baseline output
+        assert np.array_equal(x_p, x_ref)
+        assert np.array_equal(y_p, y_ref)
+        # And still exactly one pixel changed
+        changes = np.sum(x_p != x, axis=(1, 2))
+        assert np.all(changes == 1)
+    except Exception as e:
+        logger.warning("test_class_skipping_when_no_samples failed: %s", e)
+        raise ARTTestException("Class-skipping branch handling failed") from e