Save the predicted label in nifti format but shows nothing on the 3D Slicer

[mikami520]

Issue

Hi all, I am working on the Deep Atlas. When I make inference on the segmentation network, I firstly convert the predicted label from 3D numpy array into nii.gz format and then save them to the output folder. However, when I open the 3D Slicer to check labels, some show nothing but some can show the segmentation. When I use the np.unique(predict) to get number of labels, those who cannot show the segmentation on Slicer have only [0] which means they only have background. The prediction code comes from Deep Atlas tutorial, does anyone has any suggestions, thanks in advance!

Code script

import monai
import torch
import itk
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
import os
import nibabel as nib
import sys
import json

def seg_inference(seg_net, device, model_path, json_path, output_path):
    json_file = load_json(json_path)
    raw_data = json_file['total_test'] # get data list of test data
    headers, affines, ids, num_labels = get_nii_info(raw_data) # get header and affine of ground truth
    seg_net.load_state_dict(torch.load(model_path)) # load best model
    seg_net.to(device)
    seg_net.eval() # change to evaluation condition
    dice_loss = dice_loss_func2() # load dice loss function
    data = load_seg_dataset(raw_data) # load test data
    k = 0
    eval_losses = []
    for i in data:
        header1 = headers[k]
        affine1 = affines[k]
        id = ids[k]
        data_item = i
        test_input = data_item['img']
        test_gt = data_item['seg']
        with torch.no_grad():
            test_seg_predicted = seg_net(test_input.unsqueeze(0).cuda()).cpu() # get segmentation prediction
            loss = dice_loss(test_seg_predicted, test_gt.unsqueeze(0)).item()

        eval_loss = f"Scan ID: {id}, dice loss: {loss}"
        eval_losses.append(eval_loss)
        prediction = torch.argmax(torch.softmax(
            test_seg_predicted, dim=1), dim=1, keepdim=True)[0, 0] # softmax prediction
        k += 1
        pred_np = prediction.detach().cpu().numpy() # convert from torch to numpy
        pred_np = pred_np.astype('int16') # change data type to int16
        nii = nib.Nifti1Image(pred_np, affine=affine1, header=header1) # convert to nifti format
        nii.header.get_xyzt_units()
        nib.save(nii, (os.path.join(output_path, id + '.nii.gz')))

def get_nii_info(data):
    headers = []
    affines = []
    ids = []
    for i in range(len(data)):
        item = data[i]
        id = os.path.basename(item['seg']).split('.')[0]
        seg = nib.load(item['seg'])
        num_labels = len(np.unique(seg.get_fdata()))
        headers.append(seg.header)
        affines.append(seg.affine)
        ids.append(id)
    
    return headers, affines, ids, num_labels

def load_seg_dataset(data_list):
    transform_seg_available = monai.transforms.Compose(
        transforms=[
            monai.transforms.LoadImageD(keys=['img', 'seg'], image_only=True),
            monai.transforms.AddChannelD(keys=['img', 'seg']),
            monai.transforms.ToTensorD(keys=['img', 'seg'])
        ]
    )
    itk.ProcessObject.SetGlobalWarningDisplay(False)
    dataset_seg_available_train = monai.data.CacheDataset(
        data=data_list,
        transform=transform_seg_available,
        cache_num=16,
        hash_as_key=True
    )
    return dataset_seg_available_train

[ebrahimebrahim]

When I use the np.unique(predict) to get number of labels, those who cannot show the segmentation on Slicer have only [0] which means they only have background

Are you seeing this [0] after loading the nifti file that you saved, or is this with the original output of the segmentation network? It's not clear to me whether you are getting blank segmentations from seg_net or whether you are just having issues with saving and loading into slicer.

[ebrahimebrahim]

The seg_net pretraining is just a simple unet trained using Dice loss; this is not necessarily any issue with the DeepAtlas setup.

So the blank segmentation issue happens only with some test cases but not others? Could those specific test cases be very different from what is seen during training?

[mikami520]

The seg_net pretraining is just a simple unet trained using Dice loss; this is not necessarily any issue with the DeepAtlas setup.

So the blank segmentation issue happens only with some test cases but not others? Could those specific test cases be very different from what is seen during training?

I checked both prediction results of deep atlas and pretrained segmentation network, they failed at exactly same test files. Also, I checked those specific cases, they are normal files and not much different other from test cases. But I do not know how to check if it is different from what is seen during training, do you have some suggestions?

[mikami520]

The seg_net pretraining is just a simple unet trained using Dice loss; this is not necessarily any issue with the DeepAtlas setup.

So the blank segmentation issue happens only with some test cases but not others? Could those specific test cases be very different from what is seen during training?

by the way, I have a question about registration network training. For example, if the dataset has scans with dimension 256x256x256, if one specific area appears at No.200 slice in one image but appears at No.250 in another image, will this affect the training? If all the test cases have unmatched scene in each slice, how will the training performance be?

[mikami520]

The seg_net pretraining is just a simple unet trained using Dice loss; this is not necessarily any issue with the DeepAtlas setup.

So the blank segmentation issue happens only with some test cases but not others? Could those specific test cases be very different from what is seen during training?

I found an interesting thing, I used same network and training parameters to work on the pre-trained segmentation network on Colab, those failed cases now has label back, I thought it is the problem of my computer, what do you think?

[ebrahimebrahim]

by the way, I have a question about registration network training. For example, if the dataset has scans with dimension 256x256x256, if one specific area appears at No.200 slice in one image but appears at No.250 in another image, will this affect the training? If all the test cases have unmatched scene in each slice, how will the training performance be?

CNNs do not care if a feature is translated around the image, so this should be fine.

I found an interesting thing, I used same network and training parameters to work on the pre-trained segmentation network on Colab, those failed cases now has label back, I thought it is the problem of my computer, what do you think?

Hmm, weird. No idea how that happens; it sounds like you have some debugging to do on your machine to find out which part is breaking 🙂