No output in visualizing segmented volumes

[Kiarashdnsh]

Hello,

I'm a novice in AI and image segmentation.
I have a dataset of [256, 256, 256]-sized images and labels. The input channel is 1 and the output channel is 3 (label IDs are 0, 1, and 2, respectively). I'm trying to segment the images using 4 networks (UNet, VNet, SegResNet, etc.) within the MONAI framework. Once each training and validation is done for each network, I try to visualize the segmentation using this code:

model.load_state_dict(torch.load(os.path.join("/content/drive/MyDrive/UNet", "best_metric_model.pth")))
model.eval()
with torch.no_grad():
    for i, val_data in enumerate(val_loader):
        roi_size = (96, 96, 96)
        sw_batch_size = 4
        val_outputs = sliding_window_inference(val_data["image"].to(device), roi_size, sw_batch_size, model)
        # plot the slice [:, :, 80]
        plt.figure("check", (18, 6))
        plt.subplot(1, 3, 1)
        plt.title(f"image {i}")
        plt.imshow(val_data["image"][0, 0, 80], cmap="gray")
        plt.subplot(1, 3, 2)
        plt.title(f"label {i}")
        plt.imshow(val_data["label"][0, 0, 80])
        plt.subplot(1, 3, 3)
        plt.title(f"output {i}")
        plt.imshow(torch.argmax(val_outputs, keepdim=True).detach().cpu()[0, 0, 0])
        plt.show()
        if i == 1:
            break

However, the output column would render empty after the run.
I have tried changing the numbers in the last brackets, but every try renders different errors. The above-mentioned code is the only code that renders no errors but outputs nothing.

for the training I have used the following code:

max_epochs = 20
val_interval = 1
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = []
metric_values = []
writer = SummaryWriter()

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():
            for val_data in val_loader:
                val_inputs, val_labels = (
                    val_data["image"].to(device),
                    val_data["label"].to(device),
                )
                roi_size = (96, 96, 96)
                sw_batch_size = 4
                val_outputs = sliding_window_inference(val_inputs, roi_size, sw_batch_size, model)
                val_outputs = [post_pred(i) for i in decollate_batch(val_outputs)]
                val_labels = [post_label(i) for i in decollate_batch(val_labels)]
                # 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()
            # 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(), os.path.join("/content/drive/MyDrive/UNet", "best_metric_model.pth"))
                print("saved new best metric model")
            print(
                f"current epoch: {epoch + 1} current mean dice: {metric:.4f}"
                f"\nbest mean dice: {best_metric:.4f} "
                f"at epoch: {best_metric_epoch}"
            )
            writer.add_scalar("val_mean_dice", metric, epoch + 1)

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

And for the validation I have used the following code:

if len(devices) > 1:
    model = torch.nn.DataParallel(model, device_ids=devices)

post_label = Compose([AsDiscrete(to_onehot=3)])
model.eval()
with torch.no_grad():
    for val_data in val_loader:
        val_inputs, val_labels = val_data["image"].to(device), val_data["label"].to(device)
        # define sliding window size and batch size for windows inference
        roi_size = (96, 96, 96)
        sw_batch_size = 4
        val_outputs = sliding_window_inference(val_inputs, roi_size, sw_batch_size, model)
        val_outputs = [post_trans(i) for i in decollate_batch(val_outputs)]
        val_labels = [post_label(i) for i in 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()

How do I solve this issue?

Thanks in advance,
Kiarash

[KumoLiu]

Hi @Kiarashdnsh, thanks for your interest here.

The code looks good to me.
The only thing I noticed is that you should specify the same slice ([0, 0, 80]) for the pred.
plt.imshow(torch.argmax(val_outputs, keepdim=True).detach().cpu()[0, 0, 0])
And may I know your dice metric?
Additionally, I would suggest you refer to the tutorial as the first step. It may help you localize the error.

Hope it helps, thanks!

[Kiarashdnsh]

Hi @KumoLiu,

Thanks for your answer. The mean dice value from a limited run was 0.58 (far higher than expected). I also have downloaded the segmented files from the network, which looks alright. However, when I set the slice to [0, 0, 80] I get an index error:

IndexError: Index 80 is out of bounds for dimension 2 with size 1.

So the last index only accepts 0. I tried changing other indices as well. The first two indices also take only 0. When I change each of them to 1 I get the following error:

IndexError: Index 1 is out of bounds for dimension 0 with size 1.

I'm thinking maybe the problem is related to the roi/spatial size I set in the transforms section of the code:

train_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        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,
            ),
        ToTensord(keys=["image", "label"])
    ]
)
val_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        ToTensord(keys=["image", "label"])
    ]
)

in the training phase:

                ...
                roi_size = (96, 96, 96)
                sw_batch_size = 4
                ...

and also in the validation phase:

        ...
        roi_size = (96, 96, 96)
        sw_batch_size = 4
        ...

[KumoLiu]

Hi @Kiarashdnsh, could you please try:

post_trans = Compose([AsDiscrete(argmax=True)])
val_output = post_trans(val_output[0])
plt.imshow(val_output[i, :, :, 80].detach().cpu())

The error seems that you missed the batch dim.

Hope it helps, thanks!

[Kiarashdnsh]

Hi @KumoLiu. Thanks again for your input.
I used many variations of the code you provided. With your help, I managed to solve the issue by changing only the slice-related code line to:

plt.imshow(val_outputs[0, 2, 80].detach().cpu())

now the second index accepts segmented IDs (0, 1, and 2), and the last index is the slice number. For some reason, the : syntaxing doesn't work in my pipeline.

I do appreciate your help!