Merge pull request #2686 from kieranfraser/dev_yolo_patch

Patching YOLO Ultralytics Upgrade

Author: beat-buesser

art/estimators/object_detection/pytorch_object_detector.py

@@ -66,7 +66,6 @@ def __init__(
             "loss_rpn_box_reg",
         ),
         device_type: str = "gpu",
-        is_yolov8: bool = False,
     ):
         """
         Initialization.
@@ -94,7 +93,6 @@ def __init__(
                               'loss_objectness', and 'loss_rpn_box_reg'.
         :param device_type: Type of device to be used for model and tensors, if `cpu` run on CPU, if `gpu` run on GPU
                             if available otherwise run on CPU.
-        :param is_yolov8: The flag to be used for marking the YOLOv8 model.
         """
         import torch
         import torchvision
@@ -139,11 +137,7 @@ def __init__(
 
         self._model: torch.nn.Module
         self._model.to(self._device)
-        self.is_yolov8 = is_yolov8
-        if self.is_yolov8:
-            self._model.model.eval()  # type: ignore
-        else:
-            self._model.eval()
+        self._model.eval()
 
     @property
     def native_label_is_pytorch_format(self) -> bool:
@@ -412,10 +406,7 @@ def predict(self, x: np.ndarray, batch_size: int = 128, **kwargs) -> list[dict[s
         from torch.utils.data import TensorDataset, DataLoader
 
         # Set model to evaluation mode
-        if self.is_yolov8:
-            self._model.model.eval()  # type: ignore
-        else:
-            self._model.eval()
+        self._model.eval()
 
         # Apply preprocessing and convert to tensors
         x_preprocessed, _ = self._preprocess_and_convert_inputs(x=x, y=None, fit=False, no_grad=True)

art/estimators/object_detection/pytorch_yolo.py

@@ -64,7 +64,7 @@ def __init__(
             "loss_rpn_box_reg",
         ),
         device_type: str = "gpu",
-        is_yolov8: bool = False,
+        is_ultralytics: bool = False,
         model_name: str | None = None,
     ):
         """
@@ -94,11 +94,9 @@ def __init__(
                               'loss_objectness', and 'loss_rpn_box_reg'.
         :param device_type: Type of device to be used for model and tensors, if `cpu` run on CPU, if `gpu` run on GPU
                             if available otherwise run on CPU.
-        :param is_yolov8: The flag to be used for marking the YOLOv8+ model.
         :param model_name: The name of the model (e.g., 'yolov8n', 'yolov10n') for determining loss function.
         """
-        # Wrap the model with YoloWrapper if it's a YOLO v8+ model
-        if is_yolov8:
+        if is_ultralytics:
             from art.estimators.object_detection.pytorch_yolo_loss_wrapper import PyTorchYoloLossWrapper
 
             model = PyTorchYoloLossWrapper(model, model_name)
@@ -114,7 +112,6 @@ def __init__(
             preprocessing=preprocessing,
             attack_losses=attack_losses,
             device_type=device_type,
-            is_yolov8=is_yolov8,
         )
 
     def _translate_labels(self, labels: list[dict[str, "torch.Tensor"]]) -> "torch.Tensor":

art/estimators/object_detection/pytorch_yolo_loss_wrapper.py

@@ -42,21 +42,14 @@ def __init__(self, model, name):
             raise ImportError("The 'ultralytics' package is required for YOLO v8+ models but not installed.") from e
 
     def forward(self, x, targets=None):
+        """Transforms the target to dict expected by model.loss"""
         if self.training:
-            boxes = []
-            labels = []
-            indices = []
-            for i, item in enumerate(targets):
-                boxes.append(item["boxes"])
-                labels.append(item["labels"])
-                indices = indices + ([i] * len(item["labels"]))
-            items = {
-                "boxes": torch.cat(boxes) / x.shape[2],
-                "labels": torch.cat(labels).type(torch.float32),
-                "batch_idx": torch.tensor(indices),
-            }
-            items["bboxes"] = items.pop("boxes")
-            items["cls"] = items.pop("labels")
+            if targets is None:
+                raise ValueError("Targets should not be None when training.")
+            items = {}
+            items["batch_idx"] = targets[:, 0]
+            items["bboxes"] = targets[:, 2:6]
+            items["cls"] = targets[:, 1]
             items["img"] = x
             loss, loss_components = self.model.loss(items)
             loss_components_dict = {"loss_total": loss.sum()}

tests/estimators/object_detection/test_pytorch_yolo_loss_wrapper.py

@@ -41,11 +41,23 @@ def test_yolov8_loss_wrapper():
     x = torch.randn((batch_size, 3, 640, 640))  # YOLOv8 expects (B, 3, 640, 640)
 
     # Create targets
-    targets = []
+    """targets = []
     for _ in range(batch_size):
         boxes = torch.tensor([[0.1, 0.1, 0.3, 0.3], [0.5, 0.5, 0.8, 0.8]])  # [x1, y1, x2, y2]
         labels = torch.zeros(2, dtype=torch.long)  # Use class 0 for testing
-        targets.append({"boxes": boxes, "labels": labels})
+        targets.append({"boxes": boxes, "labels": labels})"""
+    targets = torch.tensor(
+        [
+            [0.0000, 20.0000, 0.7738, 0.3919, 0.4525, 0.7582],
+            [0.0000, 20.0000, 0.2487, 0.4062, 0.4966, 0.5787],
+            [0.0000, 20.0000, 0.5667, 0.2772, 0.0791, 0.2313],
+            [0.0000, 20.0000, 0.1009, 0.1955, 0.2002, 0.0835],
+            [1.0000, 20.0000, 0.7738, 0.3919, 0.4525, 0.7582],
+            [1.0000, 20.0000, 0.2487, 0.4062, 0.4966, 0.5787],
+            [1.0000, 20.0000, 0.5667, 0.2772, 0.0791, 0.2313],
+            [1.0000, 20.0000, 0.1009, 0.1955, 0.2002, 0.0835],
+        ]
+    )
 
     # Test training mode
     losses = wrapper(x, targets)
@@ -94,11 +106,23 @@ def test_yolov10_loss_wrapper():
     x = torch.randn((batch_size, 3, 640, 640))  # Standard YOLO input size
 
     # Create targets
-    targets = []
+    """targets = []
     for _ in range(batch_size):
         boxes = torch.tensor([[0.1, 0.1, 0.3, 0.3], [0.5, 0.5, 0.8, 0.8]])  # [x1, y1, x2, y2]
         labels = torch.zeros(2, dtype=torch.long)  # Use class 0 for testing
-        targets.append({"boxes": boxes, "labels": labels})
+        targets.append({"boxes": boxes, "labels": labels})"""
+    targets = torch.tensor(
+        [
+            [0.0000, 20.0000, 0.7738, 0.3919, 0.4525, 0.7582],
+            [0.0000, 20.0000, 0.2487, 0.4062, 0.4966, 0.5787],
+            [0.0000, 20.0000, 0.5667, 0.2772, 0.0791, 0.2313],
+            [0.0000, 20.0000, 0.1009, 0.1955, 0.2002, 0.0835],
+            [1.0000, 20.0000, 0.7738, 0.3919, 0.4525, 0.7582],
+            [1.0000, 20.0000, 0.2487, 0.4062, 0.4966, 0.5787],
+            [1.0000, 20.0000, 0.5667, 0.2772, 0.0791, 0.2313],
+            [1.0000, 20.0000, 0.1009, 0.1955, 0.2002, 0.0835],
+        ]
+    )
 
     # Test training mode
     losses = wrapper(x, targets)
@@ -219,7 +243,7 @@ def loss(self, items):
     wrapper.train()
     # Dummy input and targets
     x = torch.zeros((1, 3, 416, 416))
-    targets = [{"boxes": torch.zeros((1, 4)), "labels": torch.zeros((1,))}]
+    targets = torch.tensor([[0.0000, 20.0000, 0.7738, 0.3919, 0.4525, 0.7582]])
     losses = wrapper(x, targets)
     assert set(losses.keys()) == {"loss_total", "loss_box", "loss_cls", "loss_dfl"}
     assert losses["loss_total"].item() == 6.0  # sum([1.0, 2.0, 3.0])
@@ -264,7 +288,7 @@ def loss(self, items):
     wrapper = PyTorchYoloLossWrapper(test_model, name="yolov8n")
     wrapper.train()
     x = torch.zeros((1, 3, 416, 416))
-    targets = [{"boxes": torch.zeros((1, 4)), "labels": torch.zeros((1,))}]
+    targets = torch.tensor([[0.0000, 20.0000, 0.7738, 0.3919, 0.4525, 0.7582]])
     losses = wrapper(x, targets)
     assert set(losses.keys()) == {"loss_total", "loss_box", "loss_cls", "loss_dfl"}
     assert losses["loss_total"].item() == 6.0
@@ -439,9 +463,7 @@ def loss(self, items):
     for batch_size in batch_sizes:
         for box_count in box_counts:
             x = torch.zeros((batch_size, 3, 416, 416))
-            targets = [
-                {"boxes": torch.zeros((box_count, 4)), "labels": torch.zeros(box_count)} for _ in range(batch_size)
-            ]
+            targets = torch.tensor([[0.0000, 20.0000, 0.7738, 0.3919, 0.4525, 0.7582]] * batch_size)
             losses = wrapper(x, targets)
 
             # Verify loss structure