Interpretability of DL model for 3D medical image segmentation

[Yangtting]

Hello, I noticed that you post an interpretability tutorial of 3D densenet for lung lesion classification from CT image patches. And I want to used Grad-CAM like this tutorial on my segmentation project followed by https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/spleen_segmentation_3d.ipynb. However, many problems occured, here is my code, could you please help me?

'''
path
'''
root_dir=r"C:\Now research\Auto segmentation\cal_ar"
data_dir=os.path.join(root_dir,"originaldataset")
data_dir=os.path.join(data_dir,"Automatic segment cases")

'''
image preprocessing
'''
val_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        Spacingd(
            keys=["image", "label"],
            pixdim=(1, 1, 1),
            mode=("bilinear", "nearest"),
        ),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        ScaleIntensityRanged(
            keys=["image"], a_min=100, a_max=800, b_min=0.0, b_max=1.0, clip=True
        ),
        CropForegroundd(keys=["image", "label"], source_key="image"),
        ToTensord(keys=["image", "label"]),
    ]
)

'''
load data
'''
train_images = sorted(
    glob.glob(os.path.join(data_dir, "images_1_2(medianfilter)", "*.nii.gz")))
train_labels = sorted(
    glob.glob(os.path.join(data_dir, "labels_1_2", "*.nii.gz")))
data_dicts = [
    {"image": image_name, "label": label_name}
    for image_name, label_name in zip(train_images, train_labels)
]
test_files =  data_dicts[81:82]

test_org_ds = CacheDataset(data=test_files, transform=val_transforms, cache_rate=1.0, num_workers=0)
test_org_loader = DataLoader(test_org_ds, batch_size=1, num_workers=0)

'''
test on the trained model
'''
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

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

for name, _ in model.named_modules(): print(name)

save_path=os.path.join(root_dir,"20230108")
save_path=os.path.join(save_path,"bestmodelsave")
save_path=os.path.join(save_path,"all") 
save_path=os.path.join(save_path,"3DresUNET") 
save_path=os.path.join(save_path,"1") 

model.load_state_dict(torch.load(
    os.path.join(save_path, "best_metric_model.pth")))
model.eval()

'''
interpreability
'''
cam = monai.visualize.GradCAM(nn_module=model, target_layers="model.2.1")
win_size=(32,32,32) 
print(
    "original feature shape",
    cam.feature_map_size([1, 1] + list(win_size), device),
)
print("upsampled feature shape", [1, 1] + list(win_size))

occ_sens = monai.visualize.OcclusionSensitivity(nn_module=model, mask_size=12, n_batch=1)

the_slice = test_org_ds[0]["image"].shape[-1] // 2
occ_sens_b_box = [-1, -1, -1, -1, the_slice - 1, the_slice]


with torch.no_grad():
    for i, batch in enumerate(test_org_loader):
        image, label = (batch["image"].cuda(), batch["label"].cuda())
        img_name = batch["image_meta_dict"]["filename_or_obj"][0].split("/")[-1]
        print("Inference on case {}".format(img_name))
        outputs = sliding_window_inference(image, (128,128,96), 4, model, overlap=0.8)
    
        pred_outputs = outputs.argmax(1).item()

    
        img = image.detach().cpu().numpy()[..., the_slice]
    
        # name = "actual: "
        # name += "lesion" if label == 1 else "non-lesion"
        # name += "\npred: "
        # name += "lesion" if pred_label == 1 else "non-lesion"
        # name += f"\nlesion: {y_pred[0,1]:.3}"
        # name += f"\nnon-lesion: {y_pred[0,0]:.3}"

        # run CAM
        cam_result = cam(x=image, class_idx=None)
        cam_result = cam_result[..., the_slice]
    
        # run occlusion
        occ_result, _ = occ_sens(x=image, b_box=occ_sens_b_box)
        occ_result = occ_result[0, pred_outputs][None, None, ..., -1]
    
        for row, (im, title) in enumerate(
            zip(
                [img, cam_result, occ_result],
                [name, "CAM", "Occ. sens."],
            )
        ):
            cmap = "gray" if row == 0 else "jet"
            ax = axes[row, example]
            if isinstance(im, torch.Tensor):
                im = im.cpu().detach()
            im_show = ax.imshow(im[0][0], cmap=cmap)
    
            ax.set_title(title, fontsize=25)
            ax.axis("off")
            fig.colorbar(im_show, ax=ax)
    
        example += 1
        if example == n_examples:
            break

[KumoLiu]

Hi @Yangtting, I would suggest you refer to these tutorials as the first step.

Hope it can help you, thanks!