extend bad det rma for arbitrary sizes

Signed-off-by: Farhan Ahmed <[email protected]>

Author: f4str

art/attacks/poisoning/bad_det/bad_det_rma.py

@@ -23,7 +23,7 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 
 import logging
-from typing import Dict, List, Tuple, Optional
+from typing import Dict, List, Tuple, Union, Optional
 
 import numpy as np
 from tqdm.auto import tqdm
@@ -82,36 +82,39 @@ def __init__(
 
     def poison(  # pylint: disable=W0221
         self,
-        x: np.ndarray,
+        x: Union[np.ndarray, List[np.ndarray]],
         y: List[Dict[str, np.ndarray]],
         **kwargs,
-    ) -> Tuple[np.ndarray, List[Dict[str, np.ndarray]]]:
+    ) -> Tuple[Union[np.ndarray, List[np.ndarray]], List[Dict[str, np.ndarray]]]:
         """
         Generate poisoning examples by inserting the backdoor onto the input `x` and changing the classification
         for labels `y`.
 
-        :param x: Sample images of shape `NCHW` or `NHWC`.
+        :param x: Sample images of shape `NCHW` or `NHWC` or a list of sample images of any size.
         :param y: True labels of type `List[Dict[np.ndarray]]`, one dictionary per input image. The keys and values
                   of the dictionary are:
 
                   - boxes [N, 4]: the boxes in [x1, y1, x2, y2] format, with 0 <= x1 < x2 <= W and 0 <= y1 < y2 <= H.
                   - labels [N]: the labels for each image.
-                  - scores [N]: the scores or each prediction.
         :return: An tuple holding the `(poisoning_examples, poisoning_labels)`.
         """
-        x_ndim = len(x.shape)
+        if isinstance(x, np.ndarray):
+            x_ndim = len(x.shape)
+        else:
+            x_ndim = len(x[0].shape) + 1
 
         if x_ndim != 4:
             raise ValueError("Unrecognized input dimension. BadDet RMA can only be applied to image data.")
 
-        if self.channels_first:
-            # NCHW --> NHWC
-            x = np.transpose(x, (0, 2, 3, 1))
-
-        x_poison = x.copy()
-        y_poison: List[Dict[str, np.ndarray]] = []
+        # copy images
+        x_poison: Union[np.ndarray, List[np.ndarray]]
+        if isinstance(x, np.ndarray):
+            x_poison = x.copy()
+        else:
+            x_poison = [x_i.copy() for x_i in x]
 
         # copy labels and find indices of the source class
+        y_poison: List[Dict[str, np.ndarray]] = []
         source_indices = []
         for i, y_i in enumerate(y):
             target_dict = {k: v.copy() for k, v in y_i.items()}
@@ -126,10 +129,12 @@ def poison(  # pylint: disable=W0221
 
         for i in tqdm(selected_indices, desc="BadDet RMA iteration", disable=not self.verbose):
             image = x_poison[i]
-
             boxes = y_poison[i]["boxes"]
             labels = y_poison[i]["labels"]
 
+            if self.channels_first:
+                image = np.transpose(image, (1, 2, 0))
+
             for j, (box, label) in enumerate(zip(boxes, labels)):
                 if self.class_source is None or label == self.class_source:
                     # extract the bounding box from the image
@@ -144,9 +149,10 @@ def poison(  # pylint: disable=W0221
                     # change the source label to the target label
                     labels[j] = self.class_target
 
-        if self.channels_first:
-            # NHWC --> NCHW
-            x_poison = np.transpose(x_poison, (0, 3, 1, 2))
+            # replace the original image with the poisoned image
+            if self.channels_first:
+                image = np.transpose(image, (2, 0, 1))
+            x_poison[i] = image
 
         return x_poison, y_poison