UNETR 2 channels output binary segmentation doubt

[doctordiaz]

Hi,

First of all I would like to thank all developers their fantastic task providing MONAI package. I am using as tutorial the example spleen segmentation here https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/spleen_segmentation_3d.ipynb and using an UNETR model to train brain nifti files with masks of lesions having small signal to noise ratio. As train and validations transforms I can only use these:

train_transforms = Compose(
[
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys=["image", "label"]),
Resized(keys=["image", "label"], spatial_size=(160,160,160)),
Orientationd(keys=["image", "label"], axcodes="RAS"),
ScaleIntensityRanged(keys=["image"],a_min=-1000,a_max=1000,b_min=-1.0,b_max=1.0,clip=True)
]
as my dataset shape is 181x217x181 and I have troubles reshaping to use RandCropByPosNegLabeld

Also I have defined model:

model = UNETR(
in_channels=1,
out_channels=2,
img_size=(160,160,160),
feature_size=8,
hidden_size=512,
mlp_dim=1024,
num_heads=8,
pos_embed="perceptron",
norm_name="instance",
res_block=True,
dropout_rate=0.1,
).to(device)

and

loss_function = TverskyLoss(include_background=False, to_onehot_y=True, softmax=True)
optimizer = Novograd(model.parameters(), lr=1e-3, weight_decay=0.8)
dice_metric = DiceMetric(include_background=False, reduction="mean", get_not_nans=False)

post_pred = Compose([ AsDiscrete(argmax=True, to_onehot=2)])
post_label = Compose([Activations(sigmoid=True), AsDiscrete(to_onehot=2)])

I get a model while training gets DICE over 0.7 but when predicting it returns me values between around -3 and 2 when I should like an output prediction mask only zeros and ones. I have tryed using a Compose of [Activations(sigmoid=True),AsDiscrete(threshold=0.5)] but then DICE gets 0 while all training. I have reviewed DICE use tutorial on multiclass https://github.com/Project-MONAI/tutorials/blob/main/modules/dice_loss_metric_notes.ipynb but still doubting about how to get a binary mask using a background+foreground only one class segmentation task with MONAI but still in doubt if information in both channels has to be used together.

If you could tell me please what is missing to get this binary output I would be really grateful.

[KumoLiu]

Hi @doctordiaz, thanks for your interest here.

1. The output of the network is not passed through the activation function, you need to add your own. I think after post_pred, you can already get binary output, right?
2. Don't need to add Activations in post_label.

Thanks!

[doctordiaz]

Hi @KumoLiu, I have changed the activation to post_pred but didn't work, I took out activations from post_label, thank you

[KumoLiu]

Did you get binary output after post_pred?