TypeError: argument of type 'Compose' is not iterable when trying to do inverse transform

[lorenzo-consoli]

Hello,

I have some volumes on which I trained your implementation of the SwinUnetR model.
In order to feed the volumes in the model, I need to make its spatial dimensions divisible by 2**5.

After I trained my model, I need to do inference on some volumes, which I need to predict on their entirety. To make
their spatial sizes divisible by 32, I used the monai.transforms.DivisiblePadd transform. The transforms I am applying to the test data are the following:

test_transforms = Compose(
        [
            LoadImaged(keys=["image"], ensure_channel_first = True), 
            #NormalizeIntensityd(keys=["image"]),
            ScaleIntensityRanged(
                keys=["image"],
                a_min=0,
                a_max=2**16,
                b_min=0.0,
                b_max=1.0,
                clip=True,
            ),
            Orientationd(keys=["image"], axcodes="RAS"),
            EnsureTyped(keys=["image"], device=device, track_meta=True),
            DivisiblePadd(keys= ["image","label"],k=patch_size, allow_missing_keys = True)
        ]
    )

I then created the dataset and data loader as follows:

test_files = ... # loading files
test_ds = Dataset(
    data=test_files, transform=test_transforms
)
test_loader = ThreadDataLoader(test_ds, num_workers=0, batch_size=1, shuffle=True)

Then, as the images are very big, I need to patch them before feeding them into the model, and this I do using the patchify library.

My inference loop looks like this:

for batch in test_loader:
        x = batch["image"].detach().cpu()
        x = torch.squeeze(x)
        x = x.numpy()
        dataset_shape[dataset_name] = x.shape
        x = patchify(x, patch_size, step = patch_size)
        dataset_patch_dimension[dataset_name] = x.shape[:3]
        x = x.reshape(-1, 1, patch_size, patch_size, patch_size)
        current_predictions = []
        for minibatch in x:
            minibatch = torch.from_numpy(minibatch)
            minibatch = torch.reshape(minibatch, (1, 1, patch_size, patch_size, patch_size))
            minibatch = minibatch.to(device)
            logit_map = model(minibatch).detach().cpu()
            current_predictions.append(logit_map)
            del minibatch, logit_map
            gc.collect()
            torch.cuda.empty_cache()
        del x 
        gc.collect()
        torch.cuda.empty_cache()
        current_predictions = torch.stack(current_predictions,dim=0)
        models_predictions[dataset_name].append(current_predictions)

(I iterate over each patch cause loading them all on the gpu is too much for the memory I have available)

Then I try to reconstruct them in the following loop, but this is were I am getting the error:

for dataset_name, predictions in models_predictions.items():
        target_shape = dataset_shape[dataset_name]
        target_patch_shape = dataset_patch_dimension[dataset_name]
        predictions = [torch.squeeze(prediction) for prediction in predictions]
        predictions = torch.stack(predictions,dim=0)
        predictions = torch.mean(predictions, dim=0)
        predictions = post_pred(predictions)
        predictions = predictions.numpy()
        predictions = predictions.reshape(*target_patch_shape, patch_size, patch_size, patch_size)
        predictions = unpatchify(predictions,target_shape)
        predictions = torch.from_numpy(predictions)
        test_post_pred = Compose(
            [
                test_transforms.transforms[-1]
            ]
        )
        predictions_dict = {
            "label": MetaTensor(
                    predictions, 
                    applied_operations = [test_post_pred]
                )
            }
        assert isinstance(predictions_dict, dict)
        with allow_missing_keys_mode(test_transforms):
            predictions = test_post_pred.inverse(data = predictions_dict)["label"]

The traceback is the followng:

Traceback (most recent call last):
  File "/workspace/miniconda3/envs/blood-vessel-seg/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/workspace/miniconda3/envs/blood-vessel-seg/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/workspace/blood-vessel-seg/bv-seg/__main__.py", line 57, in <module>
    predict(
  File "/workspace/blood-vessel-seg/bv-seg/core/predict.py", line 226, in predict
    predictions = test_post_pred.inverse(data = predictions_dict)["label"]
  File "/workspace/miniconda3/envs/blood-vessel-seg/lib/python3.10/site-packages/monai/transforms/compose.py", line 351, in inverse
    self._raise_if_not_invertible(data)
  File "/workspace/miniconda3/envs/blood-vessel-seg/lib/python3.10/site-packages/monai/transforms/compose.py", line 373, in _raise_if_not_invertible
    invertible, reasons = has_status_keys(
  File "/workspace/miniconda3/envs/blood-vessel-seg/lib/python3.10/site-packages/monai/transforms/utils.py", line 2052, in has_status_keys
    _, reasons = has_status_keys(d, status_key, default_message)
  File "/workspace/miniconda3/envs/blood-vessel-seg/lib/python3.10/site-packages/monai/transforms/utils.py", line 2049, in has_status_keys
    status_key_occurrences.extend(check_applied_operations(op, status_key, default_message))
  File "/workspace/miniconda3/envs/blood-vessel-seg/lib/python3.10/site-packages/monai/transforms/utils.py", line 2005, in check_applied_operations
    if TraceKeys.STATUSES in entry:
TypeError: argument of type 'Compose' is not iterable

I tried to follow your notebook tutorial, and it looked to me that they were doing pretty much the same...
I guess there is something I didn't understand about how to pass the applied_operations attribute of the MetaTensor, but I am not able to figure out how....

Thank you very much for your help and time

[KumoLiu]

Hi @FMGS666, I think the reason is that test_post_pred is Compose not a dict which can not use ["label"].
I would suggest you refer to this tutorial about how to inverse.
https://github.com/Project-MONAI/tutorials/blob/main/modules/inverse_transforms_and_test_time_augmentations.ipynb
BTW, you could use sliding_window_inference to do inference with the patch.
You can refer to this tutorial.
https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/spleen_segmentation_3d.ipynb

Hope it helps, thanks!