Skip object detection tests if pycocotools is not installed. (#257)

IntelLabs · May 16, 2024 · d349977 · d349977
1 parent 3689976
commit d349977
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -14,14 +14,17 @@
 from hydra.core.global_hydra import GlobalHydra
 from omegaconf import DictConfig, open_dict
 
-from mart.utils.imports import _HAS_TIMM, _HAS_TORCHVISION
+from mart.utils.imports import _HAS_PYCOCOTOOLS, _HAS_TIMM, _HAS_TORCHVISION
 
 root = Path(os.getcwd())
 pyrootutils.set_root(path=root, dotenv=True, pythonpath=True)
 
 experiments_require_torchvision = [
     "CIFAR10_CNN",
     "CIFAR10_CNN_Adv",
+]
+
+experiments_require_torchvision_pycocotools = [
     "COCO_TorchvisionFasterRCNN",
     "COCO_TorchvisionFasterRCNN_Adv",
     "COCO_TorchvisionRetinaNet",
@@ -35,6 +38,8 @@
 experiments_names = []
 if _HAS_TORCHVISION:
     experiments_names += experiments_require_torchvision
+if _HAS_TORCHVISION and _HAS_PYCOCOTOOLS:
+    experiments_names += experiments_require_torchvision_pycocotools
 if _HAS_TIMM and _HAS_TORCHVISION:
     experiments_names += experiments_require_torchvision_and_timm
 

diff --git a/tests/test_experiments.py b/tests/test_experiments.py
@@ -4,7 +4,7 @@
 import pytest
 from hydra.core.global_hydra import GlobalHydra
 
-from mart.utils.imports import _HAS_TIMM, _HAS_TORCHVISION
+from mart.utils.imports import _HAS_PYCOCOTOOLS, _HAS_TIMM, _HAS_TORCHVISION
 from tests.helpers.dataset_generator import FakeCOCODataset
 from tests.helpers.run_if import RunIf
 from tests.helpers.run_sh_command import run_sh_command
@@ -131,7 +131,10 @@ def test_imagenet_timm_experiment(classification_cfg, tmp_path):
 
 @RunIf(sh=True)
 @pytest.mark.slow
-@pytest.mark.skipif(not _HAS_TORCHVISION, reason="test requires that torchvision is installed")
+@pytest.mark.skipif(
+    not _HAS_TORCHVISION or not _HAS_PYCOCOTOOLS,
+    reason="test requires that torchvision and pycocotools are installed",
+)
 def test_coco_fasterrcnn_experiment(coco_cfg, tmp_path):
     """Test TorchVision FasterRCNN experiment."""
     overrides = coco_cfg["trainer"] + coco_cfg["datamodel"]
@@ -147,7 +150,10 @@ def test_coco_fasterrcnn_experiment(coco_cfg, tmp_path):
 
 @RunIf(sh=True)
 @pytest.mark.slow
-@pytest.mark.skipif(not _HAS_TORCHVISION, reason="test requires that torchvision is installed")
+@pytest.mark.skipif(
+    not _HAS_TORCHVISION or not _HAS_PYCOCOTOOLS,
+    reason="test requires that torchvision and pycocotools are installed",
+)
 def test_coco_fasterrcnn_adv_experiment(coco_cfg, tmp_path):
     """Test TorchVision FasterRCNN Adv experiment."""
     overrides = coco_cfg["trainer"] + coco_cfg["datamodel"]
@@ -163,7 +169,10 @@ def test_coco_fasterrcnn_adv_experiment(coco_cfg, tmp_path):
 
 @RunIf(sh=True)
 @pytest.mark.slow
-@pytest.mark.skipif(not _HAS_TORCHVISION, reason="test requires that torchvision is installed")
+@pytest.mark.skipif(
+    not _HAS_TORCHVISION or not _HAS_PYCOCOTOOLS,
+    reason="test requires that torchvision and pycocotools  installed",
+)
 def test_coco_retinanet_experiment(coco_cfg, tmp_path):
     """Test TorchVision RetinaNet experiment."""
     overrides = coco_cfg["trainer"] + coco_cfg["datamodel"]