Multi class segmentation of prostate images

[Muhammad-Raza156]

Hi , i am segmenting prostate mri images , my dataset consists of 120 training and 20 validation samples with each image having different size. It genrates poor results some times infinite also I am unable to resolve this i have already searched for similar discussions

this is my transformation code

class ConvertToMultiChannelBasedOnPancreaClassesd(MapTransform):
"""
Convert labels to multi channels:
label 0 is Background
label 1 is CZ(Central Zone)
label 2 is PZ(Peripheral Zone)
"""

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

train_transforms = Compose(

[
    LoadImaged(keys=['image', 'label']),
    EnsureChannelFirstd(keys=['image', 'label']),
    Spacingd(keys=['image', 'label'], pixdim=(0.5, 0.5, 0.5)),
    Orientationd(keys=['image', 'label'],axcodes='RAS'),
    ScaleIntensityd(keys='image',minv=0,maxv=1),
    NormalizeIntensityd(keys='image'),
    RandBiasFieldd(keys='image', degree=10, coeff_range=[0.0,0.01], prob=prob),
    RandGaussianSmoothd(keys='image',sigma_x=[0.25, 1.5], sigma_y=[0.25, 1.5], sigma_z=[0.25, 1.5],prob=prob),
    RandGibbsNoised(keys='image', alpha=[0.5,1], prob=prob),
    RandAffined(
            keys=['image', 'label'],rotate_range=5, shear_range=0.5, translate_range=25,mode=['bilinear' , 'nearest'],prob=prob),
    RandRotate90d(keys=['image', 'label'],spatial_axes=[0,1],prob=prob),
    RandRotated(keys=['image', 'label'],range_x=0.1, range_y=0.1, range_z=0.1,mode=['bilinear','nearest'],prob=prob),
    RandElasticd(keys=['image', 'label'],sigma_range= [0.5,1.5],magnitude_range= [0.5,1.5],rotate_range=5,shear_range=0.5,translate_range=25,
                 mode=['bilinear' , 'nearest'],prob=prob),
    #RandZoomd(keys=['image', 'label'],min_zoom=0.9,max_zoom=1.1,mode=['bilinear','nearest'],prob=prob),
    RandCropByPosNegLabeld(keys=['image', 'label'],label_key='label',spatial_size=[64, 64, 64],pos=2,neg=1,num_samples=8,
            image_key='image',image_threshold=0,),
    SqueezeDimd(keys=["label"],dim=0),
    ConvertToMultiChannelBasedOnPancreaClassesd(keys="label"),
    RandGaussianNoised(keys='image',mean=0.1,std=0.25,prob=prob),
    RandShiftIntensityd(keys='image',offsets=0.2,prob=prob),
    RandGaussianSharpend(
            keys='image',
            sigma1_x=[0.5, 1.0],sigma1_y=[0.5, 1.0],sigma1_z=[0.5, 1.0],
            sigma2_x=[0.5, 1.0],sigma2_y=[0.5, 1.0],sigma2_z=[0.5, 1.0],
            alpha=[10.0, 30.0],prob=prob),
    RandAdjustContrastd(keys='image',gamma=2.0,prob=prob),
    #CropForegroundd(keys=['image', 'label'], source_key='image'),
    EnsureTyped(keys=['image', 'label']),
    ToTensord(keys=['image', 'label'])
]

)

val_transforms = Compose(

[
    LoadImaged(keys=['image', 'label']),
    EnsureChannelFirstd(keys=['image', 'label']),
    Spacingd(keys=['image', 'label'], pixdim=(0.5, 0.5, 0.5)),
    Orientationd(keys=['image', 'label'],axcodes='RAS'),
    SqueezeDimd(keys=["label"],dim=0),
    ConvertToMultiChannelBasedOnPancreaClassesd(keys="label"),
    ScaleIntensityd(keys='image',minv=0,maxv=1),
    NormalizeIntensityd(keys='image'),
    EnsureTyped(keys=['image', 'label']),
    ToTensord(keys=['image', 'label'])
]

)

This is my rest of the code

model= build_unet().to(device)
loss_function = DiceLoss(squared_pred=True, to_onehot_y=False, sigmoid=True)
optimizer= torch.optim.Adam(model.parameters(), lr=0.001)
​
dice_metric = DiceMetric(include_background=True, reduction="mean")
dice_metric_batch = DiceMetric(include_background=True, reduction="mean_batch")
​
hd95_metric = HausdorffDistanceMetric(include_background=True, percentile=95, reduction="mean")
hd95_metric_batch = HausdorffDistanceMetric(include_background=True, percentile=95, reduction="mean_batch")
​
sd_metric = SurfaceDistanceMetric(include_background=True, symmetric=True, reduction="mean")
sd_metric_batch = SurfaceDistanceMetric(include_background=True, symmetric=True, reduction="mean_batch")
​
post_trans = Compose(
[EnsureType(), Activations(sigmoid=True), AsDiscrete(threshod=0.5)]
)

max_epochs = 40
val_interval = 2
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = []

cor_hd=0
cor_asd=0

metric_values = []
metric_values_CZ = []
metric_values_PZ = []

hd_metric_values = []
hd_metric_values_CZ = []
hd_metric_values_PZ = []

ssd_metric_values = []
ssd_metric_values_CZ = []
ssd_metric_values_PZ = []

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()
with autocast():
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
            overlap=0.5
            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)]
            
            # compute metric for current iteration
            dice_metric(y_pred=val_outputs, y=val_labels)
            dice_metric_batch(y_pred=val_outputs, y=val_labels)
            
            hd95_metric(y_pred=val_outputs, y=val_labels)
            hd95_metric_batch(y_pred=val_outputs, y=val_labels)
            
            sd_metric(y_pred=val_outputs, y=val_labels)
            sd_metric_batch(y_pred=val_outputs, y=val_labels)
            

        # aggregate the final mean dice result
        metric = dice_metric.aggregate().item()
        metric_values.append(metric)
        
        metric_batch = dice_metric_batch.aggregate()
        
        
        metric_CZ = metric_batch[1].item()
        metric_values_CZ.append(metric_CZ)
        metric_PZ = metric_batch[2].item()
        metric_values_PZ.append(metric_PZ)

        # reset the status for next validation round
        dice_metric.reset()
        dice_metric_batch.reset()
        
        hd_metric = hd95_metric.aggregate().item()
        hd_metric_values.append(hd_metric)
        hd_metric_batch = hd95_metric_batch.aggregate()
        
        hd_metric_CZ = hd_metric_batch[1].item()
        hd_metric_values_CZ.append(hd_metric_CZ)
        hd_metric_PZ = hd_metric_batch[2].item()
        hd_metric_values_PZ.append(hd_metric_PZ)
        hd95_metric.reset()
        hd95_metric_batch.reset()
        
        
        ssd_metric = sd_metric.aggregate().item()
        ssd_metric_values.append(ssd_metric)
        ssd_metric_batch = sd_metric_batch.aggregate()
        
        
        ssd_metric_CZ = ssd_metric_batch[1].item()
        ssd_metric_values_CZ.append(ssd_metric_CZ)
        ssd_metric_PZ = ssd_metric_batch[2].item()
        ssd_metric_values_PZ.append(ssd_metric_PZ)
        sd_metric.reset()
        sd_metric_batch.reset()
        
        

        metric_values.append(metric)
        if metric > best_metric:
            best_metric = metric
            cor_hd= hd_metric
            cor_asd=ssd_metric
            best_metric_epoch = epoch + 1
            torch.save(model.state_dict(), os.path.join(
                "/kaggle/working/", "best_metric_model.pth"))
            print("saved new best metric model")
        print(
            f"current epoch: {epoch + 1} current mean dice: {metric:.4f}"
            f"\n CZ: {metric_CZ:.4f} PZ: {metric_PZ:.4f}"
            f"\nbest mean dice: {best_metric:.4f} "
            f"at epoch: {best_metric_epoch}"
        )
        
        print(
            f"current mean HD: {hd_metric:.4f}"
            f"\n CZ: {hd_metric_CZ:.4f} PZ: {hd_metric_PZ:.4f}"
        )
        
        print(
            f"current mean ASD: {ssd_metric:.4f}"
            f"\n CZ: {ssd_metric_CZ:.4f} PZ: {ssd_metric_PZ:.4f}"
        )

[Muhammad-Raza156]

hi, @Nic-Ma would be thankful

[Nic-Ma]

Hi @yiheng-wang-nv ,

Could you please help take a look at this question?

Thanks in advance.

[yiheng-wang-nv]

Hi @Muhammad-Raza156 , I'm curious about what is the performance on the training set?
If it's good, it means the network is overfitting. Otherwise, it's underfitting. could you provide some more details about:

1. training metrics, 2) validation metrics

[yiheng-wang-nv]

You used sigmoid activation, so could you double confirm that this task is a multi-label segmentation task? If not, you can change it into using argmax (also for the loss function).
Thanks for the updates of the validation scores. How about the training scores? If not logged, maybe you can calculate with the trained weights.

[yiheng-wang-nv]

My last response is based on your first two messages.

I see, according to your attached figures, the task is a 3-class segmentation task rather than a multi-label segmentation task. I think sigmoid should not be used.

[Muhammad-Raza156]

Thanks for the response so u suggest to use "softmax" then i have to set to_one_hot=3 and argmax will be true is that true ?
what about include_background argument ?

[Muhammad-Raza156]

train loss at 1st epoch was 0.6059
and at epoch 31 it is 0.5972

[yiheng-wang-nv]

what is the output shape of your network? using include_background or not depends on your purpose.