Convert mutli-labels to multi-channel

[LIKP0]

Hello all,

For the turtorial 3d_segmentation/torch/unet_training_array.py, I want to try create_test_image_3d(128, 128, 128, num_seg_classes=5) and I also modify the network architecture with

model = monai.networks.nets.UNet(
        spatial_dims=3,
        in_channels=1,
        out_channels=6,
        channels=(16, 32, 64, 128, 256),
        strides=(2, 2, 2, 2),
        num_res_units=2,
    ).to(device)

But I encounter this problem:
ValueError: y_pred and y should have same shapes, got torch.Size([1, 6, 128, 128, 128]) and torch.Size([1, 1, 128, 128, 128]). at dice_metric(y_pred=val_outputs, y=val_labels):

                for val_data in val_loader:
                    val_images, val_labels = val_data[0].to(device), val_data[1].to(device)
                    roi_size = (96, 96, 96)
                    sw_batch_size = 4
                    val_outputs = sliding_window_inference(val_images, roi_size, sw_batch_size, model)
                    val_outputs = [post_trans(i) for i in decollate_batch(val_outputs)]
                    # compute metric for current iteration
                    dice_metric(y_pred=val_outputs, y=val_labels)

I have searched the similar questions like Segmentation example on multiple labels, but I can't find a similar one_hot option for dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)

Could anyone tell me how to solve this problem? For example, using built-in transforms to convert multi-lables to multi-channel?
Really appreciate your help.

[KumoLiu]

Hi @LIKP0, you can simply add a post-transform for the label like this:

y_trans = Compose([AsDiscrete(to_onehot=num_class)])
y_onehot = [y_trans(i) for i in decollate_batch(y, detach=False)]

You can also refer to this tutorial.
Hope it can help you, thanks!

[KumoLiu]

Hi @LIKP0, could you please check your labels' shape, simply do it like this:

for i in decollate_batch(val_labels):
    print(i.shape)

If it does not contain a channel dim, you should add a channel dim by using EnsureChannelFirst.
Thanks!

[LIKP0]

I got it! Thank you very much!
I find the shape is (1, 1, 128, 128, 128) and AsDiscrete(to_onehot=6, dim=1) will help!
The dim=1 is important otherwise AsDiscrete will change to (6, 1, 128, 128, 128) instead of (1, 6, 128, 128, 128)

[LIKP0]

By the way, I find the easiest way to train multi-label case for tutorial 3d_segmentation/torch/unet_training_array.py is to change three points:

1. Change num_seg_classes=N, which denotes the number of label categories you have:

    for i in range(40):
        im, seg = create_test_image_3d(128, 128, 128, num_seg_classes=5) # Change this

        n = nib.Nifti1Image(im, np.eye(4))
        nib.save(n, os.path.join(tempdir, f"im{i:d}.nii.gz"))

        n = nib.Nifti1Image(seg, np.eye(4))
        nib.save(n, os.path.join(tempdir, f"seg{i:d}.nii.gz"))

2. Change transform of segs to convert multi lables to one-hot label with multi channels (N+1 for background). And I think maybe transform on train and validation seg set in advance is better than transform at loss function like loss_function = monai.losses.DiceLoss(sigmoid=True, to_onehot_y=True) or post_trans=Compose([AsDiscrete(to_onehot=6, dim=1)]) afterwards.

# define transforms for image and segmentation
    train_imtrans = Compose(
        [
            ScaleIntensity(),
            EnsureChannelFirst(),
            RandSpatialCrop((96, 96, 96), random_size=False),
            RandRotate90(prob=0.5, spatial_axes=(0, 2)),
        ]
    )
    train_segtrans = Compose(
        [
            EnsureChannelFirst(),
            AsDiscrete(to_onehot=6),  # Add this
            RandSpatialCrop((96, 96, 96), random_size=False),
            RandRotate90(prob=0.5, spatial_axes=(0, 2)),
        ]
    )
    val_imtrans = Compose([ScaleIntensity(), EnsureChannelFirst()])
    val_segtrans = Compose([EnsureChannelFirst(),
                            AsDiscrete(to_onehot=6),  # Add this
                            ])

3. Change the network output channels to N+1 for including background.

model = monai.networks.nets.UNet(
        spatial_dims=3,
        in_channels=1,
        out_channels=6, # Change this
        channels=(16, 32, 64, 128, 256),
        strides=(2, 2, 2, 2),
        num_res_units=2,
    ).to(device)

Hope this could help somebody.