unable to produce appropriate segmentation output

[qwang0225]

Hi, I have been trying to make a 3D binary segmentation model to segment intracranial hemorrhage from CT scans by following the monai tutorials. However, I cannot produce appropriate output even though the mean dice reaches 0.7.

images = sorted(glob.glob(os.path.join("./data/ICH_dataset", "ct_scans", ".nii")))
labels = sorted(glob.glob(os.path.join("./data/ICH_dataset", "masks", ".nii")))
data_dicts = [{"image": image_name, "label": label_name} for image_name, label_name in zip(images, labels)]
train_files, val_files = data_dicts[8:], data_dicts[:8]

train_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        ScaleIntensityRanged(
            keys=["image"],
            a_min=-20,
            a_max=100,
            b_min=0.0,
            b_max=1.0,
            clip=True,
        ),
            # Original label values 0 or 255, convert to 0 or 1 
            ScaleIntensityRanged(
            keys=["label"],
            a_min=0.0, a_max=255.0,
            b_min=0.0, b_max=1.0,
            clip=True,
        ),
        Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 1.0), mode=("bilinear", "nearest")),
        RandCropByPosNegLabeld(
            keys=["image", "label"],
            label_key="label",
            spatial_size=(96, 96, 96),
            pos=3,
            neg=1,
            num_samples=4,
            image_key="image",
            image_threshold=0,
            allow_smaller=True
        ),
     ]
)
val_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        ScaleIntensityRanged(
            keys=["image"],
            a_min=-20,
            a_max=100,
            b_min=0.0,
            b_max=1.0,
            clip=True,
        ),
         ScaleIntensityRanged(
            keys=["label"],
            a_min=0.0, a_max=255.0,
            b_min=0.0, b_max=1.0, 
            clip=True,
        ),
        Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 1.0), mode=("bilinear", "nearest")),
    ]
)

train_ds = CacheDataset(data=train_files, transform=train_transforms, cache_rate=1.0, num_workers=4)
train_loader = DataLoader(train_ds, batch_size=4, shuffle=True, num_workers=4)

val_ds = CacheDataset(data=val_files, transform=val_transforms, cache_rate=1.0, num_workers=4)
val_loader = DataLoader(val_ds, batch_size=1, num_workers=4)

device = torch.device("cuda:0")
model = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
    norm=Norm.BATCH,
).to(device)

weight = torch.tensor([0.2, 0.8]).to(device)
loss_function = DiceLoss(to_onehot_y=True, softmax=True, weight=weight)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=300)
dice_metric = DiceMetric(include_background=True, reduction="mean")
post_pred = Compose([EnsureType(), AsDiscrete(argmax=True, to_onehot=2)])
post_label = Compose([EnsureType(), AsDiscrete(to_onehot=2)])

max_epochs = 300
val_interval = 2
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = []
metric_values = []

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),
        )
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
        #print(outputs.shape)
        #print(outputs[0][0])
        #print(outputs[1][1])
        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 = decollate_batch(val_labels)
                # compute metric for current iteration
                val_outputs = [i.to(device) for i in val_outputs]
                val_labels = [i.to(device) for i in val_labels]
                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(dir, "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}"

model.load_state_dict(torch.load(os.path.join(dir, "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 [:, :, 65]
        for j in range(2):
            plt.figure("check", (18, 6))
            plt.subplot(1, 3, 1)
            plt.title(f"image {i}")
            plt.imshow(val_data["image"][0, 0, :, :, j + 65], cmap="gray")
            plt.subplot(1, 3, 2)
            plt.title(f"label {i}")
            plt.imshow(val_data["label"][0, 0, :, :, j + 65])
            plt.subplot(1, 3, 3)
            plt.title(f"output {i}")
            plt.imshow(torch.argmax(val_outputs, dim=1).detach().cpu()[0, :, :, j+65])
            plt.show()
        
        if i == 2:
            break

Any inputs would be deeply appreciated. Thanks!

[KumoLiu]

Hi @qwang0225, I noticed that you specified include_background=True when calculate the DiceMetric, this may make your results artificially high because you have a large proportion of the background. I would suggest you set it to False and see your real metric.
Thanks!

[qwang0225]

You are right. The dice coefficient became very low after I set dice_metric = DiceMetric(include_background=False, reduction="mean")

The dataset I used is a public dataset link
Any suggestions on how to improve the model? Thanks!

[KumoLiu]

The results could be due to a variety of factors. First and foremost, you should double-check your input. The proportion of the foreground appears to be quite small, and your ROI setting appears to be a little large. You could try 64. Also, different networks and hyperparameters are also worth investigating.
Thanks!