Multi-label segmentation predictions not binarised

[ghost]

Hi,

I've successfully used MONAI for some single-label segmentation and have now moved on to a more complicated multi-label problem for which I've been following the BRATS Brain tutorial.

It runs without any errors and seemingly works fine, the training loader images/labels look good, loss and Dice score both progressing as I'd expect but then have a problem when I get to the end and use the trained model on validation cases. For now I'm using only a few cases at a low res so don't worry about the actual segmentation quality, but I don't know why the result is not binarised?

I have two labels, with greyscale intensity 1 and 2 following the convention as in the BRATS database using "ConvertToMultiChannelBasedOnBratsClassesd" as in the tutorial.

Here is the model:

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = UNet(
    dimensions=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
).to(device)
loss_function = DiceLoss(to_onehot_y=False, sigmoid=True, squared_pred=True)
optimizer = torch.optim.Adam(
    model.parameters(), 1e-4, weight_decay=1e-5, amsgrad=True
)

Edit to add:

Here is the code I used for creating the segmentation figure:

for i in range(2):
        plt.subplot(1, 2, i + 1)
        plt.title(f"output channel {i}")
        plt.imshow(val_output[0, i, :, :, example_slice].detach().cpu())

I've tried using argmax & AsDiscrete but think the problem I'm facing occurs before this point.

I'm sure I'm missing something simple but have spent quite a few hours now looking for it! Thanks

[Nic-Ma]

Hi @jnc13 ,

Thanks for your interest here.
The output of all the MONAI networks is logits, please add Activations and AsDiscrete for it.

Thanks.

[ghost]

Thanks Nic. Yes, here you go:

import os
import shutil
import tempfile
import glob

import matplotlib.pyplot as plt
import numpy as np
from monai.apps import DecathlonDataset
from monai.data import CacheDataset, DataLoader, Dataset, NiftiSaver
from monai.config import print_config
from monai.data import DataLoader
from monai.losses import DiceLoss
from monai.metrics import DiceMetric
from monai.handlers import SegmentationSaver
from monai.networks.nets import UNet
from monai.transforms import (
    Activations,
    AsChannelFirstd,
    AsDiscrete,
    CenterSpatialCropd,
    Compose,
    LoadImaged,
    MapTransform,
    NormalizeIntensityd,
    Orientationd,
    RandFlipd,
    RandScaleIntensityd,
    RandShiftIntensityd,
    RandSpatialCropd,
    Spacingd,
    ToTensord,
    
    RandCropByPosNegLabeld,
    AddChanneld,
)
from monai.utils import set_determinism

import torch

### Directory setup

cur_dir = os.getcwd()
print('Working directory: ', cur_dir)

data_dir = os.path.join(cur_dir, "Data/HR_quick")
print('Data directory: ', data_dir)
set_determinism(seed=0)


### Define a new transform to convert brain tumor labels
#Here we convert the multi-classes labels into multi-labels segmentation task in One-Hot format.

class ConvertToMultiChannelBasedOnBratsClassesd(MapTransform):
    """
    Convert labels to multi channels based on Jenny's classes:
    label 1 is the pulmonary artery
    label 2 is the aorta
    The possible classes are PA or AO.

    """
    def __call__(self, data):
        d = dict(data)
        for key in self.keys:
            result = []
            # label 1 is PA
            result.append(d[key] == 1)
            # label 2 is Aorta
            result.append(d[key] == 2)
            d[key] = np.stack(result, axis=0).astype(np.float32)
        return d  


### Setup transforms for training and validation

train_transform = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        AddChanneld(keys=["image"]),
        ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
        Spacingd(
            keys=["image", "label"],
            pixdim=(2, 2, 2),
            mode=("bilinear", "nearest"),
        ),
        Orientationd(
            keys=["image", "label"],
            axcodes="PLS"
        ), ###Originally RAS, eventually use LPS not PLS
        CenterSpatialCropd(keys=["image", "label"],
                           roi_size=[48, 48, 48]),
        NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
        ToTensord(keys=["image", "label"]),
    ]
)
val_transform = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        AddChanneld(keys=["image"]),
        ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
        Spacingd(
            keys=["image", "label"],
            pixdim=(2, 2, 2),
            mode=("bilinear", "nearest"),
        ),
        Orientationd(
            keys=["image", "label"],
            axcodes="PLS"
        ), ### PLS easier to look at, use LPS for final runs
        CenterSpatialCropd(keys=["image", "label"],
                           roi_size=[48, 48, 48]),
        NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
        ToTensord(keys=["image", "label"]),
    ]
)

train_images = sorted(glob.glob(os.path.join(data_dir, "imagesTr", "*.nii.gz")))
train_labels = sorted(glob.glob(os.path.join(data_dir, "labelsTr", "*.nii.gz")))
val_images = sorted(glob.glob(os.path.join(data_dir, "imagesVa", "*.nii.gz")))
val_labels = sorted(glob.glob(os.path.join(data_dir, "labelsVa", "*.nii.gz")))


print('Training images found: ',len(train_images))
print('Training segmentations found: ',len(train_labels))
print('Validation images found: ',len(val_images))
print('Validation segmentations found: ',len(val_labels))

train_files = [
    {"image": image_name, "label": label_name}
    for image_name, label_name in zip(train_images, train_labels)
]

val_files = [
    {"image": image_name, "label": label_name}
    for image_name, label_name in zip(val_images, val_labels)
]


train_ds = CacheDataset(
    data=train_files,
    transform=train_transform,
    cache_rate=1.0,
    num_workers=0)

train_loader = DataLoader(
    train_ds,
    batch_size=config['trainbatchsize'],
    shuffle=True,
    num_workers=0) ###Originally num_workers=4

val_ds = CacheDataset(
    data=val_files,
    transform=val_transform,
    cache_rate=1.0,
    num_workers=0)

val_loader = DataLoader(
    val_ds,
    batch_size=config['validationbatchsize'],
    shuffle=False,
    num_workers=0) ###Originally num_workers=4

print(f"train image shape: {train_ds[0]['image'].shape}")
print(f"train label shape: {train_ds[0]['label'].shape}")

print(f"val image shape: {val_ds[0]['image'].shape}")
print(f"val label shape: {val_ds[0]['label'].shape}")

### Visualise input data 

print(f"image shape: {train_ds[0]['image'].shape}")
plt.figure("Train image", (14, 6))
for i in range(1):
    plt.subplot(1, 3, i + 1)
    plt.title(f"image channel {i}")
    plt.imshow(train_ds[0]["image"][i, :, :, 37].detach().cpu(), cmap="gray")
plt.show()
# also visualize the label channels corresponding to this image
print(f"label shape: {train_ds[0]['label'].shape}")
plt.figure("Train label", (10, 4))
for i in range(2):
    plt.subplot(1, 2, i + 1)
    plt.title(f"label channel {i}")
    plt.imshow(train_ds[0]["label"][i, :, :, 37].detach().cpu())
plt.show()


# pick one val image to visualize and check the 2 label channels
print(f"image shape: {val_ds[0]['image'].shape}")
plt.figure("Val image", (14, 6))
for i in range(1):
    plt.subplot(1, 3, i + 1)
    plt.title(f"image channel {i}")
    plt.imshow(val_ds[0]["image"][i, :, :, 44].detach().cpu(), cmap="gray")
plt.show()
# also visualize the 3 channels label corresponding to this image
print(f"label shape: {val_ds[0]['label'].shape}")
plt.figure("Val label", (10, 4))
for i in range(2):
    plt.subplot(1, 2, i + 1)
    plt.title(f"label channel {i}")
    plt.imshow(val_ds[0]["label"][i, :, :, 44].detach().cpu())
plt.show()


### standard PyTorch program style: create UNet, DiceLoss and Adam optimizer
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = UNet(
    dimensions=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
).to(device)
loss_function = DiceLoss(to_onehot_y=False, sigmoid=True, squared_pred=True)
optimizer = torch.optim.Adam(
    model.parameters(), 1e-4, weight_decay=1e-5, amsgrad=True
)

###Training

max_epochs = 20
val_interval = 2
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = []
metric_values = []
metric_values_pa = []
metric_values_ao = []

for epoch in range(max_epochs):
    print("-" * 10)
    print(f"epoch {epoch + 1}/{max_epochs}")
    model.train()
    epoch_loss = 0
    step = 0
    for batch_data in train_loader:
        step += 1
        inputs, labels = (
            batch_data["image"].to(device),
            batch_data["label"].to(device),
        )
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
        print(
            f"{step}/{len(train_ds) // train_loader.batch_size}"
            f", train_loss: {loss.item():.4f}"
        )
    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():
            dice_metric = DiceMetric(include_background=True, reduction="mean")
            post_trans = Compose(
                [Activations(sigmoid=True), AsDiscrete(threshold_values=True)]
            )
            metric_sum = metric_sum_pa = metric_sum_ao = 0.0
            metric_count = (
                metric_count_pa
            ) = metric_count_ao = 0
            for val_data in val_loader:
                val_inputs, val_labels = (
                    val_data["image"].to(device),
                    val_data["label"].to(device),
                )
                val_outputs = model(val_inputs)
                val_outputs = post_trans(val_outputs)
                # compute overall mean dice
                value, not_nans = dice_metric(y_pred=val_outputs, y=val_labels)
                not_nans = not_nans.item()
                metric_count += not_nans
                metric_sum += value.item() * not_nans
                # compute mean dice for PA
                value_pa, not_nans = dice_metric(
                    y_pred=val_outputs[:, 0:1], y=val_labels[:, 0:1]
                )
                not_nans = not_nans.item()
                metric_count_pa += not_nans
                metric_sum_pa += value_pa.item() * not_nans
                # compute mean dice for Ao
                value_ao, not_nans = dice_metric(
                    y_pred=val_outputs[:, 1:2], y=val_labels[:, 1:2]
                )
                not_nans = not_nans.item()
                metric_count_ao += not_nans
                metric_sum_ao += value_ao.item() * not_nans
            metric = metric_sum / metric_count
            metric_values.append(metric)
            metric_pa = metric_sum_pa / metric_count_pa
            metric_values_pa.append(metric_pa)
            metric_ao = metric_sum_ao / metric_count_ao
            metric_values_ao.append(metric_ao)

            if metric > best_metric:
                best_metric = metric
                best_metric_epoch = epoch + 1
                torch.save(
                    model.state_dict(),
                    os.path.join(cur_dir, "best_metric_model.pth"),
                )
                print("saved new best metric model")
            print(
                f"current epoch: {epoch + 1} current mean dice: {metric:.4f}"
                f" pa: {metric_pa:.4f} ao: {metric_ao:.4f}"
                f"\nbest mean dice: {best_metric:.4f}"
                f" at epoch: {best_metric_epoch}"
            )

print(
    f"train completed, best_metric: {best_metric:.4f}"
    f" at epoch: {best_metric_epoch}"
)


### Plot the loss and metric

plt.figure("train", (10, 4))
plt.subplot(1, 2, 1)
plt.title("Epoch Average Loss")
x = [i + 1 for i in range(len(epoch_loss_values))]
y = epoch_loss_values
plt.xlabel("epoch")
plt.plot(x, y, color="red")
plt.subplot(1, 2, 2)
plt.title("Val Mean Dice")
x = [val_interval * (i + 1) for i in range(len(metric_values))]
y = metric_values
plt.xlabel("epoch")
plt.plot(x, y, color="green")
plt.show()

plt.figure("train", (10, 4))
plt.subplot(1, 2, 1)
plt.title("Val Mean Dice PA")
x = [val_interval * (i + 1) for i in range(len(metric_values_pa))]
y = metric_values_pa
plt.xlabel("epoch")
plt.plot(x, y, color="blue")
plt.subplot(1, 2, 2)
plt.title("Val Mean Dice Ao")
x = [val_interval * (i + 1) for i in range(len(metric_values_ao))]
y = metric_values_ao
plt.xlabel("epoch")
plt.plot(x, y, color="brown")
plt.show()


#### Check best model output with the input image and label
model.load_state_dict(
    torch.load(os.path.join(cur_dir, "best_metric_model.pth"))
)
model.eval()
with torch.no_grad():
    # select one image to evaluate and visualize the model output
    val_image_to_check = 0 #choose which to observe
    example_slice = 40
    saver_monai = SegmentationSaver(output_dir="./172modeloutput")
    val_input = val_ds[val_image_to_check]["image"].unsqueeze(0).to(device)
    val_output = model(val_input)
    plt.figure("image", (10, 4))
    for i in range(1):
        plt.subplot(1, 1, i + 1)
        plt.title(f"image channel {i}")
        plt.imshow(val_ds[val_image_to_check]["image"][i, :, :, example_slice].detach().cpu(), cmap="gray")
    plt.show()
    # visualize the 2 channels label corresponding to this image
    plt.figure("label", (10, 4))
    for i in range(2):
        plt.subplot(1, 2, i + 1)
        plt.title(f"label channel {i}")
        plt.imshow(val_ds[val_image_to_check]["label"][i, :, :, example_slice].detach().cpu())
    plt.show()
    # visualize the 2 channels model output corresponding to this image
    plt.figure("output", (10, 4))
    for i in range(2):
        plt.subplot(1, 2, i + 1)
        plt.title(f"output channel {i}")
        plt.imshow(val_output[0, i, :, :, example_slice].detach().cpu())

    
    plt.figure("output binarised", (10, 4))
    for i in range(2):
        plt.subplot(1, 2, i + 1)
        plt.title(f"output channel {i}")
        argmax = AsDiscrete(argmax=True,
                            #n_classes=(2),
                            logit_thresh=0.1,
                            )(val_output)
        plt.imshow(argmax[0, 0, :, :, example_slice].detach().cpu())
        

Thanks in advance!

[rijobro]

Hi, from your code, I'm finding it hard to understand what is being visualised. You have two plots that have the title output channel {i}. One plots val_output (logits) and the other plots argmax (binary). Could you double check that your original screenshot is of the latter, not the former?

You should also be able to check with something like torch.unique(argmax), which should hopefully return [0, 1] for binary segmentation.

[Nic-Ma]

Hi @jnc13 ,

Thanks for providing the test code, please make sure to use both Activations and AsDiscrete transforms in below code block before visualization:

model.eval()
with torch.no_grad():
    # select one image to evaluate and visualize the model output
    val_image_to_check = 0 #choose which to observe
    example_slice = 40
    saver_monai = SegmentationSaver(output_dir="./172modeloutput")
    val_input = val_ds[val_image_to_check]["image"].unsqueeze(0).to(device)
    val_output = model(val_input)
    plt.figure("image", (10, 4))
    for i in range(1):
        plt.subplot(1, 1, i + 1)
        plt.title(f"image channel {i}")
        plt.imshow(val_ds[val_image_to_check]["image"][i, :, :, example_slice].detach().cpu(), cmap="gray")
    plt.show()
    # visualize the 2 channels label corresponding to this image
    plt.figure("label", (10, 4))
    for i in range(2):
        plt.subplot(1, 2, i + 1)
        plt.title(f"label channel {i}")
        plt.imshow(val_ds[val_image_to_check]["label"][i, :, :, example_slice].detach().cpu())
    plt.show()
    # visualize the 2 channels model output corresponding to this image
    plt.figure("output", (10, 4))
    for i in range(2):
        plt.subplot(1, 2, i + 1)
        plt.title(f"output channel {i}")
        plt.imshow(val_output[0, i, :, :, example_slice].detach().cpu())

    
    plt.figure("output binarised", (10, 4))
    for i in range(2):
        plt.subplot(1, 2, i + 1)
        plt.title(f"output channel {i}")
        argmax = AsDiscrete(argmax=True,
                            #n_classes=(2),
                            logit_thresh=0.1,
                            )(val_output)
        plt.imshow(argmax[0, 0, :, :, example_slice].detach().cpu())

Thanks.

[ghost]

Thanks so much @rijobro and @Nic-Ma as ever for coming to the rescue.

You're right @rijobro sorry I didn't make it clear in the code. The original screenshot was of the former visualisation block.

I was getting a little confused because from how I interpret the BRATS tutorial the transformation "post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold_values=True)])" is used within the model creation loop and after which for visualiation just "val_output = model(val_input)" was used without the transforms, I'm not sure how it works in that scenario and not mine?

Nevertheless, I've added the transforms for visualisation and it is working as expected. :)
In case it helps anyone else I now have:

model.load_state_dict(
    torch.load(os.path.join(cur_dir, "best_metric_model.pth"))
)
model.eval()
with torch.no_grad():
    # select one image to evaluate and visualize the model output
    val_image_to_check = 0 #choose which to observe
    example_slice = 40
    saver = NiftiSaver(output_dir="./modeloutput", dtype=(['numpy.float16']))
    val_input = val_ds[val_image_to_check]["image"].unsqueeze(0).to(device)
    post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold_values=True)]
            )
    val_output = model(val_input)
    val_output_bin = post_trans(val_output)
    plt.figure("image", (10, 4))
    for i in range(1):
        plt.subplot(1, 1, i + 1)
        plt.title(f"image channel {i}")
        plt.imshow(val_ds[val_image_to_check]["image"][i, :, :, example_slice].detach().cpu(), cmap="gray")
    plt.show()

    # visualize the 2 channels label corresponding to this image
    plt.figure("label", (10, 4))
    for i in range(2):
        plt.subplot(1, 2, i + 1)
        plt.title(f"label channel {i}")
        plt.imshow(val_ds[val_image_to_check]["label"][i, :, :, example_slice].detach().cpu())
    plt.show()
    # visualize the 2 channels model output corresponding to this image
    plt.figure("output", (10, 4))
    for i in range(2):
        plt.subplot(1, 2, i + 1)
        plt.title(f"output channel val outputs {i}")
        plt.imshow(val_output[0, i, :, :, example_slice].detach().cpu())
     
    # visualize the 2 final binarised channels model output corresponding to this image
    plt.figure("output binarised", (10, 4))
    for i in range(2):
        plt.subplot(1, 2, i + 1)
        plt.title(f"output binarised argmax {i}")
        plt.imshow(val_output_bin[0, i, :, :, example_slice].detach().cpu())
    print("val_output:", torch.unique(val_output))   
    print("val_output_bin:", torch.unique(val_output_bin))   
        #largest = KeepLargestConnectedComponent(applied_labels=[1])(argmax)
        #plt.imshow(argmax.detach().cpu()[0, i, :, :, example_slice])   

I chose to create a new variable val_output_bin because I quite like seeing the untransformed version output too but you you could instead overwrite val_output.

Ignore the poor quality, I've only trained on a few images and didn't use KeepLargestConnectedComponent yet, the important thing is that the pipeline now works.

Thanks again both of you, us novices really do appreciate you helping us out!

[Nic-Ma]

Sounds great!