StanfordMIMI / skm-tea

Repository for the Stanford Knee MRI Multi-Task Evaluation (SKM-TEA) Dataset
MIT License
77 stars 15 forks source link

Compared to what was written in your article, the pre-trained Unet model provided very poor results. #29

Closed beeozfamous closed 1 year ago

beeozfamous commented 1 year ago

Hi, I've been training a segmentation network on the SKM-TEA dataset for a while now. I came across your SKM-TEA tutorial Colab notebook which provides a section to perform segmentation using a pre-trained Unet on a slide of MTR201. However, the segmentation results appear to be very poor. I made some changes to the code to process all the slides of MTR201. When I calculated the DSC (Dice Similarity Coefficient) between the predicted and ground truth masks, the results were disappointing.

FEMORAL DICE SCORE = 0.9056336561510854 TIBIAL DICE SCORE = 0.8161735476184412 MENISCUS DICE SCORE = 0.013693482102011822 PATELLA DICE SCORE = 0.00209708260035654

Could you please explain why the results I computed differ from those presented in your article? I would appreciate it if you could provide an explanation.

I computed the Dice Similarity Coefficient (DSC) using the SKM-TEA tutorial notebook with some modifications made in the Segmentation section. The modifications are given below.

def calculate_dice_coefficient(segmentation, ground_truth,epsilon=1e-6):
    intersection = np.sum(segmentation * ground_truth)
    union = np.sum(segmentation) + np.sum(ground_truth)
    return ((2. * intersection)+epsilon) / (union + epsilon)

    return np.mean((numerator + epsilon) / (denominator + epsilon)) # average over classes and batch
def keep_ya_head_up(source,one):
    one = one + 2
    the = source.copy() + 2
    the[the == one] = 1
    the[the != 1] = 0
    return the

# DICOM data + segmentation
image_file = scan["image_file"]
with h5py.File(image_file, "r") as f:
    echo1 = f["echo1"][()]  # Shape: (x, y, z)
    segmentation = f["seg"][()]  # Shape: (x, y, z, #classes)

echo1 = torch.as_tensor(echo1).unsqueeze(0).unsqueeze(0).float()  # Shape: (B, C, H, W)
gt_seg_sl = segmentation
gt_seg_sl = collect_mask(gt_seg_sl, (0, 1, (2, 3), (4, 5)), out_channel_first=False)
gt_seg_sl = oF.one_hot_to_categorical(gt_seg_sl, channel_dim=-1)
save_nif=np.zeros((512, 512,0))
with torch.no_grad():
    for x in range(160):
        print(x)
        echo1_sl = echo1[..., x]
        logits = model({"image": echo1_sl})["sem_seg_logits"]
        save_nif=np.concatenate((save_nif,oF.pred_to_categorical(logits, activation='sigmoid').cpu().data.numpy()[0][...,np.newaxis]),axis=2)
groundtruth_nif=gt_seg_sl

background = keep_ya_head_up(save_nif,0)
femoral = keep_ya_head_up(save_nif,1)
tibial = keep_ya_head_up(save_nif,2)
meniscus = keep_ya_head_up(save_nif,3)
patella = keep_ya_head_up(save_nif,4)

background_gt = keep_ya_head_up(groundtruth_nif,0)
femoral_gt = keep_ya_head_up(groundtruth_nif,1)
tibial_gt = keep_ya_head_up(groundtruth_nif,2)
meniscus_gt = keep_ya_head_up(groundtruth_nif,3)
patella_gt = keep_ya_head_up(groundtruth_nif,4)

bgds=calculate_dice_coefficient(background,background_gt)
fds=calculate_dice_coefficient(femoral,femoral_gt)
tds=calculate_dice_coefficient(tibial,tibial_gt)
mds=calculate_dice_coefficient(meniscus,meniscus_gt)
pds=calculate_dice_coefficient(patella,patella_gt)
superds=(fds+tds+mds+pds)/4

print("BACKGROUND=DICE=SCORE","=",str(bgds))
print("FEMORAL=DICE=SCORE","=",str(fds))
print("TIBIAL=DICE=SCORE","=",str(tds))
print("MENISCUS=DICE=SCORE","=",str(mds))
print("PATELLA=DICE=SCORE","=",str(pds))
print("SUPER=DICE=SCORE","=",str(superds))
ad12 commented 1 year ago

Hi @beeozfamous - there are two key processing steps that have to happen to the input and the output:

  1. The input need to be preprocessed prior to use with the model - the model was trained with zero-mean unit-standard deviation normalization.
  2. The order of segmentation categories in the model outputs are not the same as the ground truth segmentation. the output segmentations have to be reordered

I have a minimum reproducible example this colab

beeozfamous commented 1 year ago

Thank you for your response. You explained your answer in great detail.