CI: download sam2 checkpoints for tests

Author: ajwitmer

.github/workflows/ci.yml

@@ -26,9 +26,21 @@ jobs:
     - name: install and lint
       run: |
         python -m pip install -e ./neat_ml/sam2/
-        python -m pip install -v ".[dev]"
+        python -m pip install -v -e ".[dev]"
         mypy neat_ml
         ruff check neat_ml
+    - name: Cache SAM2 checkpoint
+      id: cache-sam2
+      uses: actions/cache@v4
+      with:
+        path: neat_ml/sam2/checkpoints/sam2_hiera_large.pt
+        key: sam2-hiera-large-checkpoint
+    - name: Download SAM2 checkpoint
+      if: steps.cache-sam2.outputs.cache-hit != 'true'
+      run: |
+        mkdir -p neat_ml/sam2/checkpoints
+        wget -q https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt \
+          -O neat_ml/sam2/checkpoints/sam2_hiera_large.pt
     - name: install fonts
       if: runner.os == 'Linux'
       run: |

README.md

@@ -18,7 +18,7 @@ pip install -e ".[notebooks]"
 cd checkpoints
 # modify the provided bash script to download the appropriate checkpoints
 git apply ../../checkpoint.diff
-sh download_chkpts.sh
+sh download_ckpts.sh
 ```    
 
 NOTE:

neat_ml/bubblesam/SAM.py

@@ -5,6 +5,7 @@
 import logging
 import os
 os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+from importlib.abc import Traversable
 
 from sam2.build_sam import build_sam2
 from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
@@ -21,7 +22,7 @@ class SAMModel:
     def __init__(
         self, 
         model_config: str, 
-        checkpoint_path: str, 
+        checkpoint_path: Traversable, 
         device: str
         ) -> None:
         """
@@ -33,7 +34,7 @@ def __init__(
         ----------
         model_config : str
                 YAML cfg describing network architecture.
-        checkpoint_path : str
+        checkpoint_path : Path
                 Path to *.pt checkpoint with learned weights.
         device : str
                 Torch device ('cuda' | 'cpu' | 'cuda:0', …).

neat_ml/bubblesam/bubblesam.py

@@ -10,6 +10,7 @@
 import logging
 from matplotlib.axes import Axes
 from numpy.random import Generator
+from importlib.resources import files
 
 from skimage.measure import label, regionprops, find_contours
 from matplotlib.patches import Rectangle
@@ -18,7 +19,6 @@
 
 memory = joblib.Memory("joblib_cache", verbose=0)
 
-logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # these settings are used to parametrize 
@@ -44,9 +44,11 @@
     "use_m2m": True,
 }
 
+checkpoint_path = files("neat_ml.sam2").joinpath("checkpoints/sam2_hiera_large.pt")
+
 DEFAULT_MODEL_CFG = {
     "model_config": "sam2_hiera_l.yaml",
-    "checkpoint_path": "./neat_ml/sam2/checkpoints/sam2_hiera_large.pt",
+    "checkpoint_path": checkpoint_path,
     "device": ("cuda" if torch.cuda.is_available()
         else "mps" if torch.backends.mps.is_available()
         else "cpu"

neat_ml/tests/test_bubblesam.py

@@ -11,6 +11,8 @@
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 from matplotlib.testing.compare import compare_images
+from importlib.resources import files
+from importlib.abc import Traversable
 
 from neat_ml.bubblesam.bubblesam import (
     show_anns,
@@ -21,14 +23,14 @@
 )
 from neat_ml.bubblesam.SAM import SAMModel
 
-CHECKPOINT = "./neat_ml/sam2/checkpoints/sam2_hiera_large.pt"
+CHECKPOINT = files("neat_ml.sam2").joinpath("checkpoints/sam2_hiera_large.pt")
 
-def _skip_unless_available(model_chkpt: str = CHECKPOINT) -> None:
+def _skip_unless_available(model_chkpt: Traversable = CHECKPOINT) -> None:
     """
     Abort the whole module if we cannot load sam2 or the checkpoint.
     """
     pytest.importorskip("sam2", reason="sam2 package is required for SAM-2 tests")
-    if not Path(model_chkpt).exists():
+    if not model_chkpt.is_file():
         pytest.skip(
             f"SAM-2 checkpoint not found at {model_chkpt}. "
             "Install it to run integration tests.",
@@ -42,7 +44,7 @@ def _skip_unless_available(model_chkpt: str = CHECKPOINT) -> None:
     reason="This test is intended for systems without GPU support"
 )
 def test_setup_cuda_does_not_crash_on_cpu(
-    model_chkpt: str = CHECKPOINT,
+    model_chkpt: Traversable = CHECKPOINT,
 ):
     """
     Ensures that calling setup_cuda() in an environment with no GPU
@@ -62,7 +64,7 @@ def test_setup_cuda_does_not_crash_on_cpu(
     not torch.cuda.is_available(),
     reason="This test requires a CUDA-enabled GPU"
 )
-def test_setup_cuda_on_real_gpu(model_chkpt = CHECKPOINT):
+def test_setup_cuda_on_real_gpu(model_chkpt: Traversable = CHECKPOINT):
     """
     Verifies that setup_cuda() correctly configures torch backends on
     a live GPU. This test only runs if a CUDA device is found.
@@ -81,7 +83,7 @@ def test_setup_cuda_on_real_gpu(model_chkpt = CHECKPOINT):
         assert torch.backends.cudnn.allow_tf32
 
 @pytest.fixture(scope="module")
-def real_sam_model(model_chkpt: str = CHECKPOINT) -> SAMModel:
+def real_sam_model(model_chkpt: Traversable = CHECKPOINT) -> SAMModel:
     """
     Actual SAM-2 network on CPU
     """

neat_ml/tests/test_workflow.py

@@ -179,7 +179,7 @@ def test_stage_detect_pipeline_runs(
             chkpt_dir.mkdir(parents=True)
             shutil.copy(
                 files(neat_ml) / "sam2/checkpoints/sam2_hiera_large.pt", 
-                Path(tmpdir) / chkpt_dir,
+                chkpt_dir,
             )  # type: ignore[call-overload]
 
         ds["detection"]["img_dir"] = img_dir