SliceInferer problem - dimensions out of range

[Cerann]

Dear community,
I have a problem while using SliceInferer with already trained model of MONAI VAE class. The model is a 2D VAE trained on 256x512 size slices, extracted from a larger 3D volume. While trying to reconstruct a 'new' test 3D image (with SliceInferer) I come across an following error:

The same example with AE MONAI model works ok.

torch.Size([1, 512, 256])
{'im': tensor([[[[[0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           ...,
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.]],

          [[0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           ...,
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.]],

          [[0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           ...,
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.]],

          ...,

          [[0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           ...,
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.]],

          [[0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           ...,
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.]],

          [[0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           ...,
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.],
           [0., 0., 0.,  ..., 0., 0., 0.]]]]]), 'im_meta_dict': {'sizeof_hdr': tensor([348], dtype=torch.int32), 'extents': tensor([0], dtype=torch.int32), 'session_error': tensor([0], dtype=torch.int16), 'dim_info': tensor([0], dtype=torch.uint8), 'dim': tensor([[  3, 512, 256, 256,   1,   1,   1,   1]], dtype=torch.int16), 'intent_p1': tensor([0.]), 'intent_p2': tensor([0.]), 'intent_p3': tensor([0.]), 'intent_code': tensor([0], dtype=torch.int16), 'datatype': tensor([16], dtype=torch.int16), 'bitpix': tensor([32], dtype=torch.int16), 'slice_start': tensor([0], dtype=torch.int16), 'pixdim': tensor([[1., 1., 1., 1., 0., 0., 0., 0.]]), 'vox_offset': tensor([0.]), 'scl_slope': tensor([nan]), 'scl_inter': tensor([nan]), 'slice_end': tensor([0], dtype=torch.int16), 'slice_code': tensor([0], dtype=torch.uint8), 'xyzt_units': tensor([2], dtype=torch.uint8), 'cal_max': tensor([0.]), 'cal_min': tensor([0.]), 'slice_duration': tensor([0.]), 'toffset': tensor([0.]), 'glmax': tensor([0], dtype=torch.int32), 'glmin': tensor([0], dtype=torch.int32), 'qform_code': tensor([1], dtype=torch.int16), 'sform_code': tensor([0], dtype=torch.int16), 'quatern_b': tensor([0.]), 'quatern_c': tensor([0.]), 'quatern_d': tensor([1.]), 'qoffset_x': tensor([255.4852]), 'qoffset_y': tensor([83.7779]), 'qoffset_z': tensor([-829.9689]), 'srow_x': tensor([[0., 0., 0., 0.]]), 'srow_y': tensor([[0., 0., 0., 0.]]), 'srow_z': tensor([[0., 0., 0., 0.]]), 'affine': tensor([[[  -1.0000,    0.0000,    0.0000,  255.4852],
         [   0.0000,   -1.0000,    0.0000,   83.7779],
         [   0.0000,    0.0000,    1.0000, -829.9689],
         [   0.0000,    0.0000,    0.0000,    1.0000]]], dtype=torch.float64), 'original_affine': tensor([[[  -1.0000,    0.0000,    0.0000,  255.4852],
         [   0.0000,   -1.0000,    0.0000,   83.7779],
         [   0.0000,    0.0000,    1.0000, -829.9689],
         [   0.0000,    0.0000,    0.0000,    1.0000]]], dtype=torch.float64), 'as_closest_canonical': tensor([False]), 'spatial_shape': tensor([[512, 256, 256]], dtype=torch.int16), 'original_channel_dim': ['no_channel'], 'filename_or_obj': ['/theia/scratch/brussel/vo/000/bvo00025/vsc10215/MONAI_Autoencoder_Data/Meta_Pelvis_Pad_Res/CUSL-002/T1_c_reg.nii.gz']}, 'im_transforms': [{'class': ['EnsureTyped'], 'id': tensor([46921324750112]), 'orig_size': [tensor([512]), tensor([256]), tensor([256])]}, {'class': ['ToTensord'], 'id': tensor([46921324750304]), 'orig_size': [tensor([512]), tensor([256]), tensor([256])]}]}
Printing image name:
CUSL-002
val_data["im"].shape: 
torch.Size([1, 1, 512, 256, 256])
Traceback (most recent call last):
  File "/vscmnt/brussel_pixiu_home/_user_brussel/102/vsc10215/MONAI_VAE/reconstruct_image_VAE_2D_slice.py", line 197, in <module>
    recon = slice_inferer(val_inputs, model)
  File "/apps/brussel/CO7/zen2-ib/software/MONAI/0.9.0-foss-2021a-CUDA-11.3.1/lib/python3.9/site-packages/monai/inferers/inferer.py", line 310, in __call__
    return super().__call__(inputs=inputs, network=lambda x: self.network_wrapper(network, x, *args, **kwargs))
  File "/apps/brussel/CO7/zen2-ib/software/MONAI/0.9.0-foss-2021a-CUDA-11.3.1/lib/python3.9/site-packages/monai/inferers/inferer.py", line 192, in __call__
    return sliding_window_inference(  # type: ignore
  File "/apps/brussel/CO7/zen2-ib/software/MONAI/0.9.0-foss-2021a-CUDA-11.3.1/lib/python3.9/site-packages/monai/inferers/utils.py", line 176, in sliding_window_inference
    seg_prob_out = predictor(window_data, *args, **kwargs)  # batched patch segmentation
  File "/apps/brussel/CO7/zen2-ib/software/MONAI/0.9.0-foss-2021a-CUDA-11.3.1/lib/python3.9/site-packages/monai/inferers/inferer.py", line 310, in <lambda>
    return super().__call__(inputs=inputs, network=lambda x: self.network_wrapper(network, x, *args, **kwargs))
  File "/apps/brussel/CO7/zen2-ib/software/MONAI/0.9.0-foss-2021a-CUDA-11.3.1/lib/python3.9/site-packages/monai/inferers/inferer.py", line 336, in network_wrapper
    return tuple(out_i.unsqueeze(dim=self.spatial_dim + 2) for out_i in out)
  File "/apps/brussel/CO7/zen2-ib/software/MONAI/0.9.0-foss-2021a-CUDA-11.3.1/lib/python3.9/site-packages/monai/inferers/inferer.py", line 336, in <genexpr>
    return tuple(out_i.unsqueeze(dim=self.spatial_dim + 2) for out_i in out)
IndexError: Dimension out of range (expected to be in range of [-3, 2], but got 4)

The script I use:

`val_ds = CacheDataset(data=val_files, transform=val_transforms, cache_rate=1.0, num_workers=4)

val_loader = DataLoader(val_ds, batch_size=1, num_workers=4, collate_fn=pad_list_data_collate)

#-------------------------------------------------------------------------------------

### The in_shape has to be collapsed to 2D slice
num_samples = 1
patch_func = RandSpatialCropSamplesd(keys=["im"], roi_size=[-1, -1, 1], num_samples=num_samples, random_size=False,)

patch_transform = Compose(
    [
        SqueezeDimd(keys=["im"], dim=-1), # squeeze the last dim
        SqueezeDimd(keys=["im"], dim=0) # squeeze the first dim
    ]
)

patch_ds_train = PatchDataset(
    val_ds,
    transform = patch_transform,
    patch_func = patch_func,
    samples_per_image = num_samples,
)

print('Length of the created patch_ds_train dataset: {}'.format(len(patch_ds_train)))

patch_train_loader = DataLoader(patch_ds_train, batch_size = 1, shuffle=False, num_workers=4)

print('Length of the created patch_train_loader dataset: {}'.format(len(patch_train_loader)))

check_data = monai.utils.misc.first(patch_train_loader)
print("Slice's shape (train): ", check_data["im"].shape)

in_shape = check_data["im"].shape

#Remove extra first dimension which was 1
#in_shape = torch.squeeze(in_shape, 0)

print('NEW IN SHAPE: ')
print(in_shape)

#-------------------------------------------------------------------------------------
latent_size = 2

model = VarAutoEncoder(
    spatial_dims=2,
    in_shape=in_shape,
    out_channels=1,
    latent_size=latent_size,
    channels=(32, 64, 128, 128),
    strides=(2, 2, 2, 2),
).to(device)

#model.load_state_dict(torch.load('/user/brussel/102/vsc10215/MONAI_Autoencoder/model_full_500'))
model.load_state_dict(torch.load('/user/brussel/102/vsc10215/MONAI_VAE/model_full_VAE_2D_lat2_beta100_LR0.001_1000', map_location=torch.device('cpu')))

with torch.no_grad():
    
    for val_data in val_loader:
        print(val_data)
        
        image_name = val_data['im_meta_dict']['filename_or_obj'][0].split('/')[-2]
      
        print('Printing image name:')
        print(image_name)

        val_inputs = val_data["im"].to(device)
        
        print('val_data["im"].shape: ')
        print(val_data["im"].shape)
        
        roi_size = [-1, -1]
        sw_batch_size = 1

        slice_inferer = SliceInferer(
            roi_size=roi_size,
            sw_batch_size=sw_batch_size,
            spatial_dim=2,  # Spatial dim to slice along is defined here
            #padding_mode="replicate",
        )

        recon = slice_inferer(val_inputs, model)
      
        #{input_image_name}_{output_postfix}{output_ext}
        saver = NiftiSaver(output_dir='/theia/scratch/brussel/vo/000/bvo00025/vsc10215/MONAI_Autoencoder_Data/Reconstructed_VAE_2D_Slice', mode="bilinear", separate_folder=False, output_dtype=np.single, output_postfix=image_name)
        saver.save_batch(recon, val_data['im_meta_dict'])
        
        print('Finished')`

[wyli]

the network's output has 4 elements, so it's not directly usable with the slice inferer,

MONAI/monai/networks/nets/varautoencoder.py

Line 154 in e4a0de9

return self.decode_forward(z, self.use_sigmoid), mu, logvar, z

you can define a new model, which returns the first element of the original model:

def _model(x):
    outputs = model(x)
    return outputs[0]

recon = slice_inferer(val_inputs, _model)