Fixes for Train+Infer (#1156)

* Remove deepedit nuclei and fix validation dataloader transforms

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

* Fix infer decollate

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

* remove deepedit from pathology

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

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

Author: SachidanandAlle

monailabel/tasks/infer/basic_infer.py

@@ -477,8 +477,11 @@ def run_inferer(self, data: Dict[str, Any], convert_to_batch=True, device="cuda"
                 torch.cuda.empty_cache()
 
             if convert_to_batch:
-                outputs_d = decollate_batch(outputs)
-                outputs = outputs_d[0]
+                if data.get("decollate_batch", False):  # TODO:: Use automatically depending on the input/output
+                    outputs_d = decollate_batch(outputs)
+                    outputs = outputs_d[0]
+                else:
+                    outputs = outputs[0]
 
             if isinstance(outputs, dict):
                 data.update(outputs)

monailabel/tasks/train/basic_train.py

@@ -198,15 +198,19 @@ def lr_scheduler_handler(self, context: Context):
         lr_scheduler = torch.optim.lr_scheduler.StepLR(context.optimizer, step_size=1000, gamma=0.1)
         return LrScheduleHandler(lr_scheduler, print_lr=True)
 
-    def _dataset(self, context, datalist, replace_rate=0.25):
+    def _dataset(self, context, datalist, is_train, replace_rate=0.25):
         if context.multi_gpu:
             world_size = torch.distributed.get_world_size()
             if len(datalist) // world_size:  # every gpu gets full data when datalist is smaller
                 datalist = partition_dataset(data=datalist, num_partitions=world_size, even_divisible=True)[
                     context.local_rank
                 ]
 
-        transforms = self._validate_transforms(self.train_pre_transforms(context), "Training", "pre")
+        transforms = (
+            self._validate_transforms(self.train_pre_transforms(context), "Training", "pre")
+            if is_train
+            else self._validate_transforms(self.val_pre_transforms(context), "Validation", "pre")
+        )
         dataset = (
             CacheDataset(datalist, transforms)
             if context.dataset_type == "CacheDataset"
@@ -234,8 +238,8 @@ def _dataloader(self, context, dataset, batch_size, num_workers, shuffle=False):
             num_workers=num_workers,
         )
 
-    def train_data_loader(self, context, num_workers=0, shuffle=False):
-        dataset, datalist = self._dataset(context, context.train_datalist)
+    def train_data_loader(self, context, num_workers=0, shuffle=True):
+        dataset, datalist = self._dataset(context, context.train_datalist, is_train=True)
         logger.info(f"{context.local_rank} - Records for Training: {len(datalist)}")
         logger.debug(f"{context.local_rank} - Training: {datalist}")
 
@@ -288,7 +292,7 @@ def train_additional_metrics(self, context: Context):
         return None
 
     def val_data_loader(self, context: Context, num_workers=0):
-        dataset, datalist = self._dataset(context, context.val_datalist)
+        dataset, datalist = self._dataset(context, context.val_datalist, is_train=False)
         logger.info(f"{context.local_rank} - Records for Validation: {len(datalist)}")
         logger.debug(f"{context.local_rank} - Validation: {datalist}")
 

runtests.sh

@@ -468,7 +468,7 @@ function run_integration_tests() {
 # network training/inference/eval integration tests
 if [ $doNetTests = true ]; then
   run_integration_tests "radiology" "tests/data/dataset/local/spleen" "deepedit,segmentation_spleen,segmentation,deepgrow_2d,deepgrow_3d"
-  run_integration_tests "pathology" "tests/data/pathology" "deepedit_nuclei,segmentation_nuclei,nuclick"
+  run_integration_tests "pathology" "tests/data/pathology" "segmentation_nuclei,nuclick"
   run_integration_tests "monaibundle" "tests/data/dataset/local/spleen" "spleen_ct_segmentation_v0.1.0,spleen_deepedit_annotation_v0.1.0,swin_unetr_btcv_segmentation_v0.1.0"
   run_integration_tests "endoscopy" "tests/data/endoscopy" "tooltracking,inbody,deepedit"
 fi

sample-apps/pathology/README.md

@@ -24,15 +24,16 @@ weights/model (UNET).
     - Connective/Soft tissue cells
     - Dead Cells
     - Epithelial
-- **DeepEdit Nuclei** - It is a combination of
-  both [Interaction + Auto Segmentation](https://github.com/Project-MONAI/MONAILabel/wiki/DeepEdit) model which is
-  trained to segment Nuclei cells that combines all above labels as *Nuclei*.
-- **NuClick** - This is NuClick implementation (UNet model) as provided at: https://github.com/mostafajahanifar/nuclick_torch.  Training task for monailabel is not yet supported.
+- **NuClick** - This is NuClick implementation (UNet model) as provided at: https://github.com/mostafajahanifar/nuclick_torch.
+- **Classification Nuclei** - It is a simple classification model which can be used along with NuClick model to classify Nuclei cells.
 
 ### Dataset
 
-Above _Nuclei_ models are trained
-on [PanNuke Dataset for Nuclei Instance Segmentation and Classification](https://warwick.ac.uk/fac/cross_fac/tia/data/pannuke)
+Above _Nuclei_ models are trained either on
+ - [PanNuke Dataset for Nuclei Instance Segmentation and Classification](https://warwick.ac.uk/fac/cross_fac/tia/data/pannuke)
+ - [CoNSeP Dataset](https://warwick.ac.uk/fac/cross_fac/tia/data/hovernet)
+
+Pass `--conf consep true` option while starting MONAILabel server to use models trained on CoNSeP Dataset
 
 ### Inputs
 
@@ -86,7 +87,7 @@ The current version of plugin comes with **limited features** to support basic a
 - Make sure MONAILabel Server URL is correctly through `Preferences`.
 - Open Sample Whole Slide Image in QuPath (which is shared as studies for MONAILabel server)
 - Add/Select Rectangle ROI to run annotations using MONAI Label models.
-- For Interative model (e.g. DeepEdit) you can choose to provide `Positive` and `Negative` points through Annotation panel.
+- For Interactive model you can choose to provide `Positive` and `Negative` points through Annotation panel.
 
 ![image](../../docs/images/qupath.jpg)
 

sample-apps/pathology/lib/configs/deepedit_nuclei.py

@@ -10,6 +10,5 @@
 # limitations under the License.
 
 from .classification_nuclei import ClassificationNuclei
-from .deepedit_nuclei import DeepEditNuclei
 from .nuclick import NuClick
 from .segmentation_nuclei import SegmentationNuclei

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

@@ -10,6 +10,5 @@
 # limitations under the License.
 
 from .classification_nuclei import ClassificationNuclei
-from .deepedit_nuclei import DeepEditNuclei
 from .nuclick import NuClick
 from .segmentation_nuclei import SegmentationNuclei