How to correctly mutli-channel labels in single channel after onehot encoding?

[OmarAshkar]

Hello,
I have a struggle here. I started from the unet_2d tutorial and changed few things:
1- included 5 clasess (4 labels + 0 bg)
2- used reduction in Dice as "None"
3- Changed the post-process in the evaluation to have argmax.
4- change to swin-unetr

Here is the training code:

       train_transforms = Compose(
        [
            LoadImaged(keys=["img", "seg"], image_only=True),
            EnsureChannelFirstd(keys=["img", "seg"], channel_dim=-1),
            ScaleIntensityd(keys=["img", "seg"]),
            AsDiscreted(keys=["seg"], to_onehot=[5]), 
            # FIXME this caues error
            # RandCropByPosNegLabeld(
            #     keys=["img", "seg"], label_key="seg", spatial_size=[96, 96], pos=1, neg=1, num_samples=4
            # ),
            # RandRotate90d(keys=["img", "seg"], prob=0.5, spatial_axes=[0, 1]),
        ]
    )
    val_transforms = Compose(
        [
            LoadImaged(keys=["img", "seg"], image_only= True),
            EnsureChannelFirstd(keys=["img", "seg"], channel_dim=-1),
            ScaleIntensityd(keys=["img", "seg"]),
            AsDiscreted(keys=["seg"], to_onehot=[5]), 
        ]
    )

    from monai.data import Dataset
    # define dataset, data loader
    check_ds = monai.data.Dataset(data=train_files, transform=train_transforms)
    # use batch_size=2 to load images and use RandCropByPosNegLabeld to generate 2 x 4 images for network training
    check_loader = DataLoader(check_ds, batch_size=2, num_workers=4)
    check_data = monai.utils.misc.first(check_loader)
    print(check_data["img"].shape, check_data["seg"].shape)

    # create a training data loader
    train_ds = monai.data.Dataset(data=train_files, transform=train_transforms)
    # use batch_size=2 to load images and use RandCropByPosNegLabeld to generate 2 x 4 images for network training
    train_loader = DataLoader(
        train_ds,
        batch_size=16,
        shuffle=True,
        num_workers=4,
        collate_fn=list_data_collate,
        pin_memory=torch.cuda.is_available(),
    )
    # create a validation data loader
    val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
    val_loader = DataLoader(val_ds, batch_size=16, num_workers=4)
    dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
    post_trans = Compose([Activations(softmax=True), AsDiscrete(threshold=0.5,argmax  = True, to_onehot=5)])
    # create UNet, DiceLoss and Adam optimizer
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = monai.networks.nets.SwinUNETR(
        img_size=(512, 512),
        spatial_dims=2,
        in_channels=3,
        out_channels=5,
    ).to(device)
    loss_function = monai.losses.DiceLoss(sigmoid=True)
    optimizer = torch.optim.Adam(model.parameters(), 1e-3)

    # start a typical PyTorch training
    val_interval = 2
    best_metric = -1
    best_metric_epoch = -1
    epoch_loss_values = list()
    metric_values = list()
    writer = SummaryWriter()
    n_epoch = 100
    for epoch in range(n_epoch):
        print("-" * 10)
        print(f"epoch {epoch + 1}/{n_epoch}")
        model.train()
        epoch_loss = 0
        step = 0
        for batch_data in train_loader:
            step += 1
            inputs, labels = batch_data["img"].to(device), batch_data["seg"].to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = loss_function(outputs, labels)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
            epoch_len = len(train_ds) // train_loader.batch_size
            print(f"{step}/{epoch_len}, train_loss: {loss.item():.4f}")
            writer.add_scalar("train_loss", loss.item(), epoch_len * epoch + step)
        epoch_loss /= step
        epoch_loss_values.append(epoch_loss)
        print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")

        if (epoch + 1) % val_interval == 0:
            model.eval()
            with torch.no_grad():
                val_images = None
                val_labels = None
                val_outputs = None
                for val_data in val_loader:
                    val_images, val_labels = val_data["img"].to(device), val_data["seg"].to(device)
                    roi_size = (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)
                # aggregate the final mean dice result
                metric = dice_metric.aggregate().item()
                print(f'Validation score is {dice_metric.aggregate("mean_batch")}')
                # reset the status for next validation round
                dice_metric.reset()
                metric_values.append(metric)
                if metric > best_metric:
                    best_metric = metric
                    best_metric_epoch = epoch + 1
                    torch.save(model.state_dict(), "best_metric_model_swinunetr.pth")
                    print("saved new best metric model")
                print(
                    "current epoch: {} current mean dice: {:.4f} best mean dice: {:.4f} at epoch {}".format(
                        epoch + 1, metric, best_metric, best_metric_epoch
                    )
                )

The training was finished successfully but I need to evaluate and save the output inference

here is the eval code

 # define transforms for image and segmentation
    val_transforms = Compose(
        [
            LoadImaged(keys=["img", "seg"], image_only=True),
            EnsureChannelFirstd(keys=["img", "seg"], channel_dim=-1),
            ScaleIntensityd(keys=["img", "seg"]),
            AsDiscreted(keys=["seg"], to_onehot=[5]), 
        ]
    )
    val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
    # sliding window inference need to input 1 image in every iteration
    val_loader = DataLoader(val_ds, batch_size=1, num_workers=4, collate_fn=list_data_collate)
    dice_metric = DiceMetric(include_background=True, reduction="None", get_not_nans=False)
    post_trans = Compose([Activations(softmax = True), 
                          AsDiscrete(argmax = True)])
    saver = SaveImage(output_dir="./output", output_ext=".png", output_postfix="seg")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = SwinUNETR(
        img_size=(512, 512),
        spatial_dims=2,
        in_channels=3,
        out_channels=5,
    ).to(device)

    model.load_state_dict(torch.load("unet/best_metric_model_swinunetr.pth"))

    model.eval()
    with torch.no_grad():
        for val_data in val_loader:
            val_images, val_labels = val_data["img"].to(device), val_data["seg"].to(device)
            # define sliding window size and batch size for windows inference
            roi_size = (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)]
            val_labels = decollate_batch(val_labels)
            # compute metric for current iteration
            dice_metric(y_pred=val_outputs, y=val_labels)
            for val_output in val_outputs:
                saver(val_output)
        # aggregate the final mean dice result
        print("evaluation metric:", dice_metric.aggregate().item())
        # reset the status
        dice_metric.reset()

By adding argmax before dice_metric(), I get size mismatch
ValueError: y_pred and y should have same shapes, got torch.Size([1, 1, 512, 512]) and torch.Size([1, 5, 512, 512]).

I need to return and save one channel with the class predicted.

Thanks in advance.

[KumoLiu]

Hi @OmarAshkar, you can also add an AsDiscrete for the label.

post_trans_lbl = Compose([ AsDiscrete(argmax = True)])
val_labels = [post_trans_lbl(i) for i in decollate_batch(val_labels)]

Thanks!

[OmarAshkar]

Thanks @KumoLiu. I just keep getting inconsitencies here and there. For example, in this tutorial https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/brats_segmentation_3d.ipynb

3 labels were passed as onehot (not using AsDiscrete for that process). and the post_trans looked like that and it resulted in 3 label channels again.

post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])

If I used used this same config, I get only one channel with my 5 labels that causes error on calculate dice between the outputs and labels.

But if I passed to_onehot, I get the correct shape. So why the tutorial data gives different results?

Also, would it differ here if I used threshold or argmax? Also did the loss function looks correctly like that for multi-label?

loss_function = monai.losses.DiceLoss(sigmoid=True)

Many thanks!

[KumoLiu]

Hi @OmarAshkar, sorry for the later response.

In the tutorial, it uses post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)]) which will also keep the channel dim.
But in your case, you used post_trans = Compose([Activations(softmax = True), AsDiscrete(argmax = True)]), so the channel dim will be one.
You only need to consider whether the predictions and labels are one-hot.

Hope it helps, thanks!