3D segmentation configuration

[LSnyd]

Hi there,

I am trying to get a binary 3D segmentation running based on the two examples of the tutorials, but it seems like there are still some problems in my setup. I hope that someone might be able to help me.

1. I am performing the following transformations on my data:

train_transforms = Compose(
            [
                LoadImaged(keys=("image", "label")),
                Clipping(keys=("label")),
                EnsureChannelFirstd(keys=["image", "label"]),
                Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 2.0),mode=("bilinear", "nearest")),
                NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                CropForegroundd(keys=("image", "label"), source_key="image"),
                Resized(keys=("image", "label"), spatial_size=(160,160,160)),
                ToTensord(keys=("image", "label")),
                EnsureTyped(keys=["image", "label"])
           ]
     )

val_transforms = Compose(
            [
                LoadImaged(keys=("image", "label")),
                Clipping(keys=("label")),
                EnsureChannelFirstd(keys=["image", "label"]),
                Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 2.0),mode=("bilinear", "nearest"), )
                NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                ToTensord(keys=("image", "label")),
                EnsureTyped(keys=["image", "label"])
            ]
       )

2. I used the 3D Unet as a model. I took the same approach as in the spleen segmentation example with 2 outputs for a binary segmentation task.

model = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
    norm=Norm.BATCH,
).to(device)

Hence my model output at this point for batch size one should be of dimension (1, 2,160,160,160).

3. For the loss, I've also used the same approach as in the example. I am assuming that the arguments make the output (1, 2,160,160,160) compatible with the ground truth of dimension (1, 1,160,160,160).

loss_function = DiceLoss(to_onehot_y=True,softmax=True)

My loss is lowering during training:

4. I am applying post-transformations on the output data and labels to make them compatible with the dice_metric:

post_trans = Compose(
    [EnsureType(), AsDiscrete(argmax=True, to_onehot=2)]
)

post_trans_labels = Compose(
    [EnsureType(), AsDiscrete(to_onehot=2)]
) 

5. I am configuring the dice metric as follows, as my dataset is imbalanced and the background is the majority.

dice_metric = DiceMetric(include_background=False, reduction="mean")

For the dice_metric I am receiving a value of 0 almost throughout the whole training and the final output segmentations look almost like the opposite of the ground truth:

Are there any configurations that do not make sense in my scenario? Thanks a lot for the help!

[LSnyd]

@yiheng-wang-nv sorry for the direct tag, but could you take a look at this by chance?
As you were working on the 3D Segmentation example. Thanks a lot! :)

[yiheng-wang-nv]

Just for my understanding, RandCropByPosNegLabeld takes the whole sample and crops out parts with positive and negative labels so it reaches, in the end, a total of the defined spatial_size?

Yes. Looking forward to seeing your updated training results : )

[LSnyd]

Hi @yiheng-wang-nv,
unfortunately, the results don't look much better.

The model seemed again to lower the loss. This time the dice score got better, but only reached 0.29 for the validation set (0.55 for train set):

When I am checking the output of the model for the validation set, it just returns empty masks for almost all channels.
I haven't tried many hyperparameters yet, so that could be cause as well. Do you have any other idea though?
I will post my main configurations below again:

train_transforms = Compose(
            [
                LoadImaged(keys=("image", "label")),
                Clipping(keys=("label")),
                EnsureChannelFirstd(keys=["image", "label"]),
                Spacingd(keys=["image", "label"],pixdim=(1.5, 1.5, 2.0),mode=("bilinear", "nearest")),
                NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                CropForegroundd(keys=("image", "label"), source_key="image"),
                SpatialPadd(keys=("image", "label"),spatial_size=(96,96,96))
                RandCropByPosNegLabeld(
                  keys=["image", "label"],
                  label_key="label",
                  spatial_size=(96,96,96),
                  pos=1,
                  neg=1,
                  num_samples=4,
                  image_key="image",
                  image_threshold=0,
                  allow_smaller=True,
              ),
                ToTensord(keys=("image", "label")),
                EnsureTyped(keys=["image", "label"])

    ]
)
val_transforms = Compose(
            [
                LoadImaged(keys=("image", "label")),
                Clipping(keys=("label")),
                EnsureChannelFirstd(keys=["image", "label"]),
                Spacingd(keys=["image", "label"],pixdim=(1.5, 1.5, 2.0),mode=("bilinear", "nearest") ),
                NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                ToTensord(keys=("image", "label")),
                EnsureTyped(keys=["image", "label"])
            ]
)



model = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
    norm=Norm.BATCH,
).to(device)

loss_function = DiceLoss(to_onehot_y=True,softmax=True)
optimizer = torch.optim.Adam(model.parameters(),  lr=lr, weight_decay=1e-5)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_epochs)

dice_metric = DiceMetric(include_background=False, reduction="mean")
dice_metric_batch = DiceMetric(include_background=False, reduction="mean_batch")

post_trans = Compose(
    [EnsureType(), AsDiscrete(argmax=True, to_onehot=2)]
)

post_trans_labels = Compose(
    [EnsureType(), AsDiscrete(to_onehot=2)]
)

[yiheng-wang-nv]

Hi @LSnyd , sorry for late reply, I did not receive the notification email (or missed).
In train_transforms , put NormalizeIntensityd before CropForegroundd may introduce some issues I think:
CropForegroundd uses margin = 0 in default, but NormalizeIntensityd will convert all values that less than mean into negative values.
Could you try to use ScaleIntensityRanged to replace NormalizeIntensityd and have a try?

[TuanDTr]

Hi @yiheng-wang-nv and @LSnyd, I have a similar problem. I am training a lung tumor segmentation task (Medical Segmentation Decathlon task 6). What I notice is that the tumor size is pretty small which could likely affect the computation of dice score because a small shift of several pixels has a significant on the dice score. On my first attempt, I included the background in both dice loss and validation dice score, but I soon realized that the improvement came only with the background, not the foreground. When I switched to using only the foreground, the validation dice score stucks at 0.00xx.

Fyi, I trained the UNETR network but using the transforms pipeline by nnUNet.
I just wonder if you guys have any update on this thread?

[yiheng-wang-nv]

Hi @beantr96 , how about the dice score for the training set? If it is high, it means overfit too much, you may need to add more augmentation. If the training dice score is also low, it means the task is hard to learn (underfit) with current parameters.

For the task6, it does hard to train according to my experience when using DynUNet (https://github.com/Project-MONAI/tutorials/tree/ff586e0cf24ac164553327b9453edab147cffa9e/modules/dynunet_pipeline/commands/task06), I used sgd with lr=1e-2 for 3000 epochs and then finetune with 1e-3 for 1000 epochs for the whole training. I would suggest to use the same pipeline except the network part, and try to reproduce similar scores first.