strange and bad results from training 3d segmentation model

[esther1262]

Hi, I followed the tutorial spleen_segmentation_3d.ipynb tried to train the model with my own data. In the binary labelmap, 1=brain, 0=non-brain areas or background. The images and labelmaps only have one channel, dimension 256x256x16. In some slices, the brain area is relatively small compared to the whole image. I got very strange results and trying to figure out why. So I have couple questions. If anyone can help or give me some hints, it will be highly appreciated!

1. I used the same transformation for training and validation sets. Will the labelmaps be transformed too?

augmentation_transforms = Compose(
    [   LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        EnsureTyped(keys=["image", "label"]),
        RandFlipd(keys=["image", "label"], prob=0.8),
        RandScaleIntensityd(keys=["image", "label"], prob=0.8, factors=(5,10)),
        RandShiftIntensityd(keys=["image", "label"], prob=0.8, offsets=(5,10)),
        RandBiasFieldd(keys="image", coeff_range=(0.2,0.3), prob=0.5),
        RandRotated(keys=["image", "label"], prob=0.8, range_z=1, mode=['bilinear','nearest']), 
        RandZoomd(keys=["image", "label"], prob=0.8, min_zoom=0.8, max_zoom=1.5, mode=['area','nearest'])]
)

The rest codes are pretty much the same as the tutorial. so I only attached the parts that I made changes. I set in_channels and out_channels to 1. I added Gradient scaler. If I changed the values of roi_size (sliding window inference), I would get a dimensionality error ValueError: Expected more than 1 spatial element when training, got input size torch.Size([2, 128, 1, 1, 1]). Changing include_background=True or removing gradient scaler didnt improve the model performance.

UNet_meatdata = {
    "spatial_dims": 3,
    "in_channels": 1,
    "out_channels": 1,
    "channels": (16, 32, 64, 128, 256), # 5 layers
    "strides": (2, 2, 2, 2)
    }

model = UNet(**UNet_meatdata).to(device)
loss_function = DiceLoss(include_background=False, smooth_nr=1e-05, smooth_dr=1e-5, to_onehot_y=False, sigmoid=True)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-6) # Adam optimizer
dice_metric = DiceMetric(include_background=False, reduction="mean")
...# codes skipped
scaler = torch.cuda.amp.GradScaler()
...# codes skipped
post_pred = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
post_label = Compose([AsDiscrete()])

...# codes skipped
                roi_size = (32,32,16)
                sw_batch_size = 4

The tutorial also models a binary segmentation, but why they can use out_channels =2? When I put 2 in my model, I got a dimensionality error. AssertionError: ground truth has different shape (torch.Size([4, 0, 256, 256, 16])) from input (torch.Size([4, 1, 256, 256, 16])) Also, why can they DiceLoss(to_onehot_y=True, softmax=True) with binary label?

2. The dice coefficient was very low like around 0.10, even though I completed all 800 epochs. Example output:

epoch 6/800
epoch 6 average loss: 0.0200
current epoch: 6 current mean dice: 0.0649 
best mean dice: 0.0779  at epoch: 2
time consuming of epoch 6 is: 31.6528
----------
epoch 7/800
epoch 7 average loss: 0.1239
current epoch: 7 current mean dice: 0.0000 
best mean dice: 0.0779  at epoch: 2
time consuming of epoch 7 is: 31.8521
----------
epoch 8/800
epoch 8 average loss: 0.0929
current epoch: 8 current mean dice: 0.0390 
best mean dice: 0.0779  at epoch: 2
time consuming of epoch 8 is: 31.8742

epoch 24/800
epoch 24 average loss: -0.0154
current epoch: 24 current mean dice: 0.0482 
best mean dice: 0.0971  at epoch: 16
time consuming of epoch 24 is: 33.2549
----------
epoch 25/800
epoch 25 average loss: 0.0466
current epoch: 25 current mean dice: 0.0721 
best mean dice: 0.0971  at epoch: 16
time consuming of epoch 25 is: 32.1383
----------
epoch 26/800
epoch 26 average loss: 0.0058
current epoch: 26 current mean dice: 0.0000 
best mean dice: 0.0971  at epoch: 16
time consuming of epoch 26 is: 31.3328

3. Every time I used the following codes (exactly the same code) to check the model output with input image and label, I would get a different image and label. Also the image was shown with some transformation, but the label showed another degree of transformation. They were no longer matched. In some cases, the predicted segmentation was the whole head (brain + non brain tissue) rather than just the brain area.

with torch.no_grad():
    val_input = val_ds[6]["image"].unsqueeze(0).to(device)
    roi_size = (32,32,16)
    sw_batch_size = 4
    val_outputs = sliding_window_inference(val_input, roi_size, sw_batch_size, model)
    val_outputs = post_pred(val_outputs[0])

    plt.figure("check", (18, 6))
    plt.subplot(1, 3, 1)
    plt.title('image')
    plt.imshow(val_ds[6]["image"][0, :, :, 10], cmap="gray")
    plt.subplot(1, 3, 2)
    plt.title('label')
    plt.imshow(val_ds[6]["label"][0, :, :, 10])
    plt.subplot(1, 3, 3)
    plt.title('predicted label')
    plt.imshow(val_outputs[0,:,:,10].detach().cpu())
    plt.show()

[diazandr3s]

Hi @esther1262,

Thanks for opening this discussion.
With regard to your questions:

I used the same transformation for training and validation sets. Will the labelmaps be transformed too?

That's correct. That's being specified in the arguments you've used in the transforms: keys=["image", "label"] - label being the labelmap

The tutorial also models a binary segmentation, but why they can use out_channels =2? When I put 2 in my model, I got a dimensionality error. AssertionError: ground truth has different shape (torch.Size([4, 0, 256, 256, 16])) from input (torch.Size([4, 1, 256, 256, 16])) Also, why can they DiceLoss(to_onehot_y=True, softmax=True) with binary label?

I'd recommend keeping the output channels in 2. One channel is for the background and the other is for the foreground.
For binary segmentation tasks, you could also use Softmax and the AsDiscrete transform to get only one channel as output (https://github.com/Project-MONAI/MONAILabel/blob/main/sample-apps/radiology/lib/trainers/segmentation_spleen.py#L83-L88).
The other way is by using sigmoid activation and keeping one channel as output. We don't do this in the example/tutorial as we wanted to be more generic. Remember that binary segmentation is a special case of multilabel segmentation.

Every time I used the following codes (exactly the same code) to check the model output with input image and label, I would get a different image and label. Also the image was shown with some transformation, but the label showed another degree of transformation. They were no longer matched. In some cases, the predicted segmentation was the whole head (brain + non brain tissue) rather than just the brain area.

The model may need more training. However, it is difficult to asses without seeing the results/images. How many images are you using for training? Are you using MR or CT images for this task? If using MR images, you may want to consider a different type of intensity normalization.

Hope that helps,

[esther1262]

Hi Andres, when I followed the tutorial codes exactly, it would give me errors, likely from the transformation ScaleIntensityRanged and RandCropByPosNegLabelsd, so I removed these and used my own transformations e.g. RandScaleIntensityd and RandShiftIntensityd. For some reasons, these settings also caused some dimensionality error or index error (I already changed out_channel to 2). Upon editing, the performance of model has been improving. Do you know why these settings are more prone to errors or how to choose the corresponding parameters? (I tested couple, none were successful.)

I used the same PyTorch training process codes but the predicted label has opposite representation.

This is an example where the predicted label shows the brain + non-brain regions, rather than just the brain region. (I didnt finish all epochs, that's why the predicted one is like this, but you can roughly see the edges of the head.)

Also, is that possible to know the true size of training set after augmentation?
Thanks

[diazandr3s]

Hi @esther1262,

It is difficult for me to help without seeing the errors you're getting when using the transforms ScaleIntensityRanged and RandCropByPosNegLabelsd.

There is one typical error that occurs when you use a transform that tries to crop an image that is smaller than the actual crop in any of the dimensions (i.e. X, Y, or Z). So I recommend you check that.

Also, is that possible to know the true size of training set after augmentation?

You could create a breakpoint and check the dictionary that is passed through the different transforms.

Hope this helps,

[esther1262]

Hi @diazandr3s ,
Say I added RandScaleIntensityd(keys=["image", "label"], prob=1, factors=(5,10)) in the train_transforms, I got the error RuntimeError: index 10 is out of bounds for dimension 1 with size 2. Thanks

[diazandr3s]

Hi @esther1262,

Thanks for the details.
You are trying to scale the intensity of the image and the label (arguments keys=["image", "label"] in the RandScaleIntensityd transform). The label only has values of zeros and ones. That's why the error.

When scaling intensities, you don't usually apply those transforms to the label.

Hope it makes sense,

[esther1262]

Hi @diazandr3s, That totally makes sense! I dont know what I was thinking. One last question, do you know why the predicted label and the true label have opposite representations, like the picture shown in previous comment, background becomes 1, region=0? All the codes are exactly the same as tutorial like softmax, to_onehot, everything except the transformation parameters. dice value can be up to 0.9 and epoch loss can be 0.02.

[diazandr3s]

Hi @esther1262,

One last question, do you know why the predicted label and the true label have opposite representations, like the picture shown in previous comment, background becomes 1, region=0?

This is strange. It shouldn't happen.
The predicted mask should represent the background with zero and the foreground with one. Based on the Dice score you're getting, I suspect this is what's happening here.
Did you check the values from the numpy/torch array before plotting the mask? It may be something related to the way you're plotting the predicted and GT masks.

[esther1262]

I might have figured out the issue. I previously saw in the tutorial they just wanted to plot i=1 and thus created a break. I replaced val_data["image"].to(device) in sliding_window_inference to val_ds[0]["image"].unsqueeze(0).to(device) and removed the for loop.

Thanks again for all your kind help!