Initial commit: Hugging Face Hub MONAI integration demo

app.py

@@ -0,0 +1,226 @@
+from glob import glob
+import logging
+from matplotlib import pyplot as plt
+import os
+
+from monai.bundle.scripts import upload_zoo_bundle_to_hf
+import streamlit as st
+import torch
+
+
+def main():
+    st.title("MONAI 🤗 Hugging Face Integration")
+
+    st.write("""\
+Here's a demo of a prototype integration between
+[MONAI](https://monai.io/) and the [Hugging Face
+Hub](https://huggingface.co/docs/hub/index), which allows for
+uploading models to the Hub and downloading them.  The integration
+itself is implemented in [this
+branch](https://github.com/dnouri/MONAI/tree/dnouri/huggingface-support)
+of MONAI.
+""")
+
+    st.write("""\
+## Uploading models to the Hub ⬆
+
+The new `upload_zoo_bundle_to_hf` command allows us to upload models
+from the existing [MONAI Model
+Zoo](https://github.com/Project-MONAI/model-zoo) on Github directly
+onto the Hugging Face Hub.
+
+The `--name` option specifies the [filename of an existing
+model](https://github.com/Project-MONAI/model-zoo/releases/tag/hosting_storage_v1)
+in the MONAI Model Zoo, while the `--hf_organization` specifies the
+name of the organization to upload to, in the Hugging Face Hub,
+whereas `--hf_token` is the [HF user access
+token](https://huggingface.co/docs/hub/security-tokens).
+
+An additional `--hf_card_data` option allows us to specify [model card
+metadata](https://huggingface.co/docs/hub/models-cards#model-card-metadata)
+to be added to the Hugging Face model card.
+
+An example call to the `upload_zoo_bundle_to_hf` script looks like
+this:
+
+```bash
+  python -m monai.bundle upload_zoo_bundle_to_hf \\
+      --name spleen_ct_segmentation_v0.1.0 \\
+      --hf_organization dnouri --hf_token mytoken \\
+      --hf_card_data '{"lang": "en"}'
+```
+
+An example of a thus automatically uploaded model can be found
+[here](https://huggingface.co/dnouri/spleen_ct_segmentation).
+
+### Try it out!
+
+To try out uploading your own model, please provide the information below:
+""")
+    filename = st.text_input("Filename of MONAI Model Zoo model "
+                             "(e.g. ventricular_short_axis_3label_v0.1.0.zip)")
+    username = st.text_input("Hub organization or user name (e.g. dnouri)")
+    card_data = st.text_input("Optional model card metadata",
+                              value='{"tags": ["MONAI"]}')
+    token = st.text_input("Hugging Face user access token")
+
+    if filename and username and token:
+        st.write("Please wait...")
+        upload_zoo_bundle_to_hf(
+            name=filename,
+            hf_organization=username,
+            hf_token=token,
+            hf_card_data=card_data or None,
+            )
+        st.write(f"""\
+Done!  You should be able to find the [result here](https://huggingface.co/{username}/{filename.rsplit("_", 1)[0]}).
+""")
+
+    st.write("""\
+## Downloading models from the Hub ⬇
+
+Uploading isn't much fun if you can't also download the models from
+the Hub!  To help with that, we've added support for the Hugging Face
+Hub to the existing MONAI bundle `download` command.
+
+The `download` command's default `--source` is `github`.  We'll choose
+`huggingface` instead to download from the Hub.
+
+The `--name` of the model is the name of your model on the Hub,
+e.g. `ventricular_short_axis_3label`.  Note that as per MONAI
+convention, we do not specify the version name here.  (Future versions of
+this command might allow for downloading specific versions, or tags.)
+
+The `--repo` normally points to the MONAI Model Zoo's ['hosting
+storage' release page on
+Github](https://github.com/Project-MONAI/model-zoo/releases/tag/hosting_storage_v1).
+When we call `download` with the `huggingface` source, we'll require
+the `--repo` argument to point to the organization or user name that
+hosts the model, e.g. `dnouri`.  (While this choice is a bit
+confusing, it also reflects an attempt to pragmatically blend concepts
+from both MONAI bundles and the Hub.  Future versions might improve on
+this.)
+
+An example call to the `upload_zoo_bundle_to_hf` script that perhaps
+downloads the model that we uploaded previously, looks like this:
+
+```bash
+  python -m monai.bundle download \\
+      --name spleen_ct_segmentation \\
+      --source huggingface --repo dnouri
+```
+""")
+
+    st.write("""\
+## Use model for inference 🧠
+
+To use the `spleen_ct_segmentation` pretrained model to do inference,
+we'll first load it into memory (as a TorchScript module) using the
+`load` function below.  This will download the model from the Hugging
+Face Hub, as `load` uses the aforementioned `download` under the hood:
+""")
+    # The next line is a workaround against a buggy interaction
+    # between how streamlit sets up stderr and how tqdm uses it:
+    logging.getLogger().setLevel(logging.NOTSET)
+
+    with st.echo():
+        from monai.bundle.scripts import load
+
+        model, metadata, extra = load(
+            name="spleen_ct_segmentation",
+            source="huggingface",
+            repo="dnouri",
+            load_ts_module=True,
+            progress=False,
+            )
+
+    st.write("""\
+This will produce a model, but we'll also need the corresponding
+transforms.  These are defined in the MONAI bundle configuration
+files.  There's unfortunately not a convenient way to do this using a
+MONAI bundle script function, so we'll have to reach into the MONAI
+bowels for a bit:
+""")
+    with st.echo():
+        from monai.bundle.config_parser import ConfigParser
+        from monai.bundle.scripts import _process_bundle_dir
+
+        model_dir = _process_bundle_dir() / "spleen_ct_segmentation"
+        config_paths = [
+            model_dir / "configs" / "train.json",
+            model_dir / "configs" / "evaluate.json",
+            ]
+        config = ConfigParser(
+            ConfigParser.load_config_files(config_paths),
+            )
+        preprocess = config.get_parsed_content("validate#preprocessing")
+
+    st.write("""\
+We'll borrow code from the MONAI [Spleen 3D segmentation with MONAI
+tutorial](https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/spleen_segmentation_3d.ipynb)
+to download the data that our `spleen_ct_segmentation` model was
+trained with:
+""")
+
+    with st.echo():
+        from monai.apps import download_and_extract
+
+        root_dir = os.environ.get(
+            "MONAI_DATA_DIRECTORY",
+            os.path.expanduser("~/.cache/monai_data_directory")
+            )
+        os.makedirs(root_dir, exist_ok=True)
+        resource = "https://msd-for-monai.s3-us-west-2.amazonaws.com/Task09_Spleen.tar"
+        md5 = "410d4a301da4e5b2f6f86ec3ddba524e"
+        compressed_file = os.path.join(root_dir, "Task09_Spleen.tar")
+        data_dir = os.path.join(root_dir, "Task09_Spleen")
+        if not os.path.exists(data_dir):
+            download_and_extract(resource, compressed_file, root_dir, md5)
+
+        train_images = sorted(
+            glob(os.path.join(data_dir, "imagesTr", "*.nii.gz")))
+        train_labels = sorted(
+            glob(os.path.join(data_dir, "labelsTr", "*.nii.gz")))
+        data_dicts = [
+            {"image": image_name, "label": label_name}
+            for image_name, label_name in zip(train_images, train_labels)
+        ]
+        files = data_dicts
+        st.write(f"Downloaded {len(files)} files.")
+
+    st.write("""\
+Finally, we can run inference and plot some results: 🥳
+""")
+    image_idx = st.slider("Image number", 0, len(files))
+
+    with st.echo():
+        from monai.inferers import sliding_window_inference
+
+        data = preprocess(files[image_idx])
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        output = sliding_window_inference(
+            inputs=data["image"].to(device)[None, ...],
+            roi_size=(160, 160, 160),
+            sw_batch_size=4,
+            predictor=model.eval(),
+            )
+
+        fig, (ax1, ax2, ax3) = plt.subplots(1, 3)
+        ax1.set_title("Image")
+        ax2.set_title("Label")
+        ax3.set_title("Output")
+        ax1.imshow(data["image"][0, :, :, 80], cmap="gray")
+        ax2.imshow(data["label"][0, :, :, 80], cmap="gray")
+        output_img = (
+            torch.argmax(output, dim=1)[0, :, :, 80]
+            .cpu().detach().numpy()
+            )
+        ax3.imshow(
+            output_img,
+            cmap="gray",
+            )
+        st.pyplot(fig)
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,9 @@
+git+https://github.com/dnouri/MONAI.git@dnouri/huggingface-support
+dicom
+itk
+huggingface_hub
+modelcards
+nibabel
+pynrrd
+streamlit
+# trigger rebuild of cached layer with some random message :-)