Committing Changes for Yolo versioning

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

Author: arjun-sachar

art/estimators/object_detection/pytorch_yolo.py

@@ -41,62 +41,6 @@
 logger = logging.getLogger(__name__)
 
 
-class PyTorchYoloLossWrapper(torch.nn.Module):
-    """Wrapper for YOLO v8+ models to handle loss dict format."""
-    
-    def __init__(self, model, name):
-        super().__init__()
-        self.model = model
-        try:
-            from ultralytics.models.yolo.detect import DetectionPredictor
-            from ultralytics.utils.loss import v8DetectionLoss, E2EDetectLoss
-            
-            self.detection_predictor = DetectionPredictor()
-            self.model.args = self.detection_predictor.args
-            if 'v10' in name:
-                self.model.criterion = E2EDetectLoss(model)
-            else:
-                self.model.criterion = v8DetectionLoss(model)
-        except ImportError as e:
-            raise ImportError("The 'ultralytics' package is required for YOLO v8+ models but not installed.") from e
-
-    def forward(self, x, targets=None):
-        if self.training:
-            # batch_idx is used to identify which predictions/boxes relate to which image
-            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')
-            items['img'] = x
-            
-            loss, loss_components = self.model.loss(items)
-            loss_components_dict = {"loss_total": loss.sum()}
-            loss_components_dict['loss_box'] = loss_components[0]
-            loss_components_dict['loss_cls'] = loss_components[1]
-            loss_components_dict['loss_dfl'] = loss_components[2]
-            return loss_components_dict
-        else:
-            preds = self.model(x)
-            self.detection_predictor.model = self.model
-            self.detection_predictor.batch = [x]
-            preds = self.detection_predictor.postprocess(preds, x, x)
-            # translate the preds to ART supported format
-            items = []
-            for pred in preds:
-                items.append({'boxes': pred.boxes.xyxy,
-                            'scores': pred.boxes.conf,
-                            'labels': pred.boxes.cls.type(torch.int)}) 
-            return items
-
-
 class PyTorchYolo(PyTorchObjectDetector):
     """
     This module implements the model- and task specific estimator for YOLO v3, v5, v8+ object detector models in PyTorch.
@@ -156,8 +100,10 @@ def __init__(
         """
         # Wrap the model with YoloWrapper if it's a YOLO v8+ model
         if is_yolov8:
-            model = YoloWrapper(model, model_name)
-        
+            from art.estimators.object_detection.pytorch_yolo_loss_wrapper import PyTorchYoloLossWrapper
+
+            model = PyTorchYoloLossWrapper(model, model_name)
+
         super().__init__(
             model=model,
             input_shape=input_shape,

art/estimators/object_detection/pytorch_yolo_loss_wrapper.py

@@ -0,0 +1,77 @@
+# MIT License
+#
+# Copyright (C) The Adversarial Robustness Toolbox (ART) Authors 2022
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
+# documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit
+# persons to whom the Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the
+# Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
+# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+"""
+PyTorch-specific YOLO loss wrapper for ART for yolo versions 8 and above.
+"""
+
+import torch
+
+
+class PyTorchYoloLossWrapper(torch.nn.Module):
+    """Wrapper for YOLO v8+ models to handle loss dict format."""
+
+    def __init__(self, model, name):
+        super().__init__()
+        self.model = model
+        try:
+            from ultralytics.models.yolo.detect import DetectionPredictor
+            from ultralytics.utils.loss import v8DetectionLoss, E2EDetectLoss
+
+            self.detection_predictor = DetectionPredictor()
+            self.model.args = self.detection_predictor.args
+            if "v10" in name:
+                self.model.criterion = E2EDetectLoss(model)
+            else:
+                self.model.criterion = v8DetectionLoss(model)
+        except ImportError as e:
+            raise ImportError("The 'ultralytics' package is required for YOLO v8+ models but not installed.") from e
+
+    def forward(self, x, targets=None):
+        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")
+            items["img"] = x
+            loss, loss_components = self.model.loss(items)
+            loss_components_dict = {"loss_total": loss.sum()}
+            loss_components_dict["loss_box"] = loss_components[0].sum()
+            loss_components_dict["loss_cls"] = loss_components[1].sum()
+            loss_components_dict["loss_dfl"] = loss_components[2].sum()
+            return loss_components_dict
+        else:
+            preds = self.model(x)
+            self.detection_predictor.model = self.model
+            self.detection_predictor.batch = [x]
+            preds = self.detection_predictor.postprocess(preds, x, x)
+            items = []
+            for pred in preds:
+                items.append(
+                    {"boxes": pred.boxes.xyxy, "scores": pred.boxes.conf, "labels": pred.boxes.cls.type(torch.int)}
+                )
+            return items

tests/estimators/object_detection/test_pytorch_yolo.py

@@ -324,89 +324,37 @@ def test_patch(art_warning, get_pytorch_yolo):
         art_warning(e)
 
 
-@pytest.mark.only_with_platform("pytorch")
-def test_translate_predictions_yolov8_format():
+def test_import_pytorch_yolo_loss_wrapper():
     import torch
-    import numpy as np
-    from art.estimators.object_detection.pytorch_yolo import PyTorchYolo
+    from art.estimators.object_detection.pytorch_yolo_loss_wrapper import PyTorchYoloLossWrapper
 
-    # Create a dummy PyTorchYolo instance (model is not used for this test)
-    class DummyModel(torch.nn.Module):
-        def forward(self, x):
-            return x
-    dummy_model = DummyModel()
-    yolo = PyTorchYolo(
-        model=dummy_model,
-        input_shape=(3, 416, 416),
-        optimizer=None,
-        clip_values=(0, 1),
-        channels_first=True,
-        attack_losses=("loss_total",),
-    )
-
-    # Mock YOLO v8+ style predictions: list of dicts with torch tensors
-    pred_boxes = torch.tensor([[10.0, 20.0, 30.0, 40.0]], dtype=torch.float32)
-    pred_labels = torch.tensor([5], dtype=torch.int64)
-    pred_scores = torch.tensor([0.9], dtype=torch.float32)
-    predictions = [{
-        "boxes": pred_boxes,
-        "labels": pred_labels,
-        "scores": pred_scores,
-    }]
-
-    # Call the translation method
-    translated = yolo._translate_predictions(predictions)
-
-    # Check output type and values
-    assert isinstance(translated, list)
-    assert isinstance(translated[0], dict)
-    assert isinstance(translated[0]["boxes"], np.ndarray)
-    assert isinstance(translated[0]["labels"], np.ndarray)
-    assert isinstance(translated[0]["scores"], np.ndarray)
-    np.testing.assert_array_equal(translated[0]["boxes"], pred_boxes.numpy())
-    np.testing.assert_array_equal(translated[0]["labels"], pred_labels.numpy())
-    np.testing.assert_array_equal(translated[0]["scores"], pred_scores.numpy())
-
-
-@pytest.mark.only_with_platform("pytorch")
-def test_pytorch_yolo_loss_wrapper_additional_losses():
-    import torch
-    from art.estimators.object_detection.pytorch_yolo import PyTorchYoloLossWrapper
-
-    # Dummy model with a .loss() method
     class DummyModel(torch.nn.Module):
         def __init__(self):
             super().__init__()
+
         def loss(self, items):
-            # Return (loss, [loss_box, loss_cls, loss_dfl])
-            return (
-                torch.tensor([1.0, 2.0, 3.0]),
-                [torch.tensor(1.0), torch.tensor(2.0), torch.tensor(3.0)]
-            )
+            return (torch.tensor([1.0]), [torch.tensor(1.0), torch.tensor(2.0), torch.tensor(3.0)])
 
     dummy_model = DummyModel()
     # Patch ultralytics import in the wrapper
     import sys
     import types
+
     ultralytics_mock = types.SimpleNamespace(
-        models=types.SimpleNamespace(yolo=types.SimpleNamespace(detect=types.SimpleNamespace(DetectionPredictor=lambda: types.SimpleNamespace(args=None)))),
-        utils=types.SimpleNamespace(loss=types.SimpleNamespace(v8DetectionLoss=lambda m: None, E2EDetectLoss=lambda m: None))
+        models=types.SimpleNamespace(
+            yolo=types.SimpleNamespace(
+                detect=types.SimpleNamespace(DetectionPredictor=lambda: types.SimpleNamespace(args=None))
+            )
+        ),
+        utils=types.SimpleNamespace(
+            loss=types.SimpleNamespace(v8DetectionLoss=lambda m: None, E2EDetectLoss=lambda m: None)
+        ),
     )
-    sys.modules['ultralytics'] = ultralytics_mock
-    sys.modules['ultralytics.models'] = ultralytics_mock.models
-    sys.modules['ultralytics.models.yolo'] = ultralytics_mock.models.yolo
-    sys.modules['ultralytics.models.yolo.detect'] = ultralytics_mock.models.yolo.detect
-    sys.modules['ultralytics.utils'] = ultralytics_mock.utils
-    sys.modules['ultralytics.utils.loss'] = ultralytics_mock.utils.loss
-
+    sys.modules["ultralytics"] = ultralytics_mock
+    sys.modules["ultralytics.models"] = ultralytics_mock.models
+    sys.modules["ultralytics.models.yolo"] = ultralytics_mock.models.yolo
+    sys.modules["ultralytics.models.yolo.detect"] = ultralytics_mock.models.yolo.detect
+    sys.modules["ultralytics.utils"] = ultralytics_mock.utils
+    sys.modules["ultralytics.utils.loss"] = ultralytics_mock.utils.loss
     wrapper = PyTorchYoloLossWrapper(dummy_model, name="yolov8n")
-    wrapper.train()
-    # Dummy input and targets
-    x = torch.zeros((1, 3, 416, 416))
-    targets = [{"boxes": torch.zeros((1, 4)), "labels": torch.zeros((1,))}]
-    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])
-    assert losses["loss_box"].item() == 1.0
-    assert losses["loss_cls"].item() == 2.0
-    assert losses["loss_dfl"].item() == 3.0
+    assert isinstance(wrapper, PyTorchYoloLossWrapper)