Upgrade Cellpose SAM example

example/export_cellpose_model/bioimageio_export.py

@@ -1,3 +1,4 @@
+import hashlib
 import os
 from pathlib import Path
 
@@ -6,7 +7,7 @@
 from bioimageio.core import test_model
 from bioimageio.spec import save_bioimageio_package
 from bioimageio.spec.model.v0_5 import (
-    ArchitectureFromLibraryDescr,
+    ArchitectureFromFileDescr,
     AxisId,
     BatchAxis,
     BioimageioConfig,
@@ -34,11 +35,53 @@
     WeightsDescr,
 )
 
+# Cellpose version installed in this environment (`pip show cellpose`). Pinned so
+# the exported model resolves `cellpose.vit.CPnetBioImageIO` against the same code.
+CELLPOSE_VERSION = "4.2.1.1"
+
+
+def create_environment_file_for_model(building_dir: Path) -> Path:
+    """Write a conda environment.yaml pinning the deps needed to run the model.
+
+    Mirrors BiaPy's `create_environment_file_for_model` so the bioimage.io
+    `PytorchStateDictWeightsDescr.dependencies` field points at a reproducible env.
+    """
+    env_yaml = (
+        "name: cellpose-sam\n"
+        "channels:\n"
+        "  - conda-forge\n"
+        "  - nodefaults\n"
+        "dependencies:\n"
+        "  - python>=3.11\n"
+        "  - pip\n"
+        "  - pip:\n"
+        f"      - cellpose=={CELLPOSE_VERSION}\n"
+    )
+    building_dir.mkdir(parents=True, exist_ok=True)
+    env_file = building_dir / "environment.yaml"
+    with open(env_file, "w", encoding="utf8") as outfile:
+        outfile.write(env_yaml)
+    return env_file
+
+
+def sha256sum(path: Path) -> str:
+    """Return the hex sha256 digest of a file."""
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(1 << 20), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
 if __name__ == "__main__":
     logger.enable("bioimageio")
 
     os.chdir(Path(__file__).parent)
 
+    # Conda environment file pinning cellpose (and bioimageio.core) so the model
+    # works out of the box. Attached to the pytorch weights via `dependencies`.
+    env_file = create_environment_file_for_model(Path("output"))
+    env_descriptor = FileDescr(source=env_file, sha256=Sha256(sha256sum(env_file)))
+
     descr = ModelDescr(
         name="Cellpose-SAM",
         inputs=[
@@ -119,12 +162,16 @@
                 sha256=Sha256(
                     "e1440429eb384f95afe32bcba6510f90d518eaedc917ede549bed6804004abe2"
                 ),
-                architecture=ArchitectureFromLibraryDescr(
-                    callable=Identifier("Transformer"),
-                    import_from="cellpose.vit_sam",
+                architecture=ArchitectureFromFileDescr(
+                    # Bundle the fixed CPnetBioImageIO so the model runs against a
+                    # stock cellpose install (no package patch / upstream PR needed).
+                    source=Path("cellpose_vit_bioimageio.py"),
+                    sha256=Sha256(sha256sum(Path("cellpose_vit_bioimageio.py"))),
+                    callable=Identifier("CPnetBioImageIO"),
                 ),
                 pytorch_version=Version("2.10.0"),
                 strict=False,
+                dependencies=env_descriptor,
             ),
         ),
         config=Config(
@@ -145,5 +192,5 @@
         descr, output_path=Path("output/cpsam_bioimageio.zip")
     )
 
-    summary = test_model(descr, working_dir=Path("output/export_test"))
+    summary = test_model(descr, working_dir=Path("output/export_test"), absolute_tolerance=1e-3, relative_tolerance=1e-3)
     summary.display()

example/export_cellpose_model/cellpose_original.py

@@ -8,7 +8,12 @@
 import numpy as np
 
 if __name__ == "__main__":
-    cellpose_original = cellpose.models.CellposeModel(gpu=False)
+    # Pin the exact same checkpoint the bioimageio export uses (HuggingFace `cpsam`,
+    # sha256 e1440429...). CellposeModel() otherwise defaults to `cpsam_v2`, whose
+    # weights differ (347/348 tensors), producing a non-reproducible reference.
+    cellpose_original = cellpose.models.CellposeModel(
+        gpu=False, pretrained_model="/home/dfranco/.cellpose/models/cpsam"
+    )
 
     os.chdir(Path(__file__).parent)
     input_array = imageio.imread("sample_input.png").transpose(2, 0, 1)

example/export_cellpose_model/cellpose_vit_bioimageio.py

@@ -0,0 +1,49 @@
+"""BioImage.IO-compatible Cellpose-SAM architecture (bundled with the model).
+
+This module ships a fixed ``CPnetBioImageIO`` so the exported bioimage.io model runs
+against a **stock, unmodified** cellpose install. It overrides only the three methods
+of cellpose's own ``cellpose.vit.CPnetBioImageIO`` that are out of sync with current
+cellpose / bioimageio.core:
+
+* ``forward``          -> unpack the 2 tensors ``CPSAM.forward`` returns (it tried 3).
+* ``load_state_dict``  -> accept the ``strict`` kwarg bioimageio.core passes in.
+* ``load_state_dict``  -> guard the legacy ``output.2.weight`` key (cpsam uses ``out.*``).
+
+The heavy base network (``CPSAM``) is imported unchanged from the installed cellpose
+package pinned in ``environment.yaml``; only this thin subclass lives here, so no
+modification of the cellpose package (nor an upstream PR) is required.
+"""
+
+import torch
+from cellpose.vit import CPSAM
+
+
+class CPnetBioImageIO(CPSAM):
+    """CP-SAM subclass compatible with the BioImage.IO Spec (loads BMZ ``cpsam`` weights)."""
+
+    def forward(self, x):
+        output_tensor, style_tensor = super().forward(x)
+        return output_tensor, style_tensor
+
+    def load_model(self, filename, device=None):
+        if (device is not None) and (device.type != "cpu"):
+            state_dict = torch.load(filename, map_location=device, weights_only=True)
+        else:
+            self.__init__(self.nout)
+            state_dict = torch.load(
+                filename, map_location=torch.device("cpu"), weights_only=True
+            )
+        self.load_state_dict(state_dict)
+
+    def load_state_dict(self, state_dict, strict=False, **kwargs):
+        if (
+            "output.2.weight" in state_dict
+            and state_dict["output.2.weight"].shape[0] != self.nout
+        ):
+            for name in self.state_dict():
+                if "output" not in name:
+                    self.state_dict()[name].copy_(state_dict[name])
+        else:
+            super().load_state_dict(
+                {name: param for name, param in state_dict.items()}, strict=False
+            )