Repository: Project-MONAI/MONAI

Issue:
Segmentation / Classification integration tests forget to use decollate
Now, post processing transforms assume the input to be channel first tensors, but in:
https://github.com/Project-MONAI/MONAI/blob/dev/tests/test_integration_segmentation_3d.py#L149
It used batch first tensors.


Discussion:

Resolving patch:
diff --git a/monai/visualize/img2tensorboard.py b/monai/visualize/img2tensorboard.py
--- a/monai/visualize/img2tensorboard.py
+++ b/monai/visualize/img2tensorboard.py
@@ -9,11 +9,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Dict, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Union
 
 import numpy as np
 import torch
 
+from monai.config import NdarrayTensor
 from monai.transforms import rescale_array
 from monai.utils import optional_import
 
@@ -159,7 +160,7 @@ def add_animated_gif_no_channels(
 
 
 def plot_2d_or_3d_image(
-    data: Union[torch.Tensor, np.ndarray],
+    data: Union[NdarrayTensor, List[NdarrayTensor]],
     step: int,
     writer: SummaryWriter,
     index: int = 0,
@@ -174,7 +175,8 @@ def plot_2d_or_3d_image(
 
     Args:
         data: target data to be plotted as image on the TensorBoard.
-            The data is expected to have 'NCHW[D]' dimensions, and only plot the first in the batch.
+            The data is expected to have 'NCHW[D]' dimensions or a list of data with `CHW[D]` dimensions,
+            and only plot the first in the batch.
         step: current step to plot in a chart.
         writer: specify TensorBoard SummaryWriter to plot the image.
         index: plot which element in the input data batch, default is the first element.
@@ -182,7 +184,8 @@ def plot_2d_or_3d_image(
         max_frames: number of frames for 2D-t plot.
         tag: tag of the plotted image on TensorBoard.
     """
-    d = data[index].detach().cpu().numpy() if isinstance(data, torch.Tensor) else data[index]
+    data_index = data[index]
+    d: np.ndarray = data_index.detach().cpu().numpy() if isinstance(data_index, torch.Tensor) else data_index
 
     if d.ndim == 2:
         d = rescale_array(d, 0, 1)


Test patch:
diff --git a/tests/test_compute_roc_auc.py b/tests/test_compute_roc_auc.py
--- a/tests/test_compute_roc_auc.py
+++ b/tests/test_compute_roc_auc.py
@@ -15,8 +15,9 @@
 import torch
 from parameterized import parameterized
 
+from monai.data import decollate_batch
 from monai.metrics import ROCAUCMetric, compute_roc_auc
-from monai.transforms import Activations, AsDiscrete
+from monai.transforms import Activations, AsDiscrete, Compose, ToTensor
 
 TEST_CASE_1 = [
     torch.tensor([[0.1, 0.9], [0.3, 1.4], [0.2, 0.1], [0.1, 0.5]]),
@@ -85,22 +86,23 @@
 class TestComputeROCAUC(unittest.TestCase):
     @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6, TEST_CASE_7])
     def test_value(self, y_pred, y, softmax, to_onehot, average, expected_value):
-        act = Activations(softmax=softmax)
-        dis = AsDiscrete(to_onehot=to_onehot, n_classes=2)
-        y_pred = torch.stack([act(i) for i in y_pred], dim=0)
-        y = torch.stack([dis(i) for i in y], dim=0)
+        y_pred_trans = Compose([ToTensor(), Activations(softmax=softmax)])
+        y_trans = Compose([ToTensor(), AsDiscrete(to_onehot=to_onehot, n_classes=2)])
+        y_pred = torch.stack([y_pred_trans(i) for i in decollate_batch(y_pred)], dim=0)
+        y = torch.stack([y_trans(i) for i in decollate_batch(y)], dim=0)
         result = compute_roc_auc(y_pred=y_pred, y=y, average=average)
         np.testing.assert_allclose(expected_value, result, rtol=1e-5)
 
     @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6, TEST_CASE_7])
     def test_class_value(self, y_pred, y, softmax, to_onehot, average, expected_value):
-        act = Activations(softmax=softmax)
-        dis = AsDiscrete(to_onehot=to_onehot, n_classes=2)
-        y_pred = torch.stack([act(i) for i in y_pred], dim=0)
-        y = torch.stack([dis(i) for i in y], dim=0)
+        y_pred_trans = Compose([ToTensor(), Activations(softmax=softmax)])
+        y_trans = Compose([ToTensor(), AsDiscrete(to_onehot=to_onehot, n_classes=2)])
+        y_pred = [y_pred_trans(i) for i in decollate_batch(y_pred)]
+        y = [y_trans(i) for i in decollate_batch(y)]
         metric = ROCAUCMetric(average=average)
         metric(y_pred=y_pred, y=y)
         result = metric.aggregate()
+        metric.reset()
         np.testing.assert_allclose(expected_value, result, rtol=1e-5)
 
 
diff --git a/tests/test_integration_classification_2d.py b/tests/test_integration_classification_2d.py
--- a/tests/test_integration_classification_2d.py
+++ b/tests/test_integration_classification_2d.py
@@ -20,7 +20,8 @@
 
 import monai
 from monai.apps import download_and_extract
-from monai.metrics import compute_roc_auc
+from monai.data import decollate_batch
+from monai.metrics import ROCAUCMetric
 from monai.networks import eval_mode
 from monai.networks.nets import DenseNet121
 from monai.transforms import (
@@ -74,8 +75,9 @@ def run_training_test(root_dir, train_x, train_y, val_x, val_y, device="cuda:0",
     )
     train_transforms.set_random_state(1234)
     val_transforms = Compose([LoadImage(image_only=True), AddChannel(), ScaleIntensity(), ToTensor()])
-    act = Activations(softmax=True)
-    to_onehot = AsDiscrete(to_onehot=True, n_classes=len(np.unique(train_y)))
+    y_pred_trans = Compose([ToTensor(), Activations(softmax=True)])
+    y_trans = Compose([ToTensor(), AsDiscrete(to_onehot=True, n_classes=len(np.unique(train_y)))])
+    auc_metric = ROCAUCMetric()
 
     # create train, val data loaders
     train_ds = MedNISTDataset(train_x, train_y, train_transforms)
@@ -127,18 +129,21 @@ def run_training_test(root_dir, train_x, train_y, val_x, val_y, device="cuda:0",
                 # compute accuracy
                 acc_value = torch.eq(y_pred.argmax(dim=1), y)
                 acc_metric = acc_value.sum().item() / len(acc_value)
+                # decollate prediction and label and execute post processing
+                y_pred = [y_pred_trans(i) for i in decollate_batch(y_pred)]
+                y = [y_trans(i) for i in decollate_batch(y)]
                 # compute AUC
-                y_pred = torch.stack([act(i) for i in y_pred])
-                y = torch.stack([to_onehot(i) for i in y])
-                auc_metric = compute_roc_auc(y_pred, y)
-                metric_values.append(auc_metric)
-                if auc_metric > best_metric:
-                    best_metric = auc_metric
+                auc_metric(y_pred, y)
+                auc_value = auc_metric.aggregate()
+                auc_metric.reset()
+                metric_values.append(auc_value)
+                if auc_value > best_metric:
+                    best_metric = auc_value
                     best_metric_epoch = epoch + 1
                     torch.save(model.state_dict(), model_filename)
                     print("saved new best metric model")
                 print(
-                    f"current epoch {epoch +1} current AUC: {auc_metric:0.4f} "
+                    f"current epoch {epoch +1} current AUC: {auc_value:0.4f} "
                     f"current accuracy: {acc_metric:0.4f} best AUC: {best_metric:0.4f} at epoch {best_metric_epoch}"
                 )
     print(f"train completed, best_metric: {best_metric:0.4f}  at epoch: {best_metric_epoch}")
diff --git a/tests/test_integration_segmentation_3d.py b/tests/test_integration_segmentation_3d.py
--- a/tests/test_integration_segmentation_3d.py
+++ b/tests/test_integration_segmentation_3d.py
@@ -21,7 +21,7 @@
 from torch.utils.tensorboard import SummaryWriter
 
 import monai
-from monai.data import NiftiSaver, create_test_image_3d
+from monai.data import NiftiSaver, create_test_image_3d, decollate_batch
 from monai.inferers import sliding_window_inference
 from monai.metrics import DiceMetric
 from monai.networks import eval_mode
@@ -36,6 +36,7 @@
     RandRotate90d,
     ScaleIntensityd,
     Spacingd,
+    ToTensor,
     ToTensord,
 )
 from monai.utils import set_determinism
@@ -94,7 +95,7 @@ def run_training_test(root_dir, device="cuda:0", cachedataset=0):
     # create a validation data loader
     val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
     val_loader = monai.data.DataLoader(val_ds, batch_size=1, num_workers=4)
-    val_post_tran = Compose([Activations(sigmoid=True), AsDiscrete(threshold_values=True)])
+    val_post_tran = Compose([ToTensor(), Activations(sigmoid=True), AsDiscrete(threshold_values=True)])
     dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
 
     # create UNet, DiceLoss and Adam optimizer
@@ -146,7 +147,10 @@ def run_training_test(root_dir, device="cuda:0", cachedataset=0):
                 for val_data in val_loader:
                     val_images, val_labels = val_data["img"].to(device), val_data["seg"].to(device)
                     sw_batch_size, roi_size = 4, (96, 96, 96)
-                    val_outputs = val_post_tran(sliding_window_inference(val_images, roi_size, sw_batch_size, model))
+                    val_outputs = sliding_window_inference(val_images, roi_size, sw_batch_size, model)
+                    # decollate prediction into a list and execute post processing for every item
+                    val_outputs = [val_post_tran(i) for i in decollate_batch(val_outputs)]
+                    # compute metrics
                     dice_metric(y_pred=val_outputs, y=val_labels)
 
                 metric = dice_metric.aggregate().item()
@@ -191,7 +195,7 @@ def run_inference_test(root_dir, device="cuda:0"):
     val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
     # sliding window inference need to input 1 image in every iteration
     val_loader = monai.data.DataLoader(val_ds, batch_size=1, num_workers=4)
-    val_post_tran = Compose([Activations(sigmoid=True), AsDiscrete(threshold_values=True)])
+    val_post_tran = Compose([ToTensor(), Activations(sigmoid=True), AsDiscrete(threshold_values=True)])
     dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
 
     model = UNet(
@@ -213,7 +217,10 @@ def run_inference_test(root_dir, device="cuda:0"):
             val_images, val_labels = val_data["img"].to(device), val_data["seg"].to(device)
             # define sliding window size and batch size for windows inference
             sw_batch_size, roi_size = 4, (96, 96, 96)
-            val_outputs = val_post_tran(sliding_window_inference(val_images, roi_size, sw_batch_size, model))
+            val_outputs = sliding_window_inference(val_images, roi_size, sw_batch_size, model)
+            # decollate prediction into a list and execute post processing for every item
+            val_outputs = [val_post_tran(i) for i in decollate_batch(val_outputs)]
+            # compute metrics
             dice_metric(y_pred=val_outputs, y=val_labels)
             saver.save_batch(val_outputs, val_data["img_meta_dict"])
 
diff --git a/tests/test_integration_workflows.py b/tests/test_integration_workflows.py
--- a/tests/test_integration_workflows.py
+++ b/tests/test_integration_workflows.py
@@ -112,6 +112,7 @@ def run_training_test(root_dir, device="cuda:0", amp=False, num_workers=4):
 
     val_postprocessing = Compose(
         [
+            ToTensord(keys=["pred", "label"]),
             Activationsd(keys="pred", sigmoid=True),
             AsDiscreted(keys="pred", threshold_values=True),
             KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]),
@@ -152,6 +153,7 @@ def _forward_completed(self, engine):
 
     train_postprocessing = Compose(
         [
+            ToTensord(keys=["pred", "label"]),
             Activationsd(keys="pred", sigmoid=True),
             AsDiscreted(keys="pred", threshold_values=True),
             KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]),
@@ -237,6 +239,7 @@ def run_inference_test(root_dir, model_file, device="cuda:0", amp=False, num_wor
 
     val_postprocessing = Compose(
         [
+            ToTensord(keys=["pred", "label"]),
             Activationsd(keys="pred", sigmoid=True),
             AsDiscreted(keys="pred", threshold_values=True),
             KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]),