Medical Decathalon: Pancreas Multiclass segementation:

[ramchandracheke]

Hi,
Thank you for making this project open source. I am trying to train U-Net segmentation model for Task07 Pancreas segmentation.

import torch.nn as nn

import monai
from PolyLR import PolynomialLRDecay

from monai.utils import first, set_determinism
from monai.transforms import (
    Activations,
    RandCropByPosNegLabeld,
    AsDiscrete,
    SqueezeDimd,
    CropForegroundd,
    Resized,
    MapTransform,
    LoadImaged,
    ScaleIntensityRanged,
    SpatialPadd,
    Compose,
    LoadImaged,
    MapTransform,
    EnsureChannelFirstd,
    EnsureTyped,
    EnsureType,
    RandSpatialCropd
)

from monai.handlers.utils import from_engine
from monai.networks.layers import Norm
from monai.metrics import DiceMetric
from monai.losses import DiceLoss
from monai.inferers import sliding_window_inference
from monai.data import DataLoader, Dataset, decollate_batch
from monai.config import print_config
from monai.apps import download_and_extract

import time
import torch
import os
import glob
import numpy as np

# standard PyTorch program style: create VNet, DiceLoss and Adam optimizer
from monai.networks.nets import VNet
from monai.networks.nets import UNet

print_config()

data_dir ='/work/Monai/Cropped2/'
root_dir = '/work/Monai/'

train_images = sorted(
    glob.glob(os.path.join(data_dir, "Images", "*.nii.gz")))
train_labels = sorted(
    glob.glob(os.path.join(data_dir, "Label", "*.nii.gz")))
data_dicts = [
    {"image": image_name, "label": label_name}
    for image_name, label_name in zip(train_images, train_labels)
]
train_files, val_files = data_dicts[:-20], data_dicts[-20:]

set_determinism(seed=0)

class ConvertToMultiChannelBasedOnPancreaClassesd(MapTransform):
    """
    Convert labels to multi channels:
    label 0 is Background
    label 1 is pancrease
    label 2 is tumor
    """

    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"]),
        RandCropByPosNegLabeld(
            keys=["image", "label"],
            label_key="label",
            spatial_size=(128, 128, 64),
            pos=3,
            neg=0,
            num_samples=3,
            image_key="image",
            image_threshold=0,
        ),
        SqueezeDimd(keys=["label"],dim=0),
        ConvertToMultiChannelBasedOnPancreaClassesd(keys="label"),
        ScaleIntensityRanged(
            keys=["image"], a_min=-100, a_max=300,
            b_min=0.0, b_max=1.0, clip=True,
        ),
        EnsureTyped(keys=["image", "label"]),
    ]
)
val_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        RandCropByPosNegLabeld(
            keys=["image", "label"],
            label_key="label",
            spatial_size=(128, 128, 64),
            pos=3,
            neg=0,
            num_samples=3,
            image_key="image",
            image_threshold=0,
        ),
        SqueezeDimd(keys=["label"],dim=0),
        ConvertToMultiChannelBasedOnPancreaClassesd(keys="label"),
        ScaleIntensityRanged(
            keys=["image"], a_min=-100, a_max=300,
            b_min=0.0, b_max=1.0, clip=True,
        ),
        EnsureTyped(keys=["image", "label"]),
    ]
)
train_ds = monai.data.Dataset(data=train_files, transform=train_transforms)
val_ds = Dataset(data=val_files, transform=val_transforms)

max_epochs = 800
val_interval = 25

device = torch.device("cuda:0")

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

loss_function = DiceLoss(smooth_nr=0, smooth_dr=1e-5, squared_pred=True, to_onehot_y=False, sigmoid=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-2, weight_decay=1e-5)
#lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_epochs)
scheduler_poly_lr_decay = PolynomialLRDecay(optimizer, max_decay_steps=max_epochs, end_learning_rate=0.00001, power=2.0)


dice_metric = DiceMetric(include_background=False, reduction="mean")
dice_metric_batch = DiceMetric(include_background=False, reduction="mean_batch")

post_trans = Compose(
    [EnsureType(), Activations(sigmoid=True), AsDiscrete(threshold_values=True)]
)

# define inference method
def inference(input):
    def _compute(input):
        return sliding_window_inference(
            inputs=input,
            roi_size=(128, 128, 64),
            sw_batch_size=1,
            predictor=model,
        )
    return _compute(input)

train_loader = DataLoader(train_ds, batch_size=4, shuffle=True, num_workers=4)
val_loader = DataLoader(val_ds, batch_size=4, num_workers=4)

best_metric = -1
best_metric_epoch = -1
best_metrics_epochs_and_time = [[], []]
epoch_loss_values = []
metric_values = []
metric_values_pan = []
metric_values_tum = []
learning_rate_logs = []


total_start = time.time()
for epoch in range(max_epochs):
    epoch_start = time.time()
    print("-" * 30)
    print(f"epoch {epoch + 1}/{max_epochs}")
    model.train()
    epoch_loss = 0
    step = 0
    for batch_data in train_loader:
        step_start = time.time()
        step += 1

        inputs, labels = (
            batch_data["image"].to(device),
            batch_data["label"].to(device),
        )
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()

    learning_rate_logs.append(optimizer.param_groups[0]['lr'])
    scheduler_poly_lr_decay.step()
    
#    lr_scheduler.step()
    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),
                )
                val_outputs = inference(val_inputs)
                val_outputs = [post_trans(i) for i in decollate_batch(val_outputs)]
                dice_metric(y_pred=val_outputs, y=val_labels)
                dice_metric_batch(y_pred=val_outputs, y=val_labels)

            metric = dice_metric.aggregate().item()
            metric_values.append(metric)
            metric_batch = dice_metric_batch.aggregate()
            metric_pan = metric_batch[0].item()
            metric_values_pan.append(metric_pan)
            metric_tum = metric_batch[1].item()
            metric_values_tum.append(metric_tum)
            dice_metric.reset()
            dice_metric_batch.reset()

            if metric > best_metric:
                best_metric = metric
                best_metric_epoch = epoch + 1
                best_metrics_epochs_and_time[0].append(best_metric)
                best_metrics_epochs_and_time[1].append(time.time() - total_start)
                torch.save(
                    model.state_dict(),
                    os.path.join(root_dir, "best_metric_model.pth"),
                )
                print("saved new best metric model")
            print(
                f"current epoch: {epoch + 1} current mean dice: {metric:.4f}"
                f" Pan: {metric_pan:.4f} tumor: {metric_tum:.4f}"
                f"\nbest mean dice: {best_metric:.4f}"
                f" at epoch: {best_metric_epoch}"
            )
    print(f"time consuming of epoch {epoch + 1} is: {(time.time() - epoch_start):.4f}")

    if (epoch + 1) % 100 == 0:
        print('epoch_loss_values',best_metrics_epochs_and_time)
        print('metric_values',metric_values)
        print('metric_values_pan',metric_values_pan)
        print('metric_values_tum',metric_values_tum)
        print('learning_rate_logs',learning_rate_logs)

total_time = time.time() - total_start

print('total_time',total_time)

I trained model for 800 epochs. Pancreas score was fine but for tumor score was very low. can you please give me some advice is there something I am missing ?

epoch 800/800
epoch 800 average loss: 0.0690
current epoch: 800 current mean dice: 0.3960 Pan: 0.7164 tumor: 0.0504
best mean dice: 0.4190 at epoch: 625
time consuming of epoch 800 is: 80.8982

[ramchandracheke]

Also, I have one more question.
current mean dice: 0.3960 Pan: 0.7164 tumor: 0.0504
current mean dice = (0.7164+0.0504)/2= 0.3834 but its slightly less than 0.3960?

[Nic-Ma]

Hi @dongyang0122 ,

Could you please help share some experience on this training task?

Thanks in advance.

[dongyang0122]

Hi @ramchandracheke , based on the training scripts, it seems that your formulation is 3-class output. Please consider to use softmax instead of sigmoid in your loss function and post_trans.
Here is an example of 2-class segmentation.
https://github.com/Project-MONAI/tutorials/blob/master/3d_segmentation/spleen_segmentation_3d.ipynb

[ramchandracheke]

I followed the 3D segmentation Brats Tutorial. Thank you very much! I will try that and update you soon!
It would be great if you can also provide some insight about mean dice score!
Thanks a million!

[ramchandracheke]

Quick update. It worked for me!

[vincentochs]

Any chance of getting a look into your implementation @ramchandracheke ? :)

[debi2023-group13]

@ramchandracheke Where can I find data_dir ='/work/Monai/Cropped2/'?

[mmkaiser]

@debi2023-group13 - I'm looking for this information too. Did you happen to figure this out?

[KumoLiu]

Hi both, I think the dataset comes from http://medicaldecathlon.com/. The data_dir is just a path set by @ramchandracheke.
Thanks!

[mmkaiser]

Thank you @KumoLiu. I am getting this error in the first epoch:

ValueError: The size of the proposed random crop ROI is larger than the image size, got ROI size (96, 96, 96) and label image size (219, 209, 71) respectively.

Is there a cropped dataset somewhere? Or some pre-processing of this dataset that was put in this Cropped directory? I am running the Task07_Pancreas data "as is" currently when getting this error.

[KumoLiu]

Hi @mmkaiser, I noticed that the ROI size used on the code above is (128, 128, 64).
However, you can try these approaches to address the error itself.

• Try using Resize to unify the size

Or

• Set allow_smaller to true in RandCropByPosNegLabel and then add a ResizeWithPadOrCrop after it to unify the sizes

Hope it helps, Thanks.