3D UNETR model

[abhishekpatil32]

I am very new to MONAI and using the 3D UNETR model to train the model and to segment a small region of interest in a CT scan.

What changes are required in the code to run the unetr_btcv_segmentation_3d.ipynb code ??
link: [https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/unetr_btcv_segmentation_3d.ipynb]

I did not find any tutorials for binary segmentation.

I really appreciate any help you can provide.

[KumoLiu]

Hi @abhishekpatil32, thanks for your interest here.

You can refer to the spleen segmentation here which is about binary segmentation: https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/spleen_segmentation_3d.ipynb
You need adjust the output channel, activation function and loss accordingly.

Hope it helps,
Thanks.

[abhishekpatil32]

Thank you for your answer, @KumoLiu . I worked around the issue and successfully ran the UNETR model with the parameters mentioned above. However, the model's performance is poor, with a very low mean Dice score (~0.05).

Could the issue be related to the input data? My target area for segmentation is very small (a 15 mm cube), and the input data consists of CT images and labels (both in .nii.gz format). I am wondering if preprocessing the labels is necessary before training.

Please suggest measures to improve the Dice score.
Thank you in advance

[abhishekpatil32]

Hello @KumoLiu, when I try to run the script for visualization named /unetr_btcv_segmentation_3d_lightning.ipynb

slice_map = {
    "img0035.nii.gz": 170,
    "img0036.nii.gz": 230,
    "img0037.nii.gz": 204,
    "img0038.nii.gz": 204,
    "img0039.nii.gz": 204,
    "img0040.nii.gz": 180,
}
case_num = 4
net.load_from_checkpoint(os.path.join(root_dir, "best_metric_model.ckpt"))
net.eval()
net.to(device)

with torch.no_grad():
    img_name = os.path.split(net.val_ds[case_num]["image"].meta["filename_or_obj"])[1]
    img = net.val_ds[case_num]["image"]
    label = net.val_ds[case_num]["label"]
    val_inputs = torch.unsqueeze(img, 1).cuda()
    val_labels = torch.unsqueeze(label, 1).cuda()
    val_outputs = sliding_window_inference(val_inputs, (96, 96, 96), 4, net, overlap=0.8)
    plt.figure("check", (18, 6))
    plt.subplot(1, 3, 1)
    plt.title("image")
    plt.imshow(val_inputs.cpu().numpy()[0, 0, :, :, slice_map[img_name]], cmap="gray")
    plt.subplot(1, 3, 2)
    plt.title("label")
    plt.imshow(val_labels.cpu().numpy()[0, 0, :, :, slice_map[img_name]])
    plt.subplot(1, 3, 3)
    plt.title("output")
    plt.imshow(torch.argmax(val_outputs, dim=1).detach().cpu()[0, :, :, slice_map[img_name]])
    plt.show()

I get a strange error
AttributeError: 'Net' object has no attribute 'val_ds'

Has anyone resolved this ?

[KumoLiu]

The error AttributeError: 'Net' object has no attribute 'val_ds' indicates that Python is trying to access an attribute named val_ds on an object of type Net, but the Net class (or its instances) does not have an attribute with that name. So if you didn't define val_ds in your net, please check your code.

[abhishekpatil32]

Hello @KumoLiu, I am using UNETR model for segmenting a region of interest with 100 training, 28 validation and 32 testing samples. Also our image size is (512, 512, 300-400) and labels have same size but only 15 slices have foreground, all other have 0 in the background and out HU intensity is restricted to -190 to -30

For this work I am using the following code:

from monai.transforms import SpatialPadd

train_transforms = Compose([
    LoadImaged(keys=["image", "label"]),
    EnsureChannelFirstd(keys=["image", "label"]),
    Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 1.0), mode=("bilinear", "nearest")),
    ScaleIntensityRanged(keys=["image"], a_min=-190, a_max=-30, b_min=0.0, b_max=1.0, clip=True),
    CropForegroundd(keys=["image", "label"], source_key="label"),    
    SpatialPadd(keys=["image", "label"], spatial_size=(128, 128, 128)),
    
    RandCropByPosNegLabeld(
        keys=["image", "label"],
        label_key="label",
        spatial_size=(128, 128, 128),
        pos=1, neg=1,
        num_samples=1,
        allow_smaller=True
    ),
    
    RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=[0, 1, 2]),
    RandRotate90d(keys=["image", "label"], prob=0.5, spatial_axes=[0, 1]),
    RandGaussianNoised(keys=["image"], prob=0.1, mean=0, std=0.1),
    RandAdjustContrastd(keys=["image"], prob=0.2, gamma=(0.9, 1.1)),
    ToTensorD(keys=["image", "label"])
])

val_transforms = Compose([
    LoadImaged(keys=["image", "label"]),
    EnsureChannelFirstd(keys=["image", "label"]),
    Orientationd(keys=["image", "label"], axcodes="RAS"),
    Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 1.0), mode=("bilinear", "nearest")),
    ScaleIntensityRanged(keys=["image"], a_min=-190, a_max=-30, b_min=0.0, b_max=1.0, clip=True),
    CropForegroundd(keys=["image", "label"], source_key="label"),  # Ensure consistency with training
    SpatialPadd(keys=["image", "label"], spatial_size=(128, 128, 128)),  # Prevent size mismatch
    ToTensorD(keys=["image", "label"])
])

set_determinism(seed=0)

data_dir = "/content/data/"
split_json = "dataset_0.json"

datasets = data_dir + split_json
datalist = load_decathlon_datalist(datasets, True, "training")
val_files = load_decathlon_datalist(datasets, True, "validation")
train_ds = CacheDataset(
    data=datalist,
    transform=train_transforms,
    cache_num=24,
    cache_rate=1.0,
    num_workers=8,
)

train_loader = DataLoader(train_ds, batch_size=1, shuffle=True, num_workers=8, pin_memory=True)
val_ds = CacheDataset(data=val_files, transform=val_transforms, cache_num=6, cache_rate=1.0, num_workers=4)
val_loader = DataLoader(val_ds, batch_size=1, shuffle=False, num_workers=4, pin_memory=True)

from torch.optim.lr_scheduler import CosineAnnealingLR, OneCycleLR, CosineAnnealingWarmRestarts

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

device0 = torch.device("cuda:0")

model = UNETR(
    in_channels=1,
    out_channels=2,
    img_size=(128, 128, 128),
    feature_size=16,
    hidden_size=768,
    mlp_dim=3072,
    num_heads=12,
    proj_type="perceptron",
    norm_name="instance",
    res_block=True,
    dropout_rate=0.2,
).to(device0)


loss_function = TverskyLoss(to_onehot_y=True, softmax=True, alpha=0.3, beta=0.7)
torch.backends.cudnn.benchmark = True
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-4, weight_decay=1e-5)
scheduler = CosineAnnealingLR(optimizer, T_max=1000, eta_min=1e-6)

def validation(epoch_iterator_val):
    model.eval()
    with torch.no_grad():
        for batch in epoch_iterator_val:
            val_inputs, val_labels = (batch["image"].to(device0), batch["label"].to(device0))
            #val_outputs = sliding_window_inference(val_inputs, (96, 96, 96), 8, model)

            val_outputs = sliding_window_inference(val_inputs, (128, 128, 128), 4, model, overlap=0.25)
            val_labels_list = decollate_batch(val_labels)
            val_labels_convert = [post_label(val_label_tensor) for val_label_tensor in val_labels_list]
            val_outputs_list = decollate_batch(val_outputs)
            val_output_convert = [post_pred(val_pred_tensor) for val_pred_tensor in val_outputs_list]
            dice_metric(y_pred=val_output_convert, y=val_labels_convert)
            epoch_iterator_val.set_description("Validate (%d / %d Steps)" % (global_step, 10.0))  # noqa: B038
        mean_dice_val = dice_metric.aggregate().item()
        dice_metric.reset()
    return mean_dice_val

def train(global_step, train_loader, dice_val_best, global_step_best):
    model.train()
    epoch_loss = 0
    step = 0
    epoch_iterator = tqdm(train_loader, desc="Training (X / X Steps) (loss=X.X)", dynamic_ncols=True)
    for step, batch in enumerate(epoch_iterator):
        step += 1
        x, y = (batch["image"].to(device0), batch["label"].to(device0))
        logit_map = model(x)
        loss = loss_function(logit_map, y)
        loss.backward()
        epoch_loss += loss.item()
        optimizer.step()
        optimizer.zero_grad()
        epoch_iterator.set_description(  # noqa: B038
            "Training (%d / %d Steps) (loss=%2.5f)" % (global_step, max_iterations, loss)
        )
        if (global_step % eval_num == 0 and global_step != 0) or global_step == max_iterations:
            epoch_iterator_val = tqdm(val_loader, desc="Validate (X / X Steps) (dice=X.X)", dynamic_ncols=True)
            dice_val = validation(epoch_iterator_val)
            epoch_loss /= step
            epoch_loss_values.append(epoch_loss)
            metric_values.append(dice_val)
            if dice_val > dice_val_best:
                dice_val_best = dice_val
                global_step_best = global_step
                torch.save(model.state_dict(), os.path.join(root_dir, "best_metric_model.pth"))
                print(
                    "Model Was Saved ! Current Best Avg. Dice: {} Current Avg. Dice: {}".format(dice_val_best, dice_val)
                )
            else:
                print(
                    "Model Was Not Saved ! Current Best Avg. Dice: {} Current Avg. Dice: {}".format(
                        dice_val_best, dice_val
                    )
                )
        global_step += 1
    return global_step, dice_val_best, global_step_best


max_iterations = 2000
eval_num = 500
post_label = AsDiscrete(to_onehot=2)
post_pred = AsDiscrete(argmax=True, to_onehot=2)
#post_pred = lambda x: torch.softmax(x, dim=1).argmax(dim=1, keepdim=True)

dice_metric = DiceMetric(include_background=True, reduction="mean")
global_step = 0
dice_val_best = 0.0
global_step_best = 0
epoch_loss_values = []
metric_values = []
while global_step < max_iterations:
    global_step, dice_val_best, global_step_best = train(global_step, train_loader, dice_val_best, global_step_best)
model.load_state_dict(torch.load(os.path.join(root_dir, "best_metric_model.pth")))

For this code, I am able to achieve Dice similarity to be 0.85. But when I infer the model for the test set, I do not get proper segmentations (I applied the same transformations to the test set that were applied to the training still I get improper segmentation on the test set). Also I am not sure how to perform inference on the test dataset.
Any suggestions on this would be helpful