Add training support for whole brain segmentation, users can use active learning in the MONAI Label

README.md

@@ -7,7 +7,8 @@ license: unknown
 ---
 # Description
 Detailed whole brain segmentation is an essential quantitative technique in medical image analysis, which provides a non-invasive way of measuring brain regions from a clinical acquired structural magnetic resonance imaging (MRI).
-We provide the pre-trained model for inferencing whole brain segmentation with 133 structures.
+We provide the pre-trained model for training and inferencing whole brain segmentation with 133 structures.
+Training pipeline is provided to support active learning in MONAI Label and training with bundle.
 
 A tutorial and release of model for whole brain segmentation using the 3D transformer-based segmentation model UNEST.
 
@@ -27,7 +28,7 @@ Fig.1 - The demonstration of T1w MRI images registered in MNI space and the whol
 
 
 # Model Overview
-A pre-trained larger UNEST base model [1] for volumetric (3D) whole brain segmentation with T1w MR images.
+A pre-trained UNEST base model [1] for volumetric (3D) whole brain segmentation with T1w MR images.
 To leverage information across embedded sequences, ”shifted window” transformers
 are proposed for dense predictions and modeling multi-scale features. However, these
 attempts that aim to complicate the self-attention range often yield high computation
@@ -97,6 +98,11 @@ Add scripts component:  To run the workflow with customized components, PYTHONPA
 export PYTHONPATH=$PYTHONPATH: '<path to the bundle root dir>/scripts'
 ```
 
+Execute Training:
+
+```
+python -m monai.bundle run training --meta_file configs/metadata.json --config_file configs/train.json --logging_file configs/logging.conf
+```
 
 Execute inference:
 
@@ -111,6 +117,12 @@ python -m monai.bundle run evaluating --meta_file configs/metadata.json --config
 Fig.3 - The output prediction comparison with variant and ground truth
 </p>
 
+## Training/Validation Benchmarking
+A graph showing the training accuracy for fine-tuning 600 epochs.
+
+![](./training.png) <br>
+
+With 10 fine-tuned labels, the training process converges fast.
 
 ## Complete ROI of the whole brain segmentation
 133 brain structures are segmented.
@@ -154,7 +166,7 @@ Fig.3 - The output prediction comparison with variant and ground truth
 
 
 ## Bundle Integration in MONAI Lable
-The inference pipleine can be easily used by the MONAI Label server and 3D Slicer for fast labeling T1w MRI images in MNI space.
+The inference and training pipleine can be easily used by the MONAI Label server and 3D Slicer for fast labeling T1w MRI images in MNI space.
 
 ![](./3DSlicer_use.png) <br>
 

configs/metadata.json

@@ -1,11 +1,12 @@
 {
     "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_20220324.json",
-    "version": "0.1.1",
+    "version": "0.1.2",
     "changelog": {
+        "0.1.2": "Add training support for whole brain segmentation, users can use active learning in the MONAI Label",
         "0.1.1": "Fix dimension according to MONAI 1.0 and fix readme file",
         "0.1.0": "complete the model package"
     },
-    "monai_version": "0.9.1",
+    "monai_version": "1.0.0",
     "pytorch_version": "1.10.0",
     "numpy_version": "1.21.2",
     "optional_packages_version": {

configs/multi_gpu_train.json

@@ -0,0 +1,36 @@
+{
+    "device": "$torch.device(f'cuda:{dist.get_rank()}')",
+    "network": {
+        "_target_": "torch.nn.parallel.DistributedDataParallel",
+        "module": "$@network_def.to(@device)",
+        "device_ids": [
+            "@device"
+        ]
+    },
+    "train#sampler": {
+        "_target_": "DistributedSampler",
+        "dataset": "@train#dataset",
+        "even_divisible": true,
+        "shuffle": true
+    },
+    "train#dataloader#sampler": "@train#sampler",
+    "train#dataloader#shuffle": false,
+    "train#trainer#train_handlers": "$@train#handlers[: -2 if dist.get_rank() > 0 else None]",
+    "validate#sampler": {
+        "_target_": "DistributedSampler",
+        "dataset": "@validate#dataset",
+        "even_divisible": false,
+        "shuffle": false
+    },
+    "validate#dataloader#sampler": "@validate#sampler",
+    "validate#evaluator#val_handlers": "$None if dist.get_rank() > 0 else @validate#handlers",
+    "training": [
+        "$import torch.distributed as dist",
+        "$dist.init_process_group(backend='nccl')",
+        "$torch.cuda.set_device(@device)",
+        "$monai.utils.set_determinism(seed=123)",
+        "$setattr(torch.backends.cudnn, 'benchmark', True)",
+        "$@train#trainer.run()",
+        "$dist.destroy_process_group()"
+    ]
+}

configs/train.json

@@ -0,0 +1,299 @@
+{
+    "imports": [
+        "$import glob",
+        "$import os",
+        "$import ignite"
+    ],
+    "bundle_root": ".",
+    "ckpt_dir": "$@bundle_root + '/models'",
+    "output_dir": "$@bundle_root + '/eval'",
+    "dataset_dir": "$@bundle_root + '/dataset/brain'",
+    "images": "$list(sorted(glob.glob(@dataset_dir + '/images/*.nii.gz')))",
+    "labels": "$list(sorted(glob.glob(@dataset_dir + '/labels/*.nii.gz')))",
+    "val_interval": 5,
+    "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')",
+    "network_def": {
+        "_target_": "scripts.networks.unest_base_patch_4.UNesT",
+        "in_channels": 1,
+        "out_channels": 133,
+        "patch_size": 4,
+        "depths": [
+            2,
+            2,
+            8
+        ],
+        "embed_dim": [
+            128,
+            256,
+            512
+        ],
+        "num_heads": [
+            4,
+            8,
+            16
+        ]
+    },
+    "network": "$@network_def.to(@device)",
+    "loss": {
+        "_target_": "DiceCELoss",
+        "to_onehot_y": true,
+        "softmax": true,
+        "squared_pred": true,
+        "batch": true
+    },
+    "optimizer": {
+        "_target_": "torch.optim.Adam",
+        "params": "$@network.parameters()",
+        "lr": 0.0001
+    },
+    "train": {
+        "deterministic_transforms": [
+            {
+                "_target_": "LoadImaged",
+                "keys": [
+                    "image",
+                    "label"
+                ]
+            },
+            {
+                "_target_": "EnsureChannelFirstd",
+                "keys": [
+                    "image",
+                    "label"
+                ]
+            },
+            {
+                "_target_": "EnsureTyped",
+                "keys": [
+                    "image",
+                    "label"
+                ]
+            }
+        ],
+        "random_transforms": [
+            {
+                "_target_": "RandSpatialCropd",
+                "keys": [
+                    "image",
+                    "label"
+                ],
+                "roi_size": [
+                    96,
+                    96,
+                    96
+                ],
+                "random_size": false
+            },
+            {
+                "_target_": "RandFlipd",
+                "keys": [
+                    "image",
+                    "label"
+                ],
+                "spatial_axis": [
+                    0
+                ],
+                "prob": 0.1
+            },
+            {
+                "_target_": "RandFlipd",
+                "keys": [
+                    "image",
+                    "label"
+                ],
+                "spatial_axis": [
+                    1
+                ],
+                "prob": 0.1
+            },
+            {
+                "_target_": "RandFlipd",
+                "keys": [
+                    "image",
+                    "label"
+                ],
+                "spatial_axis": [
+                    2
+                ],
+                "prob": 0.1
+            },
+            {
+                "_target_": "RandRotate90d",
+                "keys": [
+                    "image",
+                    "label"
+                ],
+                "max_k": 3,
+                "prob": 0.1
+            },
+            {
+                "_target_": "NormalizeIntensityd",
+                "keys": "image",
+                "nonzero": true,
+                "channel_wise": true
+            }
+        ],
+        "preprocessing": {
+            "_target_": "Compose",
+            "transforms": "$@train#deterministic_transforms + @train#random_transforms"
+        },
+        "dataset": {
+            "_target_": "CacheDataset",
+            "data": "$[{'image': i, 'label': l} for i, l in zip(@images[:-2], @labels[:-2])]",
+            "transform": "@train#preprocessing",
+            "cache_rate": 1.0,
+            "num_workers": 2
+        },
+        "dataloader": {
+            "_target_": "DataLoader",
+            "dataset": "@train#dataset",
+            "batch_size": 1,
+            "shuffle": true,
+            "num_workers": 1
+        },
+        "inferer": {
+            "_target_": "SimpleInferer"
+        },
+        "postprocessing": {
+            "_target_": "Compose",
+            "transforms": [
+                {
+                    "_target_": "Activationsd",
+                    "keys": "pred",
+                    "softmax": true
+                },
+                {
+                    "_target_": "AsDiscreted",
+                    "keys": [
+                        "pred",
+                        "label"
+                    ],
+                    "argmax": [
+                        true,
+                        false
+                    ],
+                    "to_onehot": 133
+                }
+            ]
+        },
+        "handlers": [
+            {
+                "_target_": "ValidationHandler",
+                "validator": "@validate#evaluator",
+                "epoch_level": true,
+                "interval": "@val_interval"
+            },
+            {
+                "_target_": "StatsHandler",
+                "tag_name": "train_loss",
+                "output_transform": "$monai.handlers.from_engine(['loss'], first=True)"
+            },
+            {
+                "_target_": "TensorBoardStatsHandler",
+                "log_dir": "@output_dir",
+                "tag_name": "train_loss",
+                "output_transform": "$monai.handlers.from_engine(['loss'], first=True)"
+            }
+        ],
+        "key_metric": {
+            "train_accuracy": {
+                "_target_": "ignite.metrics.Accuracy",
+                "output_transform": "$monai.handlers.from_engine(['pred', 'label'])"
+            }
+        },
+        "trainer": {
+            "_target_": "SupervisedTrainer",
+            "max_epochs": 2000,
+            "device": "@device",
+            "train_data_loader": "@train#dataloader",
+            "network": "@network",
+            "loss_function": "@loss",
+            "optimizer": "@optimizer",
+            "inferer": "@train#inferer",
+            "postprocessing": "@train#postprocessing",
+            "key_train_metric": "@train#key_metric",
+            "train_handlers": "@train#handlers",
+            "amp": true
+        }
+    },
+    "validate": {
+        "preprocessing": {
+            "_target_": "Compose",
+            "transforms": "%train#deterministic_transforms"
+        },
+        "dataset": {
+            "_target_": "CacheDataset",
+            "data": "$[{'image': i, 'label': l} for i, l in zip(@images[-2:], @labels[-2:])]",
+            "transform": "@validate#preprocessing",
+            "cache_rate": 1.0
+        },
+        "dataloader": {
+            "_target_": "DataLoader",
+            "dataset": "@validate#dataset",
+            "batch_size": 2,
+            "shuffle": false,
+            "num_workers": 1
+        },
+        "inferer": {
+            "_target_": "SlidingWindowInferer",
+            "roi_size": [
+                96,
+                96,
+                96
+            ],
+            "sw_batch_size": 4,
+            "overlap": 0.5
+        },
+        "postprocessing": "%train#postprocessing",
+        "handlers": [
+            {
+                "_target_": "StatsHandler",
+                "iteration_log": false
+            },
+            {
+                "_target_": "TensorBoardStatsHandler",
+                "log_dir": "@output_dir",
+                "iteration_log": false
+            },
+            {
+                "_target_": "CheckpointSaver",
+                "save_dir": "@ckpt_dir",
+                "save_dict": {
+                    "model": "@network"
+                },
+                "save_key_metric": true,
+                "key_metric_filename": "model.pt"
+            }
+        ],
+        "key_metric": {
+            "val_mean_dice": {
+                "_target_": "MeanDice",
+                "include_background": false,
+                "output_transform": "$monai.handlers.from_engine(['pred', 'label'])"
+            }
+        },
+        "additional_metrics": {
+            "val_accuracy": {
+                "_target_": "ignite.metrics.Accuracy",
+                "output_transform": "$monai.handlers.from_engine(['pred', 'label'])"
+            }
+        },
+        "evaluator": {
+            "_target_": "SupervisedEvaluator",
+            "device": "@device",
+            "val_data_loader": "@validate#dataloader",
+            "network": "@network",
+            "inferer": "@validate#inferer",
+            "postprocessing": "@validate#postprocessing",
+            "key_val_metric": "@validate#key_metric",
+            "additional_metrics": "@validate#additional_metrics",
+            "val_handlers": "@validate#handlers",
+            "amp": true
+        }
+    },
+    "training": [
+        "$monai.utils.set_determinism(seed=123)",
+        "$setattr(torch.backends.cudnn, 'benchmark', True)",
+        "$@train#trainer.run()"
+    ]
+}

docs/README.md

@@ -1,6 +1,7 @@
 # Description
 Detailed whole brain segmentation is an essential quantitative technique in medical image analysis, which provides a non-invasive way of measuring brain regions from a clinical acquired structural magnetic resonance imaging (MRI).
-We provide the pre-trained model for inferencing whole brain segmentation with 133 structures.
+We provide the pre-trained model for training and inferencing whole brain segmentation with 133 structures.
+Training pipeline is provided to support active learning in MONAI Label and training with bundle.
 
 A tutorial and release of model for whole brain segmentation using the 3D transformer-based segmentation model UNEST.
 
@@ -20,7 +21,7 @@ Fig.1 - The demonstration of T1w MRI images registered in MNI space and the whol
 
 
 # Model Overview
-A pre-trained larger UNEST base model [1] for volumetric (3D) whole brain segmentation with T1w MR images.
+A pre-trained UNEST base model [1] for volumetric (3D) whole brain segmentation with T1w MR images.
 To leverage information across embedded sequences, ”shifted window” transformers
 are proposed for dense predictions and modeling multi-scale features. However, these
 attempts that aim to complicate the self-attention range often yield high computation
@@ -90,6 +91,11 @@ Add scripts component:  To run the workflow with customized components, PYTHONPA
 export PYTHONPATH=$PYTHONPATH: '<path to the bundle root dir>/scripts'
 ```
 
+Execute Training:
+
+```
+python -m monai.bundle run training --meta_file configs/metadata.json --config_file configs/train.json --logging_file configs/logging.conf
+```
 
 Execute inference:
 
@@ -104,6 +110,12 @@ python -m monai.bundle run evaluating --meta_file configs/metadata.json --config
 Fig.3 - The output prediction comparison with variant and ground truth
 </p>
 
+## Training/Validation Benchmarking
+A graph showing the training accuracy for fine-tuning 600 epochs.
+
+![](./training.png) <br>
+
+With 10 fine-tuned labels, the training process converges fast.
 
 ## Complete ROI of the whole brain segmentation
 133 brain structures are segmented.
@@ -147,7 +159,7 @@ Fig.3 - The output prediction comparison with variant and ground truth
 
 
 ## Bundle Integration in MONAI Lable
-The inference pipleine can be easily used by the MONAI Label server and 3D Slicer for fast labeling T1w MRI images in MNI space.
+The inference and training pipleine can be easily used by the MONAI Label server and 3D Slicer for fast labeling T1w MRI images in MNI space.
 
 ![](./3DSlicer_use.png) <br>