Examples: Load PlantSeg from and package Cellpose to BioImage.IO model zoo

[qin-yu]

Add a notebook to show how to search, download, and load PlantSeg models, adapted from PlantSeg's https://github.com/kreshuklab/plant-seg/pull/247

Some models are using architecture from dependencies which has no version specified. E.g. plantseg.models.model.UNet3D is moved to plantseg.training.model.UNet3D, and there is no way to find out. Therefore the very last step of loading this model depend on plantseg.

[qin-yu]

Hey @FynnBe I added two example notebooks as tutorials for loading and packaging models. This PR is ready for reviews.

[mese79]

Hi @qin-yu ,
Can you also put the cellpose wrapper module beside the notebook as well?

[qin-yu]

Hey @mese79, the wrapper is not meant to be used by anyone except who creates the model package. I created the wrapper to go around the limitation of bioimageio.spec. This notebook should pass all checks:

1. torchscript runs in bioimageio.core, and
2. the weights can be put into Cellpose directly.

I don't think we should expose the wrapper at all to introduce even more complexity in an example tutorial-like notebook. Do you still need it?

[mese79]

Thanks for the reply. I'm trying to pack two cellpose models fine-tuned for two of the first AI4Life OC projects, but without the model source code(your wrapper), this notebook won't make it to the final bmz zip file, right?

So, while I'm trying to make my own wrapper, I thought having this notebook and necessary files would save some time.

[qin-yu]

Hi @mese79 I packaged this model only because I heard your team (you are AI4Life right?) needs to play with it internally, so I created the wrapper to get around #615 by ignoring the downsampled tensors. When you mentioned that you need it for actually uploading a model, I thought I'll just add them as well before giving it to you, but then realised that my wrapper should not work for the pytorch_state_dict case because of https://github.com/bioimage-io/core-bioimage-io-python/issues/396, and is potentially how I found this bug https://github.com/bioimage-io/collection/issues/79.

I don't have time to improve this PR yet, but I'll try to fix it this week. Also, I'm not part of AI4Life so I'm not involved in any group chat or meeting. Please share these three aforementioned issues if you know who can fix them.

[qin-yu]

Hi @mese79, Cellpose now supports exporting to BioImage.IO Model Zoo after PR https://github.com/MouseLand/cellpose/pull/988.

Usage:

$ 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

[mese79]

Nice! I also tried to export a Cellpose model with all postprocessing using a wrapper: https://github.com/juglab/cellpose-bmz-wrapper

[FynnBe]

@qin-yu I think we should rather add these examples to bioimageio.core as the notebooks import bioimageio.core

[qin-yu]

Nice! I also tried to export a Cellpose model with all postprocessing using a wrapper: https://github.com/juglab/cellpose-bmz-wrapper

It looks great. I did spend some time thinking what is the minimal change to make, and found that simply splitting load_model() and unwrapping the output afterwards would be the fix!

[qin-yu]

I think we should rather add these examples to bioimageio.core as the notebooks import bioimageio.core

@FynnBe Thanks for the heads up! I'll update the notebook that exports Cellpose, and I might remove the validation section. As for the PlantSeg notebook, it's based solely on the spec － if you are keeping the other existing examples maybe this one stays here?

[qin-yu]

I'm closing this because:

1. Uploading to Cellpose can be found at Cellpose docs, Cellpose code and @mese79's code
2. Downloading from PlantSeg can be found at PlantSeg code (and PlantSeg docs in a few weeks)
3. @FynnBe suggested that examples should go to bioimageio.core