Help with Saving Inference Output

[lenlobo]

Hi,
I have a UNet which predicts with 2 classes, either an anomaly or not. So 0 or 1.
When I run the inference and I do
test_data['pred'].shape, I get (2,300,300,300)
Is it possible to use AsDiscreted and SaveImaged functions and save only the (1,300,300,300) slice? Because when I observe in MITK, I get another 4th dim slider at the bottom which lets me switch between. But I would like to avoid this.
Thank you

[KumoLiu]

Hi @lenlobo, you can use SplitDimd after AsDiscreted and only save maybe pred_1 in SaveImaged.

MONAI/monai/transforms/utility/dictionary.py

Line 382 in c4ff70b

class SplitDimd(MapTransform, MultiSampleTrait):

I hope it helps, thanks!

[lenlobo]

@KumoLiu Thank you so much! You've been an amazing help to me :)