Merge pull request #107 from KatherLab/v2.3

- change project version
- add version check on both mlp and vision models
- add version check on heatmaps
- update example checkpoint used in tests

Author: EzicStar

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "stamp"
-version = "2.2.0"
+version = "2.3.0"
 authors = [
     { name = "Omar El Nahhas", email = "omar.el_nahhas@tu-dresden.de" },
     { name = "Marko van Treeck", email = "markovantreeck@gmail.com" },

src/stamp/heatmaps/__init__.py

@@ -22,6 +22,8 @@
 from stamp.preprocessing import supported_extensions
 from stamp.preprocessing.tiling import get_slide_mpp_
 from stamp.types import DeviceLikeType, Microns, SlideMPP, TilePixels
+from packaging.version import Version
+
 
 _logger = logging.getLogger("stamp")
 
@@ -168,8 +170,6 @@ def heatmaps_(
     topk: int,
     bottomk: int,
 ) -> None:
-    model = LitVisionTransformer.load_from_checkpoint(checkpoint_path).to(device).eval()
-
     # Collect slides to generate heatmaps for
     if slide_paths is not None:
         wsis_to_process = (wsi_dir / slide for slide in slide_paths)
@@ -227,6 +227,17 @@ def heatmaps_(
         # coordinates as used by OpenSlide
         coords_tile_slide_px = torch.round(coords_um / slide_mpp).long()
 
+        model = (
+            LitVisionTransformer.load_from_checkpoint(checkpoint_path).to(device).eval()
+        )
+
+        # TODO: Update version when a newer model logic breaks heatmaps.
+        if Version(model.stamp_version) < Version("2.3.0"):
+            raise ValueError(
+                f"model has been built with stamp version {model.stamp_version} "
+                f"which is incompatible with the current version."
+            )
+
         # Score for the entire slide
         slide_score = (
             model.vision_transformer(

src/stamp/modeling/lightning_model.py

@@ -114,14 +114,15 @@ def __init__(
         self.categories = np.array(categories)
         self.train_patients = train_patients
         self.valid_patients = valid_patients
+        self.stamp_version = str(stamp_version)
 
         _ = metadata  # unused, but saved in model
 
         # Check if version is compatible.
         # This should only happen when the model is loaded,
         # otherwise the default value will make these checks pass.
         # TODO: Change this on version change
-        if stamp_version < Version("2.0.0.dev8"):
+        if stamp_version < Version("2.3.0"):
             # Update this as we change our model in incompatible ways!
             raise ValueError(
                 f"model has been built with stamp version {stamp_version} "

src/stamp/modeling/mlp_classifier.py

@@ -83,8 +83,26 @@ def __init__(
         self.total_steps = total_steps
         self.max_lr = max_lr
         self.div_factor = div_factor
-
-    # TODO: Add version check with version 2.2.1, for both MLP and Transformer
+        self.stamp_version = str(stamp_version)
+
+        # Check if version is compatible.
+        # This should only happen when the model is loaded,
+        # otherwise the default value will make these checks pass.
+        # TODO: Change this on version change
+        if stamp_version < Version("2.3.0"):
+            # Update this as we change our model in incompatible ways!
+            raise ValueError(
+                f"model has been built with stamp version {stamp_version} "
+                f"which is incompatible with the current version."
+            )
+        elif stamp_version > Version(stamp.__version__):
+            # Let's be strict with models "from the future",
+            # better fail deadly than have broken results.
+            raise ValueError(
+                "model has been built with a stamp version newer than the installed one "
+                f"({stamp_version} > {stamp.__version__}). "
+                "Please upgrade stamp to a compatible version."
+            )
 
     def forward(self, x: Tensor) -> Tensor:
         return self.model(x)

tests/test_deployment_backward_compatibility.py

@@ -13,9 +13,9 @@
 )
 def test_backwards_compatibility() -> None:
     example_checkpoint_path = download_file(
-        url="https://github.com/KatherLab/STAMP/releases/download/2.2.0/example-model.ckpt",
-        file_name="example-modelv2_3.ckpt",
-        sha256sum="1caa5b56a02d09e6df72f087ebe9f87bae4e53de0516e2291e6dd1ce7dfff054",
+        url="https://github.com/KatherLab/STAMP/releases/download/2.2.0/example-model-v2_3_0.ckpt",
+        file_name="example-modelv2_3_0.ckpt",
+        sha256sum="eb6225fcdea7f33dee80fd5dc4e7a0da6cd0d91a758e3ee9605d6869b30ab657",
     )
     example_feature_path = download_file(
         url="https://github.com/KatherLab/STAMP/releases/download/2.2.0/TCGA-AA-3877-01Z-00-DX1.36902310-bc0b-4437-9f86-6df85703e0ad.h5",

tests/test_heatmaps.py

@@ -10,9 +10,9 @@
 @pytest.mark.filterwarnings("ignore:There is a performance drop")
 def test_heatmap_integration(tmp_path: Path) -> None:
     example_checkpoint_path = download_file(
-        url="https://github.com/KatherLab/STAMP/releases/download/2.2.0/example-model.ckpt",
-        file_name="example-modelv2_3.ckpt",
-        sha256sum="1caa5b56a02d09e6df72f087ebe9f87bae4e53de0516e2291e6dd1ce7dfff054",
+        url="https://github.com/KatherLab/STAMP/releases/download/2.2.0/example-model-v2_3_0.ckpt",
+        file_name="example-modelv2_3_0.ckpt",
+        sha256sum="eb6225fcdea7f33dee80fd5dc4e7a0da6cd0d91a758e3ee9605d6869b30ab657",
     )
     example_slide_path = download_file(
         url="https://github.com/KatherLab/STAMP/releases/download/2.0.0.dev14/TCGA-G4-6625-01Z-00-DX1.0fa26667-2581-4f96-a891-d78dbc3299b4.svs",