update integration tests metatensor (#4501)

* update integration

Signed-off-by: Wenqi Li <[email protected]>

* update integration tests

Signed-off-by: Wenqi Li <[email protected]>

Author: wyli

tests/test_integration_classification_2d.py

@@ -33,7 +33,6 @@
     RandRotate,
     RandZoom,
     ScaleIntensity,
-    ToTensor,
     Transpose,
 )
 from monai.utils import set_determinism
@@ -69,15 +68,12 @@ def run_training_test(root_dir, train_x, train_y, val_x, val_y, device="cuda:0",
             RandRotate(range_x=np.pi / 12, prob=0.5, keep_size=True, dtype=np.float64),
             RandFlip(spatial_axis=0, prob=0.5),
             RandZoom(min_zoom=0.9, max_zoom=1.1, prob=0.5),
-            ToTensor(),
         ]
     )
     train_transforms.set_random_state(1234)
-    val_transforms = Compose(
-        [LoadImage(image_only=True), AddChannel(), Transpose(indices=[0, 2, 1]), ScaleIntensity(), ToTensor()]
-    )
-    y_pred_trans = Compose([ToTensor(), Activations(softmax=True)])
-    y_trans = Compose([ToTensor(), AsDiscrete(to_onehot=len(np.unique(train_y)))])
+    val_transforms = Compose([LoadImage(image_only=True), AddChannel(), Transpose(indices=[0, 2, 1]), ScaleIntensity()])
+    y_pred_trans = Compose([Activations(softmax=True)])
+    y_trans = AsDiscrete(to_onehot=len(np.unique(train_y)))
     auc_metric = ROCAUCMetric()
 
     # create train, val data loaders
@@ -132,7 +128,7 @@ def run_training_test(root_dir, train_x, train_y, val_x, val_y, device="cuda:0",
                 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)]
+                y = [y_trans(i) for i in decollate_batch(y, detach=False)]
                 # compute AUC
                 auc_metric(y_pred, y)
                 auc_value = auc_metric.aggregate()
@@ -153,7 +149,7 @@ def run_training_test(root_dir, train_x, train_y, val_x, val_y, device="cuda:0",
 
 def run_inference_test(root_dir, test_x, test_y, device="cuda:0", num_workers=10):
     # define transforms for image and classification
-    val_transforms = Compose([LoadImage(image_only=True), AddChannel(), ScaleIntensity(), ToTensor()])
+    val_transforms = Compose([LoadImage(image_only=True), AddChannel(), ScaleIntensity()])
     val_ds = MedNISTDataset(test_x, test_y, val_transforms)
     val_loader = DataLoader(val_ds, batch_size=300, num_workers=num_workers)
 

tests/test_integration_segmentation_3d.py

@@ -21,7 +21,7 @@
 from torch.utils.tensorboard import SummaryWriter
 
 import monai
-from monai.data import MetaTensor, create_test_image_3d, decollate_batch
+from monai.data import create_test_image_3d, decollate_batch
 from monai.inferers import sliding_window_inference
 from monai.metrics import DiceMetric
 from monai.networks import eval_mode
@@ -31,7 +31,6 @@
     AsDiscrete,
     Compose,
     EnsureChannelFirstd,
-    FromMetaTensord,
     LoadImaged,
     RandCropByPosNegLabeld,
     RandRotate90d,
@@ -40,7 +39,6 @@
     Spacingd,
 )
 from monai.utils import set_determinism
-from monai.utils.enums import PostFix
 from monai.visualize import plot_2d_or_3d_image
 from tests.testing_data.integration_answers import test_integration_value
 from tests.utils import DistTestCase, TimedCall, skip_if_quick
@@ -187,7 +185,6 @@ def run_inference_test(root_dir, device="cuda:0"):
             # resampling with align_corners=True or dtype=float64 will generate
             # slight different results between PyTorch 1.5 an 1.6
             Spacingd(keys=["img", "seg"], pixdim=[1.2, 0.8, 0.7], mode=["bilinear", "nearest"], dtype=np.float32),
-            FromMetaTensord(["img", "seg"]),
             ScaleIntensityd(keys="img"),
         ]
     )
@@ -225,11 +222,10 @@ def run_inference_test(root_dir, device="cuda:0"):
             val_outputs = sliding_window_inference(val_images, roi_size, sw_batch_size, model)
             # decollate prediction into a list
             val_outputs = [val_post_tran(i) for i in decollate_batch(val_outputs)]
-            val_meta = decollate_batch(val_data[PostFix.meta("img")])
             # compute metrics
             dice_metric(y_pred=val_outputs, y=val_labels)
-            for img, meta in zip(val_outputs, val_meta):  # save a decollated batch of files
-                saver(MetaTensor(img, meta=meta))
+            for img in val_outputs:  # save a decollated batch of files
+                saver(img)
 
     return dice_metric.aggregate().item()
 

tests/test_integration_sliding_window.py

@@ -19,7 +19,7 @@
 from ignite.engine import Engine, Events
 from torch.utils.data import DataLoader
 
-from monai.data import ImageDataset, MetaTensor, create_test_image_3d, decollate_batch
+from monai.data import ImageDataset, create_test_image_3d
 from monai.inferers import sliding_window_inference
 from monai.networks import eval_mode, predict_segmentation
 from monai.networks.nets import UNet
@@ -29,7 +29,7 @@
 
 
 def run_test(batch_size, img_name, seg_name, output_dir, device="cuda:0"):
-    ds = ImageDataset([img_name], [seg_name], transform=AddChannel(), seg_transform=AddChannel(), image_only=False)
+    ds = ImageDataset([img_name], [seg_name], transform=AddChannel(), seg_transform=AddChannel(), image_only=True)
     loader = DataLoader(ds, batch_size=1, pin_memory=torch.cuda.is_available())
 
     net = UNet(
@@ -47,9 +47,8 @@ def _sliding_window_processor(_engine, batch):
             return predict_segmentation(seg_probs)
 
     def save_func(engine):
-        meta_data = decollate_batch(engine.state.batch[2])
-        for m, o in zip(meta_data, engine.state.output):
-            saver(MetaTensor(o, meta=m))
+        for m in engine.state.output:
+            saver(m)
 
     infer_engine = Engine(_sliding_window_processor)
     infer_engine.add_event_handler(Events.ITERATION_COMPLETED, save_func)