Export CPnet to BioImage.IO Model Zoo

[qin-yu]

Hi @carsen-stringer!

I made a "export" module for Cellpose. Please have a look! This PR attempts to close #956 and provides an export.py module with a CLI interface to facilitate users to upload their custom Cellpose models (CPnets) to Bioimage.IO Model Zoo.

Usage

To install the dependencies:

pip install cellpose[bioimageio]

CLI interface:

$ python export.py --help
usage: export.py [-h] [--channels CHANNELS CHANNELS] --path_pretrained_model PATH_PRETRAINED_MODEL --path_readme PATH_README
                 --list_path_cover_images LIST_PATH_COVER_IMAGES [LIST_PATH_COVER_IMAGES ...] [--model_id MODEL_ID]
                 [--model_icon MODEL_ICON] --model_version MODEL_VERSION --model_name MODEL_NAME --model_documentation
                 MODEL_DOCUMENTATION --model_authors MODEL_AUTHORS --model_cite MODEL_CITE --model_tags MODEL_TAGS [MODEL_TAGS ...]
                 --model_license MODEL_LICENSE --model_repo MODEL_REPO

BioImage.IO model packaging for Cellpose

optional arguments:
  -h, --help            show this help message and exit
  --channels CHANNELS CHANNELS
                        Cyto-only = [2, 0], Cyto + Nuclei = [2, 1], Nuclei-only = [1, 0]
  --path_pretrained_model PATH_PRETRAINED_MODEL
                        Path to pretrained model file, e.g.,
                        cellpose_residual_on_style_on_concatenation_off_1135_rest_2023_05_04_23_41_31.252995
  --path_readme PATH_README
                        Path to README file
  --list_path_cover_images LIST_PATH_COVER_IMAGES [LIST_PATH_COVER_IMAGES ...]
                        List of paths to cover images
  --model_id MODEL_ID   Model ID, provide if already exists
  --model_icon MODEL_ICON
                        Model icon, provide if already exists
  --model_version MODEL_VERSION
                        Model version, new model should be 0.1.0
  --model_name MODEL_NAME
                        Model name, e.g., My Cool Cellpose
  --model_documentation MODEL_DOCUMENTATION
                        Model documentation, e.g., A cool Cellpose model trained for my cool dataset.
  --model_authors MODEL_AUTHORS
                        Model authors in JSON format, e.g., '[{"name": "Qin Yu", "affiliation": "EMBL", "github_user": "qin-yu", "orcid":
                        "0000-0002-4652-0795"}]'
  --model_cite MODEL_CITE
                        Model citation in JSON format, e.g., '[{"text": "For more details of the model itself, see the manuscript",
                        "doi": "10.1101/2024.02.19.580954", "url": null}]'
  --model_tags MODEL_TAGS [MODEL_TAGS ...]
                        Model tags, e.g., cellpose 3d 2d
  --model_license MODEL_LICENSE
                        Model license, e.g., MIT
  --model_repo MODEL_REPO
                        Model repository URL

This creates a bioimage.io package in a new directory models/MODEL_NAME/cellpose_model.zip which can be uploaded via https://bioimageio-uploader.netlify.app/#/uploader/add

Details

• To make CPnet class pass tests from bioimageio.core, one small refactor of .load_model() method was made, which only includes formal changes but keeps the semantics.
• The new class CPnetBioImageIO unpacks the nested output from CPnet and make it compatible to bioimageio.spec
• The standard Cellpose image at http://www.cellpose.org/static/data/rgb_3D.tif is used as test_input for all Cellpose models

[carsen-stringer]

thank you!

[carsen-stringer]

Hi @qin-yu, could you please write a short blurb in our readthedocs describing how to export a model to bioimage.io? I think it would make sense here: https://github.com/MouseLand/cellpose/blob/main/docs/models.rst#user-trained-models.

I can also do this but I think you're the expert here, thanks so much for adding this to cellpose

[qin-yu]

Hi @qin-yu, could you please write a short blurb in our readthedocs describing how to export a model to bioimage.io?

Sure! I'll let you know once it's done.

[carsen-stringer]

Hi @qin-yu , can we move the overwriting of load_state_dict into your subclass (

cellpose/cellpose/resnet_torch.py

Line 279 in 729b701

def load_state_dict(self, state_dict):

)? Or some other refactor because currently load_state_dict fails for our models. Thanks for your help!

[qin-yu]

Hi @qin-yu , can we move the overwriting of load_state_dict into your subclass (

cellpose/cellpose/resnet_torch.py

Line 279 in 729b701

def load_state_dict(self, state_dict):

)? Or some other refactor because currently load_state_dict fails for our models. Thanks for your help!

If we restore the old load_model() for CPnet and copy-paste both load_model() and load_state_dict() into CPnetBioImageIO, then your code shouldn't fail any more, and my export.py module should still work because it only uses CPnetBioImageIO class. I haven't tested this approach yet, but it's my initial thought based on the current situation.