2D Slice Segmentation with UNet from 3D Volume

[Spawnfile]

Hi all,

I've got a bunch of CT volumes like 512x512x301 etc. which I need to train. Despite the all my research I couldn't find any implementation of methods in order to do that. How can I train a 2D ( CT slice ) segmentation with UNet from 3D CT volume ?

Cheers

[KumoLiu]

Hi @Spawnfile, here are some tutorials you could refer to:
https://github.com/Project-MONAI/tutorials/blob/main/modules/2d_slices_from_3d_training.ipynb
https://github.com/Project-MONAI/tutorials/tree/main/modules/generate_random_permutations
https://github.com/Project-MONAI/tutorials/blob/main/modules/2d_slices_from_3d_sampling.ipynb

These are written in terms of sampling 2D from 3D, but they can generate 3D samples as well.
Hope it helps, thanks!

[Spawnfile]

Hello,

https://github.com/Project-MONAI/tutorials/blob/main/modules/2d_slices_from_3d_training.ipynb this was the tutorial I followed on my training task.

I'm using a public CT dataset from kaggle for the training with MonAI. These are the transforms I'm using on patching and loading image data from nifti volume.

monai_transforms = Compose([
        LoadImaged(keys=["img", "seg"]),
        AddChanneld(keys=["img", "seg"]),
        CropForegroundd(keys=["img", "seg"], source_key="img"),
        ScaleIntensityRanged(keys="img", a_min=-1021, a_max=147, b_min=0, b_max=1),  
        ])

patch_transform = Compose(
    [
        SqueezeDimd(keys=["img", "seg"], dim=3),
        Resized(keys=["img", "seg"], spatial_size=[256, 256]),
    ])

However, when I apply Resized on patching section, image and masks voxel values calculated wrongly which are image interval as min->0 and max-> 2.12 (should be 1) and mask values are tensor([0.0000, 0.1667, 0.2500, 0.3333, 0.5000, 0.6667, 0.7500, 0.8333, 1.0000]).

I need to segmentate the covid area on images but results on test images the DICE metric does not increase properly with above transforms and patching methods after the training despite decreasing dice loss.

Lastly, this is the image I'm getting after all transforms have been done.

---

Full training code is written below to examine.

# loading val and train files

train_files = [{"img": ct_scans_path + img, "seg": lung_seg_path + seg} for img, seg in zip(images[:3], lung_masks[:3])]
val_files = [{"img": ct_scans_path + img, "seg": lung_seg_path + seg} for img, seg in zip(images[14:16], lung_masks[14:16])]

monai_transforms = Compose([
        LoadImaged(keys=["img", "seg"]),
        AddChanneld(keys=["img", "seg"]),
        CropForegroundd(keys=["img", "seg"], source_key="img"),
        ScaleIntensityRanged(keys="img", a_min=-1021, a_max=147, b_min=0, b_max=1),
    ])
                                 
volume_ds = monai.data.CacheDataset(data=train_files, transform=monai_transforms)
patch_func = monai.data.PatchIterd(
    keys=["img", "seg"], patch_size=(None, None, 1), start_pos=(0, 0, 0))

patch_transform = Compose(
    [
        SqueezeDimd(keys=["img", "seg"], dim=3),
        Resized(keys=["img", "seg"], spatial_size=[256, 256]),
    ])

patch_ds = monai.data.GridPatchDataset(data=volume_ds, patch_iter=patch_func, transform=patch_transform, with_coordinates=False)
shuffle_ds = monai.data.ShuffleBuffer(patch_ds, buffer_size=200, seed=0)

train_loader = DataLoader(
    shuffle_ds,
    batch_size=5,
    num_workers=2,
    pin_memory=torch.cuda.is_available())

check_data = monai.utils.misc.first(train_loader)
print("first patch's shape: ", check_data["img"].shape, check_data["seg"].shape)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = monai.networks.nets.UNet(
    spatial_dims=2,
    in_channels=1,
    out_channels=1,
    channels=(16, 32, 64, 128),
    strides=(2, 2, 2),
    num_res_units=2,
).to(device)

loss_function = monai.losses.DiceLoss(sigmoid=True)
optimizer = torch.optim.Adam(model.parameters(), 5e-3)

epoch_loss_values = []
num_epochs = 15
for epoch in range(num_epochs):
    print("-" * 10)
    print(f"epoch {epoch + 1}/{num_epochs}")
    model.train()
    epoch_loss, step = 0, 0
    for batch_data in train_loader:
        step += 1
        inputs, labels = batch_data["img"].to(device), batch_data["seg"].to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
        if step % 200 == 0:
            print(f"{step}, train_loss: {loss.item():.4f}")
    epoch_loss /= step
    epoch_loss_values.append(epoch_loss)
    print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")
print("train completed")

Any help is appreciated.
Alper

[Spawnfile]

Hello @KumoLiu ,

This is the code I've been using for check intervals before patch transforms

monai_transforms = Compose([
        LoadImaged(keys=["img", "seg"]),
        AddChanneld(keys=["img", "seg"]),
        CropForegroundd(keys=["img", "seg"], source_key="img"),
        ScaleIntensityRanged(keys="img", a_min=-1021, a_max=147, b_min=0, b_max=1),
    ])

index = 2
volume_ds = monai.data.CacheDataset(data=train_files, transform=monai_transforms)
check_loader = DataLoader(volume_ds, batch_size=1)
check_data = monai.utils.misc.first(check_loader)
print("first volume's shape: ", check_data["img"].shape, check_data["seg"].shape)

print(f"Image Min : {torch.min(check_data['img'][0, 0, :, :, index])} Torch Max : {torch.max(check_data['img'][0, 0, :, :, index])}--- Mask Uniques : {torch.unique(check_data['seg'][0,0,:,:, index])}")

This is the output of code above

After applying PatchTtransforms which is located below

patch_func = monai.data.PatchIterd(
    keys=["img", "seg"], patch_size=(None, None, 1), start_pos=(0, 0, 0))

patch_transform = Compose(
    [
        SqueezeDimd(keys=["img", "seg"], dim=3),
        Resized(keys=["img", "seg"], spatial_size=[256, 256]),
    ])

patch_ds = monai.data.GridPatchDataset(data=volume_ds, patch_iter=patch_func, transform=patch_transform, with_coordinates=False)
shuffle_ds = monai.data.ShuffleBuffer(patch_ds, buffer_size=200, seed=0)

train_loader = DataLoader(
    shuffle_ds,
    batch_size=5,
    num_workers=2,
    pin_memory=torch.cuda.is_available())

check_data = monai.utils.misc.first(train_loader)
print("first patch's shape: ", check_data["img"].shape, check_data["seg"].shape)
image, label = (check_data["img"][index,0,:,:], check_data["seg"][index,0,:,:])
print(f"Image Uniques : {torch.unique(check_data['img'][index, 0, :, :])} --- Mask Uniques : {torch.unique(check_data['seg'][index,0,:,:])}")

I get

After applying patch transforms mask and image values corrupt. Is there any thing I did wrong or missed something ?
Cheers!

[KumoLiu]

It looks like since you shuffle the data, you're not checking the same data.

[Spawnfile]

But in any case, shouldn't be the mask data as [0, 1] ?

Edit 2
Resized(keys=["img", "seg"], spatial_size=[128, 128]) Changing the spatial_size here is manipulating the interval of mask range.

if spatial_size 128x128 mask interval range is 7 tensor([0.0000, 0.0833, 0.1667, 0.2500, 0.4167, 0.6667, 1.0000])
if spatial_size 64x64 mask interval range is 3 tensor([0.0000, 0.1071, 1.0000])

I do not understand the relation between these anyway.

[KumoLiu]

Hi @Spawnfile, yeah that's right. I just missed that. There is an interpolation mode used in Resize which defaults is area. You need to set it to "nearest" for the label.

MONAI/monai/transforms/spatial/array.py

Lines 725 to 726 in fabc198

	mode: {``"nearest"``, ``"nearest-exact"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``}
	The interpolation mode. Defaults to ``"area"``.

Something like:

Resized(keys=["img", "seg"], spatial_size=[128, 128], mode=["bilinear", "nearest"])

Thanks!

[Spawnfile]

This has solved the intensity problem, thank you!