From d27f5007e3f0305e7b1ef67dea9f6a1989568c35 Mon Sep 17 00:00:00 2001
From: CNOCycle <github@marginalman.org>
Date: Sat, 14 Sep 2024 09:31:23 +0000
Subject: [PATCH] Add a flag to be used for marking the YOLOv8 model

---
 .../object_detection/pytorch_object_detector.py       | 11 +++++++----
 art/estimators/object_detection/pytorch_yolo.py       |  3 +++
 notebooks/snal.ipynb                                  |  3 ++-
 tests/attacks/evasion/test_steal_now_attack_later.py  |  4 ++--
 4 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/art/estimators/object_detection/pytorch_object_detector.py b/art/estimators/object_detection/pytorch_object_detector.py
index 0a7d29f58d..0a9a2f52a8 100644
--- a/art/estimators/object_detection/pytorch_object_detector.py
+++ b/art/estimators/object_detection/pytorch_object_detector.py
@@ -65,6 +65,7 @@ def __init__(
             "loss_rpn_box_reg",
         ),
         device_type: str = "gpu",
+        is_yolov8: bool = False,
     ):
         """
         Initialization.
@@ -92,6 +93,7 @@ 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 re
         import torch
@@ -140,9 +142,10 @@ def __init__(
 
         self._model: torch.nn.Module
         self._model.to(self._device)
-        try:
+        self.is_yolov8 = is_yolov8
+        if self.is_yolov8:
             self._model.model.eval()
-        except AttributeError:
+        else:
             self._model.eval()
 
     @property
@@ -406,9 +409,9 @@ 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
-        try:
+        if self.is_yolov8:
             self._model.model.eval()
-        except AttributeError:
+        else:
             self._model.eval()
 
         # Apply preprocessing and convert to tensors
diff --git a/art/estimators/object_detection/pytorch_yolo.py b/art/estimators/object_detection/pytorch_yolo.py
index cd7fc69b55..cfe251e8bc 100644
--- a/art/estimators/object_detection/pytorch_yolo.py
+++ b/art/estimators/object_detection/pytorch_yolo.py
@@ -64,6 +64,7 @@ def __init__(
             "loss_rpn_box_reg",
         ),
         device_type: str = "gpu",
+        is_yolov8: bool = False,
     ):
         """
         Initialization.
@@ -92,6 +93,7 @@ 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.
         """
         super().__init__(
             model=model,
@@ -104,6 +106,7 @@ 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":
diff --git a/notebooks/snal.ipynb b/notebooks/snal.ipynb
index edeef5d0fd..22a663e9cc 100644
--- a/notebooks/snal.ipynb
+++ b/notebooks/snal.ipynb
@@ -71,7 +71,8 @@
     "model = YOLO('yolov8m')\n",
     "py_model = PyTorchYolo(model=model,\n",
     "                       input_shape=(3, 640, 640),\n",
-    "                       channels_first=True)\n",
+    "                       channels_first=True,\n",
+    "                       is_yolov8=True)\n",
     "\n",
     "# Define a custom function to collect patches from images\n",
     "def collect_patches_from_images(model: \"torch.nn.Module\",\n",
diff --git a/tests/attacks/evasion/test_steal_now_attack_later.py b/tests/attacks/evasion/test_steal_now_attack_later.py
index d58c49b46d..3eb3c35739 100644
--- a/tests/attacks/evasion/test_steal_now_attack_later.py
+++ b/tests/attacks/evasion/test_steal_now_attack_later.py
@@ -35,7 +35,7 @@ def test_generate(art_warning):
         import requests
 
         model = YOLO("yolov8m")
-        py_model = PyTorchYolo(model=model, input_shape=(3, 640, 640), channels_first=True)
+        py_model = PyTorchYolo(model=model, input_shape=(3, 640, 640), channels_first=True, is_yolov8=True)
 
         # Define a custom function to collect patches from images
         def collect_patches_from_images(model, imgs):
@@ -192,7 +192,7 @@ def _loader(self, path):
 def test_check_params(art_warning):
     try:
         model = YOLO("yolov8m")
-        py_model = PyTorchYolo(model=model, input_shape=(3, 640, 640), channels_first=True)
+        py_model = PyTorchYolo(model=model, input_shape=(3, 640, 640), channels_first=True, is_yolov8=True)
 
         def dummy_func(model, imags):
             candidates_patch = []