diff --git a/example/export_cellpose_model/bioimageio_export.py b/example/export_cellpose_model/bioimageio_export.py index 27263438..2f8edf29 100644 --- a/example/export_cellpose_model/bioimageio_export.py +++ b/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() diff --git a/example/export_cellpose_model/cellpose_original.py b/example/export_cellpose_model/cellpose_original.py index cd431446..3cec2164 100644 --- a/example/export_cellpose_model/cellpose_original.py +++ b/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) diff --git a/example/export_cellpose_model/cellpose_vit_bioimageio.py b/example/export_cellpose_model/cellpose_vit_bioimageio.py new file mode 100644 index 00000000..535180b5 --- /dev/null +++ b/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 + )