Save Nifti file during inferrence

[mbrzus]

Hello, I am using Ubuntu 18.04 with
torch=1.7.1
pytorch-lightning=1.2.1
monai=0.4.0

I trained a model using monai 3DResUNet for brain extraction and now I want to use the trained model for inferrence.
The problem is that the metadata from the trained model is in the wrong format and is not accepted by the NiftySaver. How can I save the images including the meta data?

I also tested just the transforms I am using (without training the model). After running transforms the meta data dictionary works fine with NiftySaver. I tried to manually change types and convert the data I got from my model to the format that I got from running just the transforms but it was also failing. I can provide additional files if needed (model, the meta data dictionaries and error messages)
Here is my inference code:

saver = NiftiSaver(output_dir=out_dir, dtype=None)
model = BrainExtraction.load_from_checkpoint(
checkpoint_path=f"{checkpoints_dir}/epoch=499-val_loss=0.04-val_dice=0.00.ckpt",
hparams_file=f"{params_dir}/hparams.yaml"
)
trainer = pl.Trainer(gpus=[0])

model.eval()
device = torch.device("cuda:0")
model.to(device)
inferer = SimpleInferer()
with torch.no_grad():
for i, test_data in enumerate(model.test_dataloader()):
test_outputs = inferer(inputs=test_data["image"].cuda(), network=model)
image = torch.argmax(test_outputs, dim=1).detach().cpu()
meta_data = test_data["image_meta_dict"]
saver.save(data=image, meta_data=new_dict)

[Nic-Ma]

Hi @mbrzus ,

Thanks for your interest here and sorry for the delayed response.
If you got data from DataLoader, it has the batch dim, so I think you should can saver.save_batch() instead of saver.save().
Refer to: https://github.com/Project-MONAI/tutorials/blob/master/3d_segmentation/torch/unet_evaluation_dict.py#L96

Thanks.