Integration of MONAI Bundle (#880)

* suppport monai bundle

Signed-off-by: Sachidanand Alle <[email protected]>

* fix for training

Signed-off-by: Sachidanand Alle <[email protected]>

* code cleanup

Signed-off-by: Sachidanand Alle <[email protected]>

* Support multi-gpu training + Fix docs

Signed-off-by: Sachidanand Alle <[email protected]>

* Fix readme

Signed-off-by: Sachidanand Alle <[email protected]>

requirements.txt

@@ -1,5 +1,5 @@
 torch>=1.7
-monai[nibabel, skimage, pillow, tensorboard, gdown, ignite, torchvision, itk, tqdm, lmdb, psutil, openslide]>=0.9.1rc4
+monai[nibabel, skimage, pillow, tensorboard, gdown, ignite, torchvision, itk, tqdm, lmdb, psutil, openslide, fire]>=0.9.1rc4
 uvicorn==0.17.6
 pydantic==1.9.1
 python-dotenv==0.20.0

sample-apps/README.md

@@ -31,14 +31,7 @@ If you are developing any examples related to pathology, you should refer this a
     - Epithelial
 
 
-### Deprecated Apps
-Following apps are deprecated/removed.  It is recommended to use [Radiology](./radiology) and [Pathology](./pathology) apps for reference.
-
-##### ~~DeepGrow~~
-
-##### ~~DeepEdit~~
-
-##### ~~Segmentation~~
-
-
+## [MONAI Bundle](./monaibundle)
 
+This app has example models to do both interactive and automated segmentation using monai-bundles defined in [MONAI ZOO](https://github.com/Project-MONAI/model-zoo/tree/dev/models).
+It can pull any bundle defined in the zoo if it is compatible and follows the checklist as defined [here](./monaibundle).

sample-apps/monaibundle/README.md

@@ -0,0 +1,58 @@
+# Integration with MONAI Bundle
+The App helps to pull any MONAI Bundle from [MONAI ZOO](https://github.com/Project-MONAI/model-zoo/tree/dev/models).
+However the following constraints has to be met for any monai bundle to directly import and use in MONAI Label.
+ - Has to meet [MONAI Bundle Specification](https://docs.monai.io/en/latest/mb_specification.html).
+ - For Inference, the bundle has defined **inference.json** or **inference.yaml** and defines [these keys](./lib/infers/bundle.py)
+ - For Training, the bundle has defined **train.json** or **train.yaml** and defines [these keys](./lib/trainers/bundle.py)
+ - For Multi-GPU Training, the bundle has defined **multi_gpu_train.json** or **multi_gpu_train.yaml**
+
+> By default models are picked from https://github.com/Project-MONAI/model-zoo/blob/dev/models/model_info.json
+
+### Structure of the App
+
+- **[lib/infers](./lib/infers)** is to define and activate inference task over monai-bundle.
+- **[lib/trainers](./lib/trainers)** is to define and activate training task over monai-bundle for single/multi gpu.
+- **[lib/activelearning](./lib/activelearning)** is the module to define the image selection techniques.
+- **[main.py](./main.py)** is the script to extend [MONAILabelApp](../../monailabel/interfaces/app.py) class
+
+> Modify Constants defined in [Infer](./lib/infers/bundle.py) and [Train](./lib/trainers/bundle.py) to customize and adopt if the basic standard/schema is not met for your bundle.
+
+### Overview
+
+
+### How To Use?
+```bash
+# skip this if you have already downloaded the app or using github repository (dev mode)
+monailabel apps --download --name monaibundle --output workspace
+
+# List all available models from zoo
+monailabel start_server --app workspace/monaibundle --studies workspace/images
+
+# Pick spleen_ct_segmentation_v0.1.0 model
+monailabel start_server --app workspace/monaibundle --studies workspace/images --conf models spleen_ct_segmentation_v0.1.0
+
+# Pick spleen_ct_segmentation_v0.1.0 model and preload
+monailabel start_server --app workspace/monaibundle --studies workspace/images --conf models spleen_ct_segmentation_v0.1.0 --conf preload true
+
+# Pick DeepEdit And Segmentation model (multiple models)
+monailabel start_server --app workspace/monaibundle --studies workspace/images --conf models "spleen_ct_segmentation_v0.1.0,spleen_deepedit_annotation_v0.1.0"
+
+# Pick All
+monailabel start_server --app workspace/monaibundle --studies workspace/images --conf models all
+
+# Pick All (Skip Training Tasks or Infer only mode)
+monailabel start_server --app workspace/monaibundle --studies workspace/images --conf models all --conf skip_trainers true
+```
+
+
+
+#### Additional Configs
+Pass them as **--conf _name_ _value_** while starting MONAILabelServer
+
+| Name          | Values          | Description                                                                                 |
+|---------------|-----------------|---------------------------------------------------------------------------------------------|
+| zoo_info      | string          | _Default value:_ https://github.com/Project-MONAI/model-zoo/blob/dev/models/model_info.json |
+| zoo_source    | string          | _Default value:_ github                                                                     |
+| zoo_repo      | string          | _Default value:_ Project-MONAI/model-zoo/hosting_storage_v1                                 |
+| preload       | true, **false** | Preload model into GPU                                                                      |
+| skip_trainers | true, **false** | Skip adding training tasks (Run in Infer mode only)                                         |

sample-apps/monaibundle/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

sample-apps/monaibundle/lib/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

sample-apps/monaibundle/lib/activelearning/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .first import First

sample-apps/monaibundle/lib/activelearning/first.py

@@ -0,0 +1,37 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+from monailabel.interfaces.datastore import Datastore
+from monailabel.interfaces.tasks.strategy import Strategy
+
+logger = logging.getLogger(__name__)
+
+
+class First(Strategy):
+    """
+    Consider implementing a first strategy for active learning
+    """
+
+    def __init__(self):
+        super().__init__("Get First Sample")
+
+    def __call__(self, request, datastore: Datastore):
+        images = datastore.get_unlabeled_images()
+        if not len(images):
+            return None
+
+        images.sort()
+        image = images[0]
+
+        logger.info(f"First: Selected Image: {image}")
+        return image

sample-apps/monaibundle/lib/infers/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .bundle import BundleInferTask

sample-apps/monaibundle/lib/infers/bundle.py

@@ -0,0 +1,122 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+import logging
+import os
+from typing import Callable, Sequence
+
+from monai.bundle import ConfigParser
+from monai.inferers import Inferer, SimpleInferer
+from monai.transforms import Compose, SaveImaged
+
+from monailabel.interfaces.tasks.infer import InferTask, InferType
+from monailabel.transform.post import Restored
+
+logger = logging.getLogger(__name__)
+
+
+class Const:
+    CONFIGS = ["inference.json", "inference.yaml"]
+    METADATA_JSON = "metadata.json"
+    MODEL_PYTORCH = "model.pt"
+    MODEL_TORCHSCRIPT = "model.ts"
+
+    KEY_DEVICE = "device"
+    KEY_BUNDLE_ROOT = "bundle_root"
+    KEY_NETWORK_DEF = "network_def"
+    KEY_PREPROCESSING = ["preprocessing", "pre_transforms"]
+    KEY_POSTPROCESSING = ["postprocessing", "post_transforms"]
+    KEY_INFERER = ["inferer"]
+
+
+class BundleInferTask(InferTask):
+    """
+    This provides Inference Engine for Monai Bundle.
+    """
+
+    def __init__(self, path, conf):
+        self.valid: bool = False
+        config_paths = [c for c in Const.CONFIGS if os.path.exists(os.path.join(path, "configs", c))]
+        if not config_paths:
+            logger.warning(f"Ignore {path} as there is no infer config {Const.CONFIGS} exists")
+            return
+
+        self.bundle_config = ConfigParser()
+        self.bundle_config.read_config(os.path.join(path, "configs", config_paths[0]))
+        self.bundle_config.config.update({Const.KEY_BUNDLE_ROOT: path})
+
+        network = None
+        model_path = os.path.join(path, "models", Const.MODEL_PYTORCH)
+        if os.path.exists(model_path):
+            network = self.bundle_config.get_parsed_content(Const.KEY_NETWORK_DEF, instantiate=True)
+        else:
+            model_path = os.path.join(path, "models", Const.MODEL_TORCHSCRIPT)
+            if not os.path.exists(model_path):
+                logger.warning(f"Ignore {path} as neither {Const.MODEL_PYTORCH} nor {Const.MODEL_TORCHSCRIPT} exists")
+                return
+
+        # https://docs.monai.io/en/latest/mb_specification.html#metadata-json-file
+        with open(os.path.join(path, "configs", Const.METADATA_JSON)) as fp:
+            metadata = json.load(fp)
+
+        self.key_image, image = next(iter(metadata["network_data_format"]["inputs"].items()))
+        self.key_pred, pred = next(iter(metadata["network_data_format"]["outputs"].items()))
+
+        labels = {v.lower(): int(k) for k, v in pred.get("channel_def", {}).items() if v.lower() != "background"}
+        description = metadata.get("description")
+        spatial_shape = image.get("spatial_shape")
+        dimension = len(spatial_shape) if spatial_shape else 3
+        type = (
+            InferType.DEEPEDIT
+            if "deepedit" in description.lower()
+            else InferType.DEEPGROW
+            if "deepgrow" in description.lower()
+            else InferType.SEGMENTATION
+        )
+
+        super().__init__(
+            path=model_path,
+            network=network,
+            type=type,
+            labels=labels,
+            dimension=dimension,
+            description=description,
+            preload=conf.get("preload", False),
+        )
+        self.valid = True
+
+    def is_valid(self) -> bool:
+        return self.valid
+
+    def pre_transforms(self, data=None) -> Sequence[Callable]:
+        pre = []
+        for k in Const.KEY_PREPROCESSING:
+            if self.bundle_config.get(k):
+                c = self.bundle_config.get_parsed_content(k, instantiate=True)
+                pre = [t for t in c.transforms] if isinstance(c, Compose) else c
+        return pre
+
+    def inferer(self, data=None) -> Inferer:
+        for k in Const.KEY_INFERER:
+            if self.bundle_config.get(k):
+                return self.bundle_config.get_parsed_content(k, instantiate=True)
+        return SimpleInferer()
+
+    def post_transforms(self, data=None) -> Sequence[Callable]:
+        post = []
+        for k in Const.KEY_POSTPROCESSING:
+            if self.bundle_config.get(k):
+                c = self.bundle_config.get_parsed_content(k, instantiate=True)
+                post = [t for t in c.transforms] if isinstance(c, Compose) else c
+
+        post = [t for t in post if not isinstance(t, SaveImaged)]
+        post.append(Restored(keys=self.key_pred, ref_image=self.key_image))
+        return post

sample-apps/monaibundle/lib/trainers/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .bundle import BundleTrainTask