Segmentation model only learning background and organ label not the cancer label #8133
Replies: 2 comments
-
Hi @faisalahm3d, We have a number of tutorials available in our tutorial repository that can help you get started. I highly recommend beginning with one of the segmentation tutorials. Once you've successfully run it, you can try replacing the dataset with your own. Thanks. |
Beta Was this translation helpful? Give feedback.
-
Hi @KumoLiu , Thank you very much for your sincere response. I followed this tutorial before writing this segmentation pipeline. Unfortunately, it is not working for the 3D CT image segmentation task with three labels. I would appreciate further suggestions. |
Beta Was this translation helpful? Give feedback.
-
I am new to the medical imaging domain and MONAI. As part of my learning, I attempted to create a segmentation pipeline for CT images with three labels: background (0), liver (1), and cancer (2). Unfortunately, the model is only learning for the background and liver, while the validation loss for cancer always returns 'NaN'. Could you please help me identify where I may have made a mistake in the segmentation pipeline from the code provided? Thank you very much.
`data_dir = '/home/Task03_Liver/imagesTr'
labels_dir = '/home/Task03_Liver/labelsTr'
images = sorted(
glob.glob(os.path.join(data_dir, "liver*.nii.gz")))
labels = sorted(
glob.glob(os.path.join(data_dir, "liver*.nii.gz")))
print('No. of images:', len(images), ' labels:', len(labels))
img = nib.load(images[0]).get_fdata()
lbl = nib.load(labels[0]).get_fdata()
print('\nimg shape:', img.shape, ' lbl shape:', lbl.shape)
print('img intensity min.:', np.min(img), ' max.:', np.max(img), ' unique labels:', np.unique(lbl))
data_dicts = [
{"image": image_name, "label": label_name}
for image_name, label_name in zip(images, labels)
]
train_files, val_files, test_files = data_dicts[:93], data_dicts[93:119], data_dicts[119:131]
print('train files:', len(train_files), ' val files:', len(val_files), ' test files:', len(test_files))
train_transforms = Compose(
[
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys=["image", "label"]),
ScaleIntensityRanged(
keys=["image"], a_min=-1000, a_max=1000,
b_min=0.0, b_max=1.0, clip=True,
),
val_transforms = Compose(
[
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys=["image", "label"]),
ScaleIntensityRanged(
keys=["image"], a_min=-1000, a_max=1000,
b_min=0.0, b_max=1.0, clip=True,
),
ScaleIntensityRanged(
keys=["label"], a_min=0, a_max=255,
b_min=0.0, b_max=1.0, clip=True,
),
CropForegroundd(keys=["image", "label"], source_key="image"),
Orientationd(keys=["image", "label"], axcodes="PLS"),
Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 2.0), mode=("bilinear", "nearest")),
])
train_ds = CacheDataset(data=train_files, transform=train_transforms, cache_rate=1.0)
train_loader = DataLoader(train_ds, batch_size=2, shuffle=True)
val_ds = CacheDataset(data=val_files, transform=val_transforms, cache_rate=1.0)
val_loader = DataLoader(val_ds, batch_size=1)
device = torch.device("cuda:0")
model = UNet(
spatial_dims=3,
in_channels=1,
out_channels=3, # 2 if using Softmax activation, 1 if using Sigmoid
channels=(32, 64, 128, 256, 320, 320),
strides=(1, 2, 2, 2, 2, 2),
num_res_units=2,
norm=Norm.BATCH,
).to(device)
loss_function = DiceLoss(to_onehot_y=True, softmax=True)
loss_function = DiceCELoss(to_onehot_y=True, softmax=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-4)
dice_metric = DiceMetric(include_background=False, reduction="mean",)
dice_metric2 = DiceMetric(include_background=True, reduction="none")
from monai.inferers import sliding_window_inference
Define hyperparameters
max_epochs = 500
val_interval = 2
best_metric = -1
classwise_best_metrics = None
best_metric_epoch = -1
epoch_loss_values = []
metric_values = []
Define post-processing transforms
post_pred = Compose([AsDiscrete(argmax=True, to_onehot=3)])
post_label = Compose([AsDiscrete(to_onehot=3)])
Training loop
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()
outputs = model(inputs)
loss = loss_function(outputs, labels)
print(
f"train completed \nbest_metric: {best_metric:.4f} \n classwise best mean dice: {classwise_best_metrics}"
f"\nat epoch: {best_metric_epoch}")
`
Beta Was this translation helpful? Give feedback.
All reactions