Fix dependencies

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

Author: beat-buesser

tests/attacks/evasion/test_overload_attack.py

@@ -35,23 +35,41 @@ def test_generate(art_warning):
         from ultralytics.nn.modules import Conv
         from ultralytics import YOLO
 
+        import importlib
+        import inspect
+        from torch.serialization import add_safe_globals
+
         # torch.serialization.add_safe_globals([torch.nn.modules.container.Sequential])
         # torch.serialization.add_safe_globals([torch.nn.modules.container.ModuleList])
         # torch.serialization.add_safe_globals([torch.nn.modules.pooling.MaxPool2d])
         # torch.serialization.add_safe_globals([torch.nn.modules.batchnorm.BatchNorm2d])
         # torch.serialization.add_safe_globals([torch.nn.modules.activation.SiLU])
         # torch.serialization.add_safe_globals([torch.nn.modules.conv.Conv2d])
         # torch.serialization.add_safe_globals([torch.nn.modules.upsampling.Upsample])
-        torch.serialization.add_safe_globals([ultralytics.nn.tasks.DetectionModel])
-        torch.serialization.add_safe_globals([ultralytics.nn.tasks.SPPF])
-        torch.serialization.add_safe_globals([ultralytics.nn.tasks.C3])
-        torch.serialization.add_safe_globals([ultralytics.nn.tasks.Bottleneck])
-        torch.serialization.add_safe_globals([ultralytics.nn.tasks.Detect])
+        # torch.serialization.add_safe_globals([ultralytics.nn.tasks.DetectionModel])
+        # torch.serialization.add_safe_globals([ultralytics.nn.tasks.SPPF])
+        # torch.serialization.add_safe_globals([ultralytics.nn.tasks.C3])
+        # torch.serialization.add_safe_globals([ultralytics.nn.tasks.Bottleneck])
+        # torch.serialization.add_safe_globals([ultralytics.nn.tasks.Detect])
         # torch.serialization.add_safe_globals([Conv])
         # torch.serialization.add_safe_globals([ultralytics.nn.modules.Conv])
         # torch.serialization.add_safe_globals([ultralytics.nn.modules.Concat])
         # torch.serialization.add_safe_globals([ultralytics.nn.modules.DFL])
 
+        yolo_modules = [
+            "ultralytics.nn.modules",
+            "ultralytics.nn.tasks",
+            "ultralytics.nn.autobackend",
+        ]
+        for module_name in yolo_modules:
+            try:
+                mod = importlib.import_module(module_name)
+                for name, obj in inspect.getmembers(mod):
+                    if inspect.isclass(obj):
+                        add_safe_globals([obj])
+            except ModuleNotFoundError:
+                pass  # Some modules may not exist in all YOLO versions
+
         model = YOLO("yolov5su.pt")
 
         # Collect all unique classes used in the model
@@ -96,23 +114,41 @@ def test_check_params(art_warning):
         from ultralytics import YOLO
         from ultralytics.nn.modules import Conv
 
+        import importlib
+        import inspect
+        from torch.serialization import add_safe_globals
+
         # torch.serialization.add_safe_globals([torch.nn.modules.container.Sequential])
         # torch.serialization.add_safe_globals([torch.nn.modules.container.ModuleList])
         # torch.serialization.add_safe_globals([torch.nn.modules.pooling.MaxPool2d])
         # torch.serialization.add_safe_globals([torch.nn.modules.batchnorm.BatchNorm2d])
         # torch.serialization.add_safe_globals([torch.nn.modules.activation.SiLU])
         # torch.serialization.add_safe_globals([torch.nn.modules.conv.Conv2d])
         # torch.serialization.add_safe_globals([torch.nn.modules.upsampling.Upsample])
-        torch.serialization.add_safe_globals([ultralytics.nn.tasks.DetectionModel])
-        torch.serialization.add_safe_globals([ultralytics.nn.tasks.SPPF])
-        torch.serialization.add_safe_globals([ultralytics.nn.tasks.C3])
-        torch.serialization.add_safe_globals([ultralytics.nn.tasks.Bottleneck])
-        torch.serialization.add_safe_globals([ultralytics.nn.tasks.Detect])
+        # torch.serialization.add_safe_globals([ultralytics.nn.tasks.DetectionModel])
+        # torch.serialization.add_safe_globals([ultralytics.nn.tasks.SPPF])
+        # torch.serialization.add_safe_globals([ultralytics.nn.tasks.C3])
+        # torch.serialization.add_safe_globals([ultralytics.nn.tasks.Bottleneck])
+        # torch.serialization.add_safe_globals([ultralytics.nn.tasks.Detect])
         # torch.serialization.add_safe_globals([Conv])
         # torch.serialization.add_safe_globals([ultralytics.nn.modules.Conv])
         # torch.serialization.add_safe_globals([ultralytics.nn.modules.Concat])
         # torch.serialization.add_safe_globals([ultralytics.nn.modules.DFL])
 
+        yolo_modules = [
+            "ultralytics.nn.modules",
+            "ultralytics.nn.tasks",
+            "ultralytics.nn.autobackend",
+        ]
+        for module_name in yolo_modules:
+            try:
+                mod = importlib.import_module(module_name)
+                for name, obj in inspect.getmembers(mod):
+                    if inspect.isclass(obj):
+                        add_safe_globals([obj])
+            except ModuleNotFoundError:
+                pass  # Some modules may not exist in all YOLO versions
+
         model = YOLO("yolov5su.pt")
 
         # Collect all unique classes used in the model