def _define_prediction_transforms(self):
"""Define and initialize all prediction data transforms.
* prediction set images transform
* prediction set images post-transform
@return True if data transforms could be instantiated, False otherwise.
"""
# Define transforms for prediction
self._prediction_image_transforms = Compose(
[
LoadImage(image_only=True),
ScaleIntensity(),
AddChannel(),
ToTensor(),
]
)
self._prediction_post_transforms = Compose(
[
Activations(softmax=True),
AsDiscrete(threshold_values=True),
]
)
def __init__(self, tranforms):
self.tranform_list = []
for tranform in tranforms:
if 'LoadImaged' == tranform:
self.tranform_list.append(LoadImaged(keys=["image", "label"]))
elif 'AsChannelFirstd' == tranform:
self.tranform_list.append(AsChannelFirstd(keys="image"))
elif 'ConvertToMultiChannelBasedOnBratsClassesd' == tranform:
self.tranform_list.append(ConvertToMultiChannelBasedOnBratsClassesd(keys="label"))
elif 'Spacingd' == tranform:
self.tranform_list.append(Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 2.0), mode=("bilinear", "nearest")))
elif 'Orientationd' == tranform:
self.tranform_list.append(Orientationd(keys=["image", "label"], axcodes="RAS"))
elif 'CenterSpatialCropd' == tranform:
self.tranform_list.append(CenterSpatialCropd(keys=["image", "label"], roi_size=[128, 128, 64]))
elif 'NormalizeIntensityd' == tranform:
self.tranform_list.append(NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True))
elif 'ToTensord' == tranform:
self.tranform_list.append(ToTensord(keys=["image", "label"]))
elif 'Activations' == tranform:
self.tranform_list.append(Activations(sigmoid=True))
elif 'AsDiscrete' == tranform:
self.tranform_list.append(AsDiscrete(threshold_values=True))
else:
raise ValueError(
f"Unsupported tranform: {tranform}. Please add it to support it."
)
super().__init__(self.tranform_list)
def test_compute(self):
auc_metric = ROCAUC()
act = Activations(softmax=True)
to_onehot = AsDiscrete(to_onehot=2)
device = f"cuda:{dist.get_rank()}" if torch.cuda.is_available(
) else "cpu"
if dist.get_rank() == 0:
y_pred = [
torch.tensor([0.1, 0.9], device=device),
torch.tensor([0.3, 1.4], device=device)
]
y = [
torch.tensor([0], device=device),
torch.tensor([1], device=device)
]
if dist.get_rank() == 1:
y_pred = [
torch.tensor([0.2, 0.1], device=device),
torch.tensor([0.1, 0.5], device=device),
torch.tensor([0.3, 0.4], device=device),
]
y = [
torch.tensor([0], device=device),
torch.tensor([1], device=device),
torch.tensor([1], device=device)
]
y_pred = [act(p) for p in y_pred]
y = [to_onehot(y_) for y_ in y]
auc_metric.update([y_pred, y])
result = auc_metric.compute()
np.testing.assert_allclose(0.66667, result, rtol=1e-4)
def run_inference_test(root_dir, device="cuda:0"):
images = sorted(glob(os.path.join(root_dir, "im*.nii.gz")))
segs = sorted(glob(os.path.join(root_dir, "seg*.nii.gz")))
val_files = [{"img": img, "seg": seg} for img, seg in zip(images, segs)]
# define transforms for image and segmentation
val_transforms = Compose([
LoadNiftid(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
# resampling with align_corners=True or dtype=float64 will generate
# slight different results between PyTorch 1.5 an 1.6
Spacingd(keys=["img", "seg"],
pixdim=[1.2, 0.8, 0.7],
mode=["bilinear", "nearest"],
dtype=np.float32),
ScaleIntensityd(keys="img"),
ToTensord(keys=["img", "seg"]),
])
val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
# sliding window inference need to input 1 image in every iteration
val_loader = monai.data.DataLoader(val_ds, batch_size=1, num_workers=4)
val_post_tran = Compose(
[Activations(sigmoid=True),
AsDiscrete(threshold_values=True)])
dice_metric = DiceMetric(include_background=True, reduction="mean")
model = UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
model_filename = os.path.join(root_dir, "best_metric_model.pth")
model.load_state_dict(torch.load(model_filename))
model.eval()
with torch.no_grad():
metric_sum = 0.0
metric_count = 0
# resampling with align_corners=True or dtype=float64 will generate
# slight different results between PyTorch 1.5 an 1.6
saver = NiftiSaver(output_dir=os.path.join(root_dir, "output"),
dtype=np.float32)
for val_data in val_loader:
val_images, val_labels = val_data["img"].to(
device), val_data["seg"].to(device)
# define sliding window size and batch size for windows inference
sw_batch_size, roi_size = 4, (96, 96, 96)
val_outputs = val_post_tran(
sliding_window_inference(val_images, roi_size, sw_batch_size,
model))
value, not_nans = dice_metric(y_pred=val_outputs, y=val_labels)
metric_count += not_nans.item()
metric_sum += value.item() * not_nans.item()
saver.save_batch(val_outputs, val_data["img_meta_dict"])
metric = metric_sum / metric_count
return metric
def test_class_value(self, y_pred, y, softmax, to_onehot, average,
expected_value):
y_pred = Activations(softmax=softmax)(y_pred)
y = AsDiscrete(to_onehot=to_onehot, n_classes=2)(y)
metric = ROCAUCMetric(average=average)
metric(y_pred=y_pred, y=y)
result = metric.aggregate()
np.testing.assert_allclose(expected_value, result, rtol=1e-5)
def run_inference_test(root_dir, device="cuda:0"):
images = sorted(glob(os.path.join(root_dir, "im*.nii.gz")))
segs = sorted(glob(os.path.join(root_dir, "seg*.nii.gz")))
val_files = [{"img": img, "seg": seg} for img, seg in zip(images, segs)]
# define transforms for image and segmentation
val_transforms = Compose(
[
LoadImaged(keys=["img", "seg"]),
EnsureChannelFirstd(keys=["img", "seg"]),
# resampling with align_corners=True or dtype=float64 will generate
# slight different results between PyTorch 1.5 an 1.6
Spacingd(keys=["img", "seg"], pixdim=[1.2, 0.8, 0.7], mode=["bilinear", "nearest"], dtype=np.float32),
ScaleIntensityd(keys="img"),
ToTensord(keys=["img", "seg"]),
]
)
val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
# sliding window inference need to input 1 image in every iteration
val_loader = monai.data.DataLoader(val_ds, batch_size=1, num_workers=4)
val_post_tran = Compose([ToTensor(), Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
model = UNet(
spatial_dims=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
model_filename = os.path.join(root_dir, "best_metric_model.pth")
model.load_state_dict(torch.load(model_filename))
with eval_mode(model):
# resampling with align_corners=True or dtype=float64 will generate
# slight different results between PyTorch 1.5 an 1.6
saver = SaveImage(
output_dir=os.path.join(root_dir, "output"),
dtype=np.float32,
output_ext=".nii.gz",
output_postfix="seg",
mode="bilinear",
)
for val_data in val_loader:
val_images, val_labels = val_data["img"].to(device), val_data["seg"].to(device)
# define sliding window size and batch size for windows inference
sw_batch_size, roi_size = 4, (96, 96, 96)
val_outputs = sliding_window_inference(val_images, roi_size, sw_batch_size, model)
# decollate prediction into a list
val_outputs = [val_post_tran(i) for i in decollate_batch(val_outputs)]
val_meta = decollate_batch(val_data[PostFix.meta("img")])
# compute metrics
dice_metric(y_pred=val_outputs, y=val_labels)
for img, meta in zip(val_outputs, val_meta): # save a decollated batch of files
saver(img, meta)
return dice_metric.aggregate().item()
def test_value(self, y_pred, y, softmax, to_onehot, average,
expected_value):
y_pred_trans = Compose([ToTensor(), Activations(softmax=softmax)])
y_trans = Compose([ToTensor(), AsDiscrete(to_onehot=to_onehot)])
y_pred = torch.stack(
[y_pred_trans(i) for i in decollate_batch(y_pred)], dim=0)
y = torch.stack([y_trans(i) for i in decollate_batch(y)], dim=0)
result = compute_roc_auc(y_pred=y_pred, y=y, average=average)
np.testing.assert_allclose(expected_value, result, rtol=1e-5)
def main(tempdir):
config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
print(f"generating synthetic data to {tempdir} (this may take a while)")
for i in range(5):
im, seg = create_test_image_3d(128, 128, 128, num_seg_classes=1)
n = nib.Nifti1Image(im, np.eye(4))
nib.save(n, os.path.join(tempdir, f"im{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(tempdir, "im*.nii.gz")))
segs = sorted(glob(os.path.join(tempdir, "seg*.nii.gz")))
# define transforms for image and segmentation
imtrans = Compose([ScaleIntensity(), AddChannel(), ToTensor()])
segtrans = Compose([AddChannel(), ToTensor()])
val_ds = ImageDataset(images, segs, transform=imtrans, seg_transform=segtrans, image_only=False)
# sliding window inference for one image at every iteration
val_loader = DataLoader(val_ds, batch_size=1, num_workers=1, pin_memory=torch.cuda.is_available())
dice_metric = DiceMetric(include_background=True, reduction="mean")
post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold_values=True)])
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
model.load_state_dict(torch.load("best_metric_model_segmentation3d_array.pth"))
model.eval()
with torch.no_grad():
metric_sum = 0.0
metric_count = 0
saver = NiftiSaver(output_dir="./output")
for val_data in val_loader:
val_images, val_labels = val_data[0].to(device), val_data[1].to(device)
# define sliding window size and batch size for windows inference
roi_size = (96, 96, 96)
sw_batch_size = 4
val_outputs = sliding_window_inference(val_images, roi_size, sw_batch_size, model)
val_outputs = post_trans(val_outputs)
value, _ = dice_metric(y_pred=val_outputs, y=val_labels)
metric_count += len(value)
metric_sum += value.item() * len(value)
saver.save_batch(val_outputs, val_data[2])
metric = metric_sum / metric_count
print("evaluation metric:", metric)
def test_class_value(self, y_pred, y, softmax, to_onehot, average,
expected_value):
y_pred_trans = Compose([ToTensor(), Activations(softmax=softmax)])
y_trans = Compose([ToTensor(), AsDiscrete(to_onehot=to_onehot)])
y_pred = [y_pred_trans(i) for i in decollate_batch(y_pred)]
y = [y_trans(i) for i in decollate_batch(y)]
metric = ROCAUCMetric(average=average)
metric(y_pred=y_pred, y=y)
result = metric.aggregate()
metric.reset()
np.testing.assert_allclose(expected_value, result, rtol=1e-5)
def test_value_shape(self, input_param, img, out, expected_shape):
result = Activations(**input_param)(img)
def _compare(ret, out, shape):
torch.testing.assert_allclose(ret, out)
self.assertTupleEqual(ret.shape, shape)
if isinstance(result, (list, tuple)):
for r, e in zip(result, out):
_compare(r, e, expected_shape)
else:
_compare(result, out, expected_shape)
def main(tempdir):
config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
print(f"generating synthetic data to {tempdir} (this may take a while)")
for i in range(5):
im, seg = create_test_image_2d(128, 128, num_seg_classes=1)
Image.fromarray((im * 255).astype("uint8")).save(os.path.join(tempdir, f"img{i:d}.png"))
Image.fromarray((seg * 255).astype("uint8")).save(os.path.join(tempdir, f"seg{i:d}.png"))
images = sorted(glob(os.path.join(tempdir, "img*.png")))
segs = sorted(glob(os.path.join(tempdir, "seg*.png")))
# define transforms for image and segmentation
imtrans = Compose([LoadImage(image_only=True), AddChannel(), ScaleIntensity(), EnsureType()])
segtrans = Compose([LoadImage(image_only=True), AddChannel(), ScaleIntensity(), EnsureType()])
val_ds = ArrayDataset(images, imtrans, segs, segtrans)
# sliding window inference for one image at every iteration
val_loader = DataLoader(val_ds, batch_size=1, num_workers=1, pin_memory=torch.cuda.is_available())
dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
post_trans = Compose([EnsureType(), Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
saver = SaveImage(output_dir="./output", output_ext=".png", output_postfix="seg")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UNet(
spatial_dims=2,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
model.load_state_dict(torch.load("best_metric_model_segmentation2d_array.pth"))
model.eval()
with torch.no_grad():
for val_data in val_loader:
val_images, val_labels = val_data[0].to(device), val_data[1].to(device)
# define sliding window size and batch size for windows inference
roi_size = (96, 96)
sw_batch_size = 4
val_outputs = sliding_window_inference(val_images, roi_size, sw_batch_size, model)
val_outputs = [post_trans(i) for i in decollate_batch(val_outputs)]
val_labels = decollate_batch(val_labels)
# compute metric for current iteration
dice_metric(y_pred=val_outputs, y=val_labels)
for val_output in val_outputs:
saver(val_output)
# aggregate the final mean dice result
print("evaluation metric:", dice_metric.aggregate().item())
# reset the status
dice_metric.reset()
def evaluate(model, data_loader, device):
metric = torch.zeros(8, dtype=torch.float, device=device)
model.eval()
with torch.no_grad():
dice_metric = DiceMetric(include_background=True, reduction="mean")
post_trans = Compose(
[Activations(sigmoid=True),
AsDiscrete(threshold_values=True)])
for val_data in data_loader:
val_inputs, val_labels = (
val_data["image"].to(device, non_blocking=True),
val_data["label"].to(device, non_blocking=True),
)
val_outputs = model(val_inputs)
val_outputs = post_trans(val_outputs)
# compute overall mean dice
value, not_nans = dice_metric(y_pred=val_outputs, y=val_labels)
value = value.squeeze()
metric[0] += value * not_nans
metric[1] += not_nans
# compute mean dice for TC
value_tc, not_nans = dice_metric(y_pred=val_outputs[:, 0:1],
y=val_labels[:, 0:1])
value_tc = value_tc.squeeze()
metric[2] += value_tc * not_nans
metric[3] += not_nans
# compute mean dice for WT
value_wt, not_nans = dice_metric(y_pred=val_outputs[:, 1:2],
y=val_labels[:, 1:2])
value_wt = value_wt.squeeze()
metric[4] += value_wt * not_nans
metric[5] += not_nans
# compute mean dice for ET
value_et, not_nans = dice_metric(y_pred=val_outputs[:, 2:3],
y=val_labels[:, 2:3])
value_et = value_et.squeeze()
metric[6] += value_et * not_nans
metric[7] += not_nans
# synchronizes all processes and reduce results
dist.barrier()
dist.all_reduce(metric, op=torch.distributed.ReduceOp.SUM)
metric = metric.tolist()
return metric[0] / metric[1], metric[2] / metric[3], metric[4] / metric[
5], metric[6] / metric[7]
def test_compute(self):
auc_metric = ROCAUC()
act = Activations(softmax=True)
to_onehot = AsDiscrete(to_onehot=True, n_classes=2)
y_pred = torch.Tensor([[0.1, 0.9], [0.3, 1.4]])
y = torch.Tensor([[0], [1]])
y_pred = act(y_pred)
y = to_onehot(y)
auc_metric.update([y_pred, y])
y_pred = torch.Tensor([[0.2, 0.1], [0.1, 0.5]])
y = torch.Tensor([[0], [1]])
y_pred = act(y_pred)
y = to_onehot(y)
auc_metric.update([y_pred, y])
auc = auc_metric.compute()
np.testing.assert_allclose(0.75, auc)
def test_compute(self):
auc_metric = ROCAUC()
act = Activations(softmax=True)
to_onehot = AsDiscrete(to_onehot=True, num_classes=2)
y_pred = [torch.Tensor([0.1, 0.9]), torch.Tensor([0.3, 1.4])]
y = [torch.Tensor([0]), torch.Tensor([1])]
y_pred = [act(p) for p in y_pred]
y = [to_onehot(y_) for y_ in y]
auc_metric.update([y_pred, y])
y_pred = [torch.Tensor([0.2, 0.1]), torch.Tensor([0.1, 0.5])]
y = [torch.Tensor([0]), torch.Tensor([1])]
y_pred = [act(p) for p in y_pred]
y = [to_onehot(y_) for y_ in y]
auc_metric.update([y_pred, y])
auc = auc_metric.compute()
np.testing.assert_allclose(0.75, auc)
def test_compute(self):
auc_metric = ROCAUC()
act = Activations(softmax=True)
to_onehot = AsDiscrete(to_onehot=True, n_classes=2)
device = f"cuda:{dist.get_rank()}" if torch.cuda.is_available(
) else "cpu"
if dist.get_rank() == 0:
y_pred = torch.tensor([[0.1, 0.9], [0.3, 1.4]], device=device)
y = torch.tensor([[0], [1]], device=device)
if dist.get_rank() == 1:
y_pred = torch.tensor([[0.2, 0.1], [0.1, 0.5], [0.3, 0.4]],
device=device)
y = torch.tensor([[0], [1], [1]], device=device)
y_pred = act(y_pred)
y = to_onehot(y)
auc_metric.update([y_pred, y])
result = auc_metric.compute()
np.testing.assert_allclose(0.66667, result, rtol=1e-4)
def run_training_test(root_dir, train_x, train_y, val_x, val_y, device="cuda:0", num_workers=10):
monai.config.print_config()
# define transforms for image and classification
train_transforms = Compose(
[
LoadImage(image_only=True),
AddChannel(),
Transpose(indices=[0, 2, 1]),
ScaleIntensity(),
RandRotate(range_x=np.pi / 12, prob=0.5, keep_size=True, dtype=np.float64),
RandFlip(spatial_axis=0, prob=0.5),
RandZoom(min_zoom=0.9, max_zoom=1.1, prob=0.5),
ToTensor(),
]
)
train_transforms.set_random_state(1234)
val_transforms = Compose(
[LoadImage(image_only=True), AddChannel(), Transpose(indices=[0, 2, 1]), ScaleIntensity(), ToTensor()]
)
y_pred_trans = Compose([ToTensor(), Activations(softmax=True)])
y_trans = Compose([ToTensor(), AsDiscrete(to_onehot=len(np.unique(train_y)))])
auc_metric = ROCAUCMetric()
# create train, val data loaders
train_ds = MedNISTDataset(train_x, train_y, train_transforms)
train_loader = DataLoader(train_ds, batch_size=300, shuffle=True, num_workers=num_workers)
val_ds = MedNISTDataset(val_x, val_y, val_transforms)
val_loader = DataLoader(val_ds, batch_size=300, num_workers=num_workers)
model = DenseNet121(spatial_dims=2, in_channels=1, out_channels=len(np.unique(train_y))).to(device)
loss_function = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), 1e-5)
epoch_num = 4
val_interval = 1
# start training validation
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = []
metric_values = []
model_filename = os.path.join(root_dir, "best_metric_model.pth")
for epoch in range(epoch_num):
print("-" * 10)
print(f"Epoch {epoch + 1}/{epoch_num}")
model.train()
epoch_loss = 0
step = 0
for batch_data in train_loader:
step += 1
inputs, labels = batch_data[0].to(device), batch_data[1].to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = loss_function(outputs, labels)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
epoch_loss /= step
epoch_loss_values.append(epoch_loss)
print(f"epoch {epoch + 1} average loss:{epoch_loss:0.4f}")
if (epoch + 1) % val_interval == 0:
with eval_mode(model):
y_pred = torch.tensor([], dtype=torch.float32, device=device)
y = torch.tensor([], dtype=torch.long, device=device)
for val_data in val_loader:
val_images, val_labels = val_data[0].to(device), val_data[1].to(device)
y_pred = torch.cat([y_pred, model(val_images)], dim=0)
y = torch.cat([y, val_labels], dim=0)
# compute accuracy
acc_value = torch.eq(y_pred.argmax(dim=1), y)
acc_metric = acc_value.sum().item() / len(acc_value)
# decollate prediction and label and execute post processing
y_pred = [y_pred_trans(i) for i in decollate_batch(y_pred)]
y = [y_trans(i) for i in decollate_batch(y)]
# compute AUC
auc_metric(y_pred, y)
auc_value = auc_metric.aggregate()
auc_metric.reset()
metric_values.append(auc_value)
if auc_value > best_metric:
best_metric = auc_value
best_metric_epoch = epoch + 1
torch.save(model.state_dict(), model_filename)
print("saved new best metric model")
print(
f"current epoch {epoch +1} current AUC: {auc_value:0.4f} "
f"current accuracy: {acc_metric:0.4f} best AUC: {best_metric:0.4f} at epoch {best_metric_epoch}"
)
print(f"train completed, best_metric: {best_metric:0.4f} at epoch: {best_metric_epoch}")
return epoch_loss_values, best_metric, best_metric_epoch
def test_value(self, y_pred, y, softmax, to_onehot, average,
expected_value):
y_pred = Activations(softmax=softmax)(y_pred)
y = AsDiscrete(to_onehot=to_onehot, n_classes=2)(y)
result = compute_roc_auc(y_pred=y_pred, y=y, average=average)
np.testing.assert_allclose(expected_value, result, rtol=1e-5)
def main(tempdir):
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
# create a temporary directory and 40 random image, mask pairs
print(f"generating synthetic data to {tempdir} (this may take a while)")
for i in range(40):
im, seg = create_test_image_2d(128, 128, num_seg_classes=1)
Image.fromarray(im.astype("uint8")).save(
os.path.join(tempdir, f"img{i:d}.png"))
Image.fromarray(seg.astype("uint8")).save(
os.path.join(tempdir, f"seg{i:d}.png"))
images = sorted(glob(os.path.join(tempdir, "img*.png")))
segs = sorted(glob(os.path.join(tempdir, "seg*.png")))
train_files = [{
"img": img,
"seg": seg
} for img, seg in zip(images[:20], segs[:20])]
val_files = [{
"img": img,
"seg": seg
} for img, seg in zip(images[-20:], segs[-20:])]
# define transforms for image and segmentation
train_imtrans = Compose([
LoadImage(image_only=True),
ScaleIntensity(),
AddChannel(),
RandSpatialCrop((96, 96), random_size=False),
RandRotate90(prob=0.5, spatial_axes=(0, 1)),
ToTensor(),
])
train_segtrans = Compose([
LoadImage(image_only=True),
AddChannel(),
RandSpatialCrop((96, 96), random_size=False),
RandRotate90(prob=0.5, spatial_axes=(0, 1)),
ToTensor(),
])
val_imtrans = Compose([
LoadImage(image_only=True),
ScaleIntensity(),
AddChannel(),
ToTensor()
])
val_segtrans = Compose(
[LoadImage(image_only=True),
AddChannel(), ToTensor()])
# define array dataset, data loader
check_ds = ArrayDataset(images, train_imtrans, segs, train_segtrans)
check_loader = DataLoader(check_ds,
batch_size=10,
num_workers=2,
pin_memory=torch.cuda.is_available())
im, seg = monai.utils.misc.first(check_loader)
print(im.shape, seg.shape)
# create a training data loader
train_ds = ArrayDataset(images[:20], train_imtrans, segs[:20],
train_segtrans)
train_loader = DataLoader(train_ds,
batch_size=4,
shuffle=True,
num_workers=8,
pin_memory=torch.cuda.is_available())
# create a validation data loader
val_ds = ArrayDataset(images[-20:], val_imtrans, segs[-20:], val_segtrans)
val_loader = DataLoader(val_ds,
batch_size=1,
num_workers=4,
pin_memory=torch.cuda.is_available())
dice_metric = DiceMetric(include_background=True, reduction="mean")
post_trans = Compose(
[Activations(sigmoid=True),
AsDiscrete(threshold_values=True)])
# create UNet, DiceLoss and Adam optimizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = monai.networks.nets.UNet(
dimensions=2,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
loss_function = monai.losses.DiceLoss(sigmoid=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-3)
# start a typical PyTorch training
val_interval = 2
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = list()
metric_values = list()
writer = SummaryWriter()
for epoch in range(10):
print("-" * 10)
print(f"epoch {epoch + 1}/{10}")
model.train()
epoch_loss = 0
step = 0
for batch_data in train_loader:
step += 1
inputs, labels = batch_data[0].to(device), batch_data[1].to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = loss_function(outputs, labels)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
epoch_len = len(train_ds) // train_loader.batch_size
print(f"{step}/{epoch_len}, train_loss: {loss.item():.4f}")
writer.add_scalar("train_loss", loss.item(),
epoch_len * epoch + 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():
metric_sum = 0.0
metric_count = 0
val_images = None
val_labels = None
val_outputs = None
for val_data in val_loader:
val_images, val_labels = val_data[0].to(
device), val_data[1].to(device)
roi_size = (96, 96)
sw_batch_size = 4
val_outputs = sliding_window_inference(
val_images, roi_size, sw_batch_size, model)
val_outputs = post_trans(val_outputs)
value, _ = dice_metric(y_pred=val_outputs, y=val_labels)
metric_count += len(value)
metric_sum += value.item() * len(value)
metric = metric_sum / metric_count
metric_values.append(metric)
if metric > best_metric:
best_metric = metric
best_metric_epoch = epoch + 1
torch.save(model.state_dict(),
"best_metric_model_segmentation2d_array.pth")
print("saved new best metric model")
print(
"current epoch: {} current mean dice: {:.4f} best mean dice: {:.4f} at epoch {}"
.format(epoch + 1, metric, best_metric, best_metric_epoch))
writer.add_scalar("val_mean_dice", metric, epoch + 1)
# plot the last model output as GIF image in TensorBoard with the corresponding image and label
plot_2d_or_3d_image(val_images,
epoch + 1,
writer,
index=0,
tag="image")
plot_2d_or_3d_image(val_labels,
epoch + 1,
writer,
index=0,
tag="label")
plot_2d_or_3d_image(val_outputs,
epoch + 1,
writer,
index=0,
tag="output")
print(
f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}"
)
writer.close()
def main(tempdir):
config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
print(f"generating synthetic data to {tempdir} (this may take a while)")
for i in range(5):
im, seg = create_test_image_3d(128, 128, 128, num_seg_classes=1)
n = nib.Nifti1Image(im, np.eye(4))
nib.save(n, os.path.join(tempdir, f"im{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(tempdir, "im*.nii.gz")))
segs = sorted(glob(os.path.join(tempdir, "seg*.nii.gz")))
# define transforms for image and segmentation
imtrans = Compose([ScaleIntensity(), AddChannel(), ToTensor()])
segtrans = Compose([AddChannel(), ToTensor()])
ds = ImageDataset(images,
segs,
transform=imtrans,
seg_transform=segtrans,
image_only=False)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
# define sliding window size and batch size for windows inference
roi_size = (96, 96, 96)
sw_batch_size = 4
post_trans = Compose(
[Activations(sigmoid=True),
AsDiscrete(threshold_values=True)])
def _sliding_window_processor(engine, batch):
net.eval()
with torch.no_grad():
val_images, val_labels = batch[0].to(device), batch[1].to(device)
seg_probs = sliding_window_inference(val_images, roi_size,
sw_batch_size, net)
seg_probs = post_trans(seg_probs)
return seg_probs, val_labels
evaluator = Engine(_sliding_window_processor)
# add evaluation metric to the evaluator engine
MeanDice().attach(evaluator, "Mean_Dice")
# StatsHandler prints loss at every iteration and print metrics at every epoch,
# we don't need to print loss for evaluator, so just print metrics, user can also customize print functions
val_stats_handler = StatsHandler(
name="evaluator",
output_transform=lambda x:
None, # no need to print loss value, so disable per iteration output
)
val_stats_handler.attach(evaluator)
# for the array data format, assume the 3rd item of batch data is the meta_data
file_saver = SegmentationSaver(
output_dir="tempdir",
output_ext=".nii.gz",
output_postfix="seg",
name="evaluator",
batch_transform=lambda x: x[2],
output_transform=lambda output: output[0],
)
file_saver.attach(evaluator)
# the model was trained by "unet_training_array" example
ckpt_saver = CheckpointLoader(
load_path="./runs_array/net_checkpoint_100.pt", load_dict={"net": net})
ckpt_saver.attach(evaluator)
# sliding window inference for one image at every iteration
loader = DataLoader(ds,
batch_size=1,
num_workers=1,
pin_memory=torch.cuda.is_available())
state = evaluator.run(loader)
print(state)
def post_transforms(self):
return Compose([
Activations(sigmoid=True),
AsDiscrete(threshold_values=True),
])
def run_training_test(root_dir, device="cuda:0", cachedataset=0):
monai.config.print_config()
images = sorted(glob(os.path.join(root_dir, "img*.nii.gz")))
segs = sorted(glob(os.path.join(root_dir, "seg*.nii.gz")))
train_files = [{"img": img, "seg": seg} for img, seg in zip(images[:20], segs[:20])]
val_files = [{"img": img, "seg": seg} for img, seg in zip(images[-20:], segs[-20:])]
# define transforms for image and segmentation
train_transforms = Compose(
[
LoadImaged(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
# resampling with align_corners=True or dtype=float64 will generate
# slight different results between PyTorch 1.5 an 1.6
Spacingd(keys=["img", "seg"], pixdim=[1.2, 0.8, 0.7], mode=["bilinear", "nearest"], dtype=np.float32),
ScaleIntensityd(keys="img"),
RandCropByPosNegLabeld(
keys=["img", "seg"], label_key="seg", spatial_size=[96, 96, 96], pos=1, neg=1, num_samples=4
),
RandRotate90d(keys=["img", "seg"], prob=0.8, spatial_axes=[0, 2]),
ToTensord(keys=["img", "seg"]),
]
)
train_transforms.set_random_state(1234)
val_transforms = Compose(
[
LoadImaged(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
# resampling with align_corners=True or dtype=float64 will generate
# slight different results between PyTorch 1.5 an 1.6
Spacingd(keys=["img", "seg"], pixdim=[1.2, 0.8, 0.7], mode=["bilinear", "nearest"], dtype=np.float32),
ScaleIntensityd(keys="img"),
ToTensord(keys=["img", "seg"]),
]
)
# create a training data loader
if cachedataset == 2:
train_ds = monai.data.CacheDataset(data=train_files, transform=train_transforms, cache_rate=0.8)
elif cachedataset == 3:
train_ds = monai.data.LMDBDataset(data=train_files, transform=train_transforms)
else:
train_ds = monai.data.Dataset(data=train_files, transform=train_transforms)
# use batch_size=2 to load images and use RandCropByPosNegLabeld to generate 2 x 4 images for network training
train_loader = monai.data.DataLoader(train_ds, batch_size=2, shuffle=True, num_workers=4)
# create a validation data loader
val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
val_loader = monai.data.DataLoader(val_ds, batch_size=1, num_workers=4)
val_post_tran = Compose([Activations(sigmoid=True), AsDiscrete(threshold_values=True)])
dice_metric = DiceMetric(include_background=True, reduction="mean")
# create UNet, DiceLoss and Adam optimizer
model = monai.networks.nets.UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
loss_function = monai.losses.DiceLoss(sigmoid=True)
optimizer = torch.optim.Adam(model.parameters(), 5e-4)
# start a typical PyTorch training
val_interval = 2
best_metric, best_metric_epoch = -1, -1
epoch_loss_values = list()
metric_values = list()
writer = SummaryWriter(log_dir=os.path.join(root_dir, "runs"))
model_filename = os.path.join(root_dir, "best_metric_model.pth")
for epoch in range(6):
print("-" * 10)
print(f"Epoch {epoch + 1}/{6}")
model.train()
epoch_loss = 0
step = 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()
epoch_len = len(train_ds) // train_loader.batch_size
print(f"{step}/{epoch_len}, train_loss:{loss.item():0.4f}")
writer.add_scalar("train_loss", loss.item(), epoch_len * epoch + step)
epoch_loss /= step
epoch_loss_values.append(epoch_loss)
print(f"epoch {epoch +1} average loss:{epoch_loss:0.4f}")
if (epoch + 1) % val_interval == 0:
model.eval()
with torch.no_grad():
metric_sum = 0.0
metric_count = 0
val_images = None
val_labels = None
val_outputs = None
for val_data in val_loader:
val_images, val_labels = val_data["img"].to(device), val_data["seg"].to(device)
sw_batch_size, roi_size = 4, (96, 96, 96)
val_outputs = val_post_tran(sliding_window_inference(val_images, roi_size, sw_batch_size, model))
value, not_nans = dice_metric(y_pred=val_outputs, y=val_labels)
metric_count += not_nans.item()
metric_sum += value.item() * not_nans.item()
metric = metric_sum / metric_count
metric_values.append(metric)
if metric > best_metric:
best_metric = metric
best_metric_epoch = epoch + 1
torch.save(model.state_dict(), model_filename)
print("saved new best metric model")
print(
f"current epoch {epoch +1} current mean dice: {metric:0.4f} "
f"best mean dice: {best_metric:0.4f} at epoch {best_metric_epoch}"
)
writer.add_scalar("val_mean_dice", metric, epoch + 1)
# plot the last model output as GIF image in TensorBoard with the corresponding image and label
plot_2d_or_3d_image(val_images, epoch + 1, writer, index=0, tag="image")
plot_2d_or_3d_image(val_labels, epoch + 1, writer, index=0, tag="label")
plot_2d_or_3d_image(val_outputs, epoch + 1, writer, index=0, tag="output")
print(f"train completed, best_metric: {best_metric:0.4f} at epoch: {best_metric_epoch}")
writer.close()
return epoch_loss_values, best_metric, best_metric_epoch
def main_worker(args):
# disable logging for processes except 0 on every node
if args.local_rank != 0:
f = open(os.devnull, "w")
sys.stdout = sys.stderr = f
if not os.path.exists(args.dir):
raise FileNotFoundError(f"missing directory {args.dir}")
# initialize the distributed training process, every GPU runs in a process
dist.init_process_group(backend="nccl", init_method="env://")
device = torch.device(f"cuda:{args.local_rank}")
torch.cuda.set_device(device)
# use amp to accelerate training
scaler = torch.cuda.amp.GradScaler()
torch.backends.cudnn.benchmark = True
total_start = time.time()
train_transforms = Compose([
# load 4 Nifti images and stack them together
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys="image"),
ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
Orientationd(keys=["image", "label"], axcodes="RAS"),
Spacingd(
keys=["image", "label"],
pixdim=(1.0, 1.0, 1.0),
mode=("bilinear", "nearest"),
),
EnsureTyped(keys=["image", "label"]),
ToDeviced(keys=["image", "label"], device=device),
RandSpatialCropd(keys=["image", "label"],
roi_size=[224, 224, 144],
random_size=False),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
RandScaleIntensityd(keys="image", factors=0.1, prob=0.5),
RandShiftIntensityd(keys="image", offsets=0.1, prob=0.5),
])
# create a training data loader
train_ds = BratsCacheDataset(
root_dir=args.dir,
transform=train_transforms,
section="training",
num_workers=4,
cache_rate=args.cache_rate,
shuffle=True,
)
# ThreadDataLoader can be faster if no IO operations when caching all the data in memory
train_loader = ThreadDataLoader(train_ds,
num_workers=0,
batch_size=args.batch_size,
shuffle=True)
# validation transforms and dataset
val_transforms = Compose([
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys="image"),
ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
Orientationd(keys=["image", "label"], axcodes="RAS"),
Spacingd(
keys=["image", "label"],
pixdim=(1.0, 1.0, 1.0),
mode=("bilinear", "nearest"),
),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
EnsureTyped(keys=["image", "label"]),
ToDeviced(keys=["image", "label"], device=device),
])
val_ds = BratsCacheDataset(
root_dir=args.dir,
transform=val_transforms,
section="validation",
num_workers=4,
cache_rate=args.cache_rate,
shuffle=False,
)
# ThreadDataLoader can be faster if no IO operations when caching all the data in memory
val_loader = ThreadDataLoader(val_ds,
num_workers=0,
batch_size=args.batch_size,
shuffle=False)
# create network, loss function and optimizer
if args.network == "SegResNet":
model = SegResNet(
blocks_down=[1, 2, 2, 4],
blocks_up=[1, 1, 1],
init_filters=16,
in_channels=4,
out_channels=3,
dropout_prob=0.0,
).to(device)
else:
model = UNet(
spatial_dims=3,
in_channels=4,
out_channels=3,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
loss_function = DiceFocalLoss(
smooth_nr=1e-5,
smooth_dr=1e-5,
squared_pred=True,
to_onehot_y=False,
sigmoid=True,
batch=True,
)
optimizer = Novograd(model.parameters(), lr=args.lr)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=args.epochs)
# wrap the model with DistributedDataParallel module
model = DistributedDataParallel(model, device_ids=[device])
dice_metric = DiceMetric(include_background=True, reduction="mean")
dice_metric_batch = DiceMetric(include_background=True,
reduction="mean_batch")
post_trans = Compose(
[EnsureType(),
Activations(sigmoid=True),
AsDiscrete(threshold=0.5)])
# start a typical PyTorch training
best_metric = -1
best_metric_epoch = -1
print(f"time elapsed before training: {time.time() - total_start}")
train_start = time.time()
for epoch in range(args.epochs):
epoch_start = time.time()
print("-" * 10)
print(f"epoch {epoch + 1}/{args.epochs}")
epoch_loss = train(train_loader, model, loss_function, optimizer,
lr_scheduler, scaler)
print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")
if (epoch + 1) % args.val_interval == 0:
metric, metric_tc, metric_wt, metric_et = evaluate(
model, val_loader, dice_metric, dice_metric_batch, post_trans)
if metric > best_metric:
best_metric = metric
best_metric_epoch = epoch + 1
if dist.get_rank() == 0:
torch.save(model.state_dict(), "best_metric_model.pth")
print(
f"current epoch: {epoch + 1} current mean dice: {metric:.4f}"
f" tc: {metric_tc:.4f} wt: {metric_wt:.4f} et: {metric_et:.4f}"
f"\nbest mean dice: {best_metric:.4f} at epoch: {best_metric_epoch}"
)
print(
f"time consuming of epoch {epoch + 1} is: {(time.time() - epoch_start):.4f}"
)
print(
f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch},"
f" total train time: {(time.time() - train_start):.4f}")
dist.destroy_process_group()
def test_test_time_augmentation(self):
input_size = (20, 40) # test different input data shape to pad list collate
keys = ["image", "label"]
num_training_ims = 10
train_data = self.get_data(num_training_ims, input_size)
test_data = self.get_data(1, input_size)
device = "cuda" if torch.cuda.is_available() else "cpu"
transforms = Compose(
[
AddChanneld(keys),
RandAffined(
keys,
prob=1.0,
spatial_size=(30, 30),
rotate_range=(np.pi / 3, np.pi / 3),
translate_range=(3, 3),
scale_range=((0.8, 1), (0.8, 1)),
padding_mode="zeros",
mode=("bilinear", "nearest"),
as_tensor_output=False,
),
CropForegroundd(keys, source_key="image"),
DivisiblePadd(keys, 4),
]
)
train_ds = CacheDataset(train_data, transforms)
# output might be different size, so pad so that they match
train_loader = DataLoader(train_ds, batch_size=2, collate_fn=pad_list_data_collate)
model = UNet(2, 1, 1, channels=(6, 6), strides=(2, 2)).to(device)
loss_function = DiceLoss(sigmoid=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-3)
num_epochs = 10
for _ in trange(num_epochs):
epoch_loss = 0
for batch_data in train_loader:
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()
epoch_loss /= len(train_loader)
post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
tt_aug = TestTimeAugmentation(
transform=transforms,
batch_size=5,
num_workers=0,
inferrer_fn=model,
device=device,
to_tensor=True,
output_device="cpu",
post_func=post_trans,
)
mode, mean, std, vvc = tt_aug(test_data)
self.assertEqual(mode.shape, (1,) + input_size)
self.assertEqual(mean.shape, (1,) + input_size)
self.assertTrue(all(np.unique(mode) == (0, 1)))
self.assertGreaterEqual(mean.min(), 0.0)
self.assertLessEqual(mean.max(), 1.0)
self.assertEqual(std.shape, (1,) + input_size)
self.assertIsInstance(vvc, float)
def main():
opt = Options().parse()
# monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
if opt.gpu_ids != '-1':
num_gpus = len(opt.gpu_ids.split(','))
else:
num_gpus = 0
print('number of GPU:', num_gpus)
# Data loader creation
# train images
train_images = sorted(glob(os.path.join(opt.images_folder, 'train', 'image*.nii')))
train_segs = sorted(glob(os.path.join(opt.labels_folder, 'train', 'label*.nii')))
train_images_for_dice = sorted(glob(os.path.join(opt.images_folder, 'train', 'image*.nii')))
train_segs_for_dice = sorted(glob(os.path.join(opt.labels_folder, 'train', 'label*.nii')))
# validation images
val_images = sorted(glob(os.path.join(opt.images_folder, 'val', 'image*.nii')))
val_segs = sorted(glob(os.path.join(opt.labels_folder, 'val', 'label*.nii')))
# test images
test_images = sorted(glob(os.path.join(opt.images_folder, 'test', 'image*.nii')))
test_segs = sorted(glob(os.path.join(opt.labels_folder, 'test', 'label*.nii')))
# augment the data list for training
for i in range(int(opt.increase_factor_data)):
train_images.extend(train_images)
train_segs.extend(train_segs)
print('Number of training patches per epoch:', len(train_images))
print('Number of training images per epoch:', len(train_images_for_dice))
print('Number of validation images per epoch:', len(val_images))
print('Number of test images per epoch:', len(test_images))
# Creation of data directories for data_loader
train_dicts = [{'image': image_name, 'label': label_name}
for image_name, label_name in zip(train_images, train_segs)]
train_dice_dicts = [{'image': image_name, 'label': label_name}
for image_name, label_name in zip(train_images_for_dice, train_segs_for_dice)]
val_dicts = [{'image': image_name, 'label': label_name}
for image_name, label_name in zip(val_images, val_segs)]
test_dicts = [{'image': image_name, 'label': label_name}
for image_name, label_name in zip(test_images, test_segs)]
# Transforms list
# Need to concatenate multiple channels here if you want multichannel segmentation
# Check other examples on Monai webpage.
if opt.resolution is not None:
train_transforms = [
LoadImaged(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
NormalizeIntensityd(keys=['image']),
ScaleIntensityd(keys=['image']),
Spacingd(keys=['image', 'label'], pixdim=opt.resolution, mode=('bilinear', 'nearest')),
RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=1),
RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=0),
RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=2),
RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1,
rotate_range=(np.pi / 36, np.pi / 36, np.pi * 2), padding_mode="zeros"),
RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1,
rotate_range=(np.pi / 36, np.pi / 2, np.pi / 36), padding_mode="zeros"),
RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1,
rotate_range=(np.pi / 2, np.pi / 36, np.pi / 36), padding_mode="zeros"),
Rand3DElasticd(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1,
sigma_range=(5, 8), magnitude_range=(100, 200), scale_range=(0.15, 0.15, 0.15),
padding_mode="zeros"),
RandAdjustContrastd(keys=['image'], gamma=(0.5, 2.5), prob=0.1),
RandGaussianNoised(keys=['image'], prob=0.1, mean=np.random.uniform(0, 0.5), std=np.random.uniform(0, 1)),
RandShiftIntensityd(keys=['image'], offsets=np.random.uniform(0,0.3), prob=0.1),
RandSpatialCropd(keys=['image', 'label'], roi_size=opt.patch_size, random_size=False),
ToTensord(keys=['image', 'label'])
]
val_transforms = [
LoadImaged(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
NormalizeIntensityd(keys=['image']),
ScaleIntensityd(keys=['image']),
Spacingd(keys=['image', 'label'], pixdim=opt.resolution, mode=('bilinear', 'nearest')),
ToTensord(keys=['image', 'label'])
]
else:
train_transforms = [
LoadImaged(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
NormalizeIntensityd(keys=['image']),
ScaleIntensityd(keys=['image']),
RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=1),
RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=0),
RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=2),
RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1,
rotate_range=(np.pi / 36, np.pi / 36, np.pi * 2), padding_mode="zeros"),
RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1,
rotate_range=(np.pi / 36, np.pi / 2, np.pi / 36), padding_mode="zeros"),
RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1,
rotate_range=(np.pi / 2, np.pi / 36, np.pi / 36), padding_mode="zeros"),
Rand3DElasticd(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1,
sigma_range=(5, 8), magnitude_range=(100, 200), scale_range=(0.15, 0.15, 0.15), padding_mode="zeros"),
RandAdjustContrastd(keys=['image'], gamma=(0.5, 2.5), prob=0.1),
RandGaussianNoised(keys=['image'], prob=0.1, mean=np.random.uniform(0, 0.5), std=np.random.uniform(0, 1)),
RandShiftIntensityd(keys=['image'], offsets=np.random.uniform(0,0.3), prob=0.1),
RandSpatialCropd(keys=['image', 'label'], roi_size=opt.patch_size, random_size=False),
ToTensord(keys=['image', 'label'])
]
val_transforms = [
LoadImaged(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
NormalizeIntensityd(keys=['image']),
ScaleIntensityd(keys=['image']),
ToTensord(keys=['image', 'label'])
]
train_transforms = Compose(train_transforms)
val_transforms = Compose(val_transforms)
# create a training data loader
check_train = monai.data.Dataset(data=train_dicts, transform=train_transforms)
train_loader = DataLoader(check_train, batch_size=opt.batch_size, shuffle=True, num_workers=opt.workers, pin_memory=torch.cuda.is_available())
# create a training_dice data loader
check_val = monai.data.Dataset(data=train_dice_dicts, transform=val_transforms)
train_dice_loader = DataLoader(check_val, batch_size=1, num_workers=opt.workers, pin_memory=torch.cuda.is_available())
# create a validation data loader
check_val = monai.data.Dataset(data=val_dicts, transform=val_transforms)
val_loader = DataLoader(check_val, batch_size=1, num_workers=opt.workers, pin_memory=torch.cuda.is_available())
# create a validation data loader
check_val = monai.data.Dataset(data=test_dicts, transform=val_transforms)
test_loader = DataLoader(check_val, batch_size=1, num_workers=opt.workers, pin_memory=torch.cuda.is_available())
# # try to use all the available GPUs
# devices = get_devices_spec(None)
# build the network
net = build_net()
net.cuda()
if num_gpus > 1:
net = torch.nn.DataParallel(net)
if opt.preload is not None:
net.load_state_dict(torch.load(opt.preload))
dice_metric = DiceMetric(include_background=True, reduction="mean")
post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold_values=True)])
# loss_function = monai.losses.DiceLoss(sigmoid=True)
# loss_function = monai.losses.TverskyLoss(sigmoid=True, alpha=0.3, beta=0.7)
loss_function = monai.losses.DiceCELoss(sigmoid=True)
optim = torch.optim.Adam(net.parameters(), lr=opt.lr)
net_scheduler = get_scheduler(optim, opt)
# start a typical PyTorch training
val_interval = 1
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = list()
metric_values = list()
writer = SummaryWriter()
for epoch in range(opt.epochs):
print("-" * 10)
print(f"epoch {epoch + 1}/{opt.epochs}")
net.train()
epoch_loss = 0
step = 0
for batch_data in train_loader:
step += 1
inputs, labels = batch_data["image"].cuda(), batch_data["label"].cuda()
optim.zero_grad()
outputs = net(inputs)
loss = loss_function(outputs, labels)
loss.backward()
optim.step()
epoch_loss += loss.item()
epoch_len = len(check_train) // train_loader.batch_size
print(f"{step}/{epoch_len}, train_loss: {loss.item():.4f}")
writer.add_scalar("train_loss", loss.item(), epoch_len * epoch + step)
epoch_loss /= step
epoch_loss_values.append(epoch_loss)
print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")
update_learning_rate(net_scheduler, optim)
if (epoch + 1) % val_interval == 0:
net.eval()
with torch.no_grad():
def plot_dice(images_loader):
metric_sum = 0.0
metric_count = 0
val_images = None
val_labels = None
val_outputs = None
for data in images_loader:
val_images, val_labels = data["image"].cuda(), data["label"].cuda()
roi_size = opt.patch_size
sw_batch_size = 4
val_outputs = sliding_window_inference(val_images, roi_size, sw_batch_size, net)
val_outputs = post_trans(val_outputs)
value, _ = dice_metric(y_pred=val_outputs, y=val_labels)
metric_count += len(value)
metric_sum += value.item() * len(value)
metric = metric_sum / metric_count
metric_values.append(metric)
return metric, val_images, val_labels, val_outputs
metric, val_images, val_labels, val_outputs = plot_dice(val_loader)
# Save best model
if metric > best_metric:
best_metric = metric
best_metric_epoch = epoch + 1
torch.save(net.state_dict(), "best_metric_model.pth")
print("saved new best metric model")
metric_train, train_images, train_labels, train_outputs = plot_dice(train_dice_loader)
metric_test, test_images, test_labels, test_outputs = plot_dice(test_loader)
# Logger bar
print(
"current epoch: {} Training dice: {:.4f} Validation dice: {:.4f} Testing dice: {:.4f} Best Validation dice: {:.4f} at epoch {}".format(
epoch + 1, metric_train, metric, metric_test, best_metric, best_metric_epoch
)
)
writer.add_scalar("Mean_epoch_loss", epoch_loss, epoch + 1)
writer.add_scalar("Testing_dice", metric_test, epoch + 1)
writer.add_scalar("Training_dice", metric_train, epoch + 1)
writer.add_scalar("Validation_dice", metric, epoch + 1)
# plot the last model output as GIF image in TensorBoard with the corresponding image and label
val_outputs = (val_outputs.sigmoid() >= 0.5).float()
plot_2d_or_3d_image(val_images, epoch + 1, writer, index=0, tag="validation image")
plot_2d_or_3d_image(val_labels, epoch + 1, writer, index=0, tag="validation label")
plot_2d_or_3d_image(val_outputs, epoch + 1, writer, index=0, tag="validation inference")
print(f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}")
writer.close()
def _define_training_transforms(self):
"""Define and initialize all training data transforms.
* training set images transform
* training set masks transform
* validation set images transform
* validation set masks transform
* validation set images post-transform
* test set images transform
* test set masks transform
* test set images post-transform
* prediction set images transform
* prediction set images post-transform
@return True if data transforms could be instantiated, False otherwise.
"""
if self._mask_type == MaskType.UNKNOWN:
raise Exception("The mask type is unknown. Cannot continue!")
# Depending on the mask type, we will need to adapt the Mask Loader
# and Transform. We start by initializing the most common types.
MaskLoader = LoadMask(self._mask_type)
MaskTransform = Identity
# Adapt the transform for the LABEL types
if self._mask_type == MaskType.TIFF_LABELS or self._mask_type == MaskType.NUMPY_LABELS:
MaskTransform = ToOneHot(num_classes=self._out_channels)
# The H5_ONE_HOT type requires a different loader
if self._mask_type == MaskType.H5_ONE_HOT:
# MaskLoader: still missing
raise Exception("HDF5 one-hot masks are not supported yet!")
# Define transforms for training
self._train_image_transforms = Compose(
[
LoadImage(image_only=True),
ScaleIntensity(),
AddChannel(),
RandSpatialCrop(self._roi_size, random_size=False),
RandRotate90(prob=0.5, spatial_axes=(0, 1)),
ToTensor()
]
)
self._train_mask_transforms = Compose(
[
MaskLoader,
MaskTransform,
RandSpatialCrop(self._roi_size, random_size=False),
RandRotate90(prob=0.5, spatial_axes=(0, 1)),
ToTensor()
]
)
# Define transforms for validation
self._validation_image_transforms = Compose(
[
LoadImage(image_only=True),
ScaleIntensity(),
AddChannel(),
ToTensor()
]
)
self._validation_mask_transforms = Compose(
[
MaskLoader,
MaskTransform,
ToTensor()
]
)
# Define transforms for testing
self._test_image_transforms = Compose(
[
LoadImage(image_only=True),
ScaleIntensity(),
AddChannel(),
ToTensor()
]
)
self._test_mask_transforms = Compose(
[
MaskLoader,
MaskTransform,
ToTensor()
]
)
# Post transforms
self._validation_post_transforms = Compose(
[
Activations(softmax=True),
AsDiscrete(threshold_values=True)
]
)
self._test_post_transforms = Compose(
[
Activations(softmax=True),
AsDiscrete(threshold_values=True)
]
)
def main(tempdir):
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
print(f"generating synthetic data to {tempdir} (this may take a while)")
for i in range(5):
im, seg = create_test_image_3d(128, 128, 128, num_seg_classes=1, channel_dim=-1)
n = nib.Nifti1Image(im, np.eye(4))
nib.save(n, os.path.join(tempdir, f"im{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(tempdir, "im*.nii.gz")))
segs = sorted(glob(os.path.join(tempdir, "seg*.nii.gz")))
val_files = [{"img": img, "seg": seg} for img, seg in zip(images, segs)]
# define transforms for image and segmentation
val_transforms = Compose(
[
LoadImaged(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
ScaleIntensityd(keys="img"),
EnsureTyped(keys=["img", "seg"]),
]
)
val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = UNet(
spatial_dims=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
# define sliding window size and batch size for windows inference
roi_size = (96, 96, 96)
sw_batch_size = 4
post_trans = Compose([EnsureType(), Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
save_image = SaveImage(output_dir="tempdir", output_ext=".nii.gz", output_postfix="seg")
def _sliding_window_processor(engine, batch):
net.eval()
with torch.no_grad():
val_images, val_labels = batch["img"].to(device), batch["seg"].to(device)
seg_probs = sliding_window_inference(val_images, roi_size, sw_batch_size, net)
seg_probs = [post_trans(i) for i in decollate_batch(seg_probs)]
val_data = decollate_batch(batch["img_meta_dict"])
for seg_prob, data in zip(seg_probs, val_data):
save_image(seg_prob, data)
return seg_probs, val_labels
evaluator = Engine(_sliding_window_processor)
# add evaluation metric to the evaluator engine
MeanDice().attach(evaluator, "Mean_Dice")
# StatsHandler prints loss at every iteration and print metrics at every epoch,
# we don't need to print loss for evaluator, so just print metrics, user can also customize print functions
val_stats_handler = StatsHandler(
name="evaluator",
output_transform=lambda x: None, # no need to print loss value, so disable per iteration output
)
val_stats_handler.attach(evaluator)
# the model was trained by "unet_training_dict" example
CheckpointLoader(load_path="./runs_dict/net_checkpoint_50.pt", load_dict={"net": net}).attach(evaluator)
# sliding window inference for one image at every iteration
val_loader = DataLoader(
val_ds, batch_size=1, num_workers=4, collate_fn=list_data_collate, pin_memory=torch.cuda.is_available()
)
state = evaluator.run(val_loader)
print(state)
def main(tempdir):
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
print(f"generating synthetic data to {tempdir} (this may take a while)")
for i in range(5):
im, seg = create_test_image_3d(128, 128, 128, num_seg_classes=1, channel_dim=-1)
n = nib.Nifti1Image(im, np.eye(4))
nib.save(n, os.path.join(tempdir, f"im{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(tempdir, "im*.nii.gz")))
segs = sorted(glob(os.path.join(tempdir, "seg*.nii.gz")))
val_files = [{"img": img, "seg": seg} for img, seg in zip(images, segs)]
# define transforms for image and segmentation
val_transforms = Compose(
[
LoadNiftid(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
ScaleIntensityd(keys="img"),
ToTensord(keys=["img", "seg"]),
]
)
val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
# sliding window inference need to input 1 image in every iteration
val_loader = DataLoader(val_ds, batch_size=1, num_workers=4, collate_fn=list_data_collate)
dice_metric = DiceMetric(include_background=True, reduction="mean")
post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold_values=True)])
# try to use all the available GPUs
devices = get_devices_spec(None)
model = UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(devices[0])
model.load_state_dict(torch.load("best_metric_model_segmentation3d_dict.pth"))
# if we have multiple GPUs, set data parallel to execute sliding window inference
if len(devices) > 1:
model = torch.nn.DataParallel(model, device_ids=devices)
model.eval()
with torch.no_grad():
metric_sum = 0.0
metric_count = 0
saver = NiftiSaver(output_dir="./output")
for val_data in val_loader:
val_images, val_labels = val_data["img"].to(devices[0]), val_data["seg"].to(devices[0])
# define sliding window size and batch size for windows inference
roi_size = (96, 96, 96)
sw_batch_size = 4
val_outputs = sliding_window_inference(val_images, roi_size, sw_batch_size, model)
val_outputs = post_trans(val_outputs)
value, _ = dice_metric(y_pred=val_outputs, y=val_labels)
metric_count += len(value)
metric_sum += value.item() * len(value)
saver.save_batch(val_outputs, val_data["img_meta_dict"])
metric = metric_sum / metric_count
print("evaluation metric:", metric)
def evaluate(args):
# initialize Horovod library
hvd.init()
# Horovod limits CPU threads to be used per worker
torch.set_num_threads(1)
if hvd.local_rank() == 0 and not os.path.exists(args.dir):
# create 16 random image, mask paris for evaluation
print(f"generating synthetic data to {args.dir} (this may take a while)")
os.makedirs(args.dir)
# set random seed to generate same random data for every node
np.random.seed(seed=0)
for i in range(16):
im, seg = create_test_image_3d(128, 128, 128, num_seg_classes=1, channel_dim=-1)
n = nib.Nifti1Image(im, np.eye(4))
nib.save(n, os.path.join(args.dir, f"img{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(args.dir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(args.dir, "img*.nii.gz")))
segs = sorted(glob(os.path.join(args.dir, "seg*.nii.gz")))
val_files = [{"img": img, "seg": seg} for img, seg in zip(images, segs)]
# define transforms for image and segmentation
val_transforms = Compose(
[
LoadImaged(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
ScaleIntensityd(keys="img"),
EnsureTyped(keys=["img", "seg"]),
]
)
# create a evaluation data loader
val_ds = Dataset(data=val_files, transform=val_transforms)
# create a evaluation data sampler
val_sampler = DistributedSampler(val_ds, shuffle=False, num_replicas=hvd.size(), rank=hvd.rank())
# when supported, use "forkserver" to spawn dataloader workers instead of "fork" to prevent
# issues with Infiniband implementations that are not fork-safe
multiprocessing_context = None
if hasattr(mp, "_supports_context") and mp._supports_context and "forkserver" in mp.get_all_start_methods():
multiprocessing_context = "forkserver"
# sliding window inference need to input 1 image in every iteration
val_loader = DataLoader(
val_ds,
batch_size=1,
shuffle=False,
num_workers=2,
pin_memory=True,
sampler=val_sampler,
multiprocessing_context=multiprocessing_context,
)
dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
post_trans = Compose([EnsureType(), Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
# create UNet, DiceLoss and Adam optimizer
device = torch.device(f"cuda:{hvd.local_rank()}")
torch.cuda.set_device(device)
model = monai.networks.nets.UNet(
spatial_dims=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
if hvd.rank() == 0:
# load model parameters for evaluation
model.load_state_dict(torch.load("final_model.pth"))
# Horovod broadcasts parameters
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
model.eval()
with torch.no_grad():
for val_data in val_loader:
val_images, val_labels = val_data["img"].to(device), val_data["seg"].to(device)
# define sliding window size and batch size for windows inference
roi_size = (96, 96, 96)
sw_batch_size = 4
val_outputs = sliding_window_inference(val_images, roi_size, sw_batch_size, model)
val_outputs = [post_trans(i) for i in decollate_batch(val_outputs)]
dice_metric(y_pred=val_outputs, y=val_labels)
metric = dice_metric.aggregate().item()
dice_metric.reset()
if hvd.rank() == 0:
print("evaluation metric:", metric)
def test_value_shape(self, input_param, img, out, expected_shape):
result = Activations(**input_param)(img)
torch.testing.assert_allclose(result, out)
self.assertTupleEqual(result.shape, expected_shape)
def initialize(self, args):
"""
`initialize` is called only once when the model is being loaded.
Implementing `initialize` function is optional. This function allows
the model to intialize any state associated with this model.
"""
# Pull model from google drive
extract_dir = "/models/monai_covid/1"
tar_save_path = os.path.join(extract_dir, model_filename)
download_and_extract(gdrive_url,
tar_save_path,
output_dir=extract_dir,
hash_val=md5_check,
hash_type="md5")
# load model configuration
self.model_config = json.loads(args['model_config'])
# create inferer engine and load PyTorch model
inference_device_kind = args.get('model_instance_kind', None)
logger.info(f"Inference device: {inference_device_kind}")
self.inference_device = torch.device('cpu')
if inference_device_kind is None or inference_device_kind == 'CPU':
self.inference_device = torch.device('cpu')
elif inference_device_kind == 'GPU':
inference_device_id = args.get('model_instance_device_id', '0')
logger.info(f"Inference device id: {inference_device_id}")
if torch.cuda.is_available():
self.inference_device = torch.device(
f'cuda:{inference_device_id}')
cudnn.enabled = True
else:
logger.error(
f"No CUDA device detected. Using device: {inference_device_kind}"
)
# create pre-transforms
self.pre_transforms = Compose([
LoadImage(reader="NibabelReader",
image_only=True,
dtype=np.float32),
AddChannel(),
ScaleIntensityRange(a_min=-1000,
a_max=500,
b_min=0.0,
b_max=1.0,
clip=True),
CropForeground(margin=5),
Resize([192, 192, 64], mode="area"),
AddChannel(),
ToTensor(),
Lambda(func=lambda x: x.to(device=self.inference_device)),
])
# create post-transforms
self.post_transforms = Compose([
Lambda(func=lambda x: x.to(device="cpu")),
Activations(sigmoid=True),
ToNumpy(),
AsDiscrete(threshold_values=True, logit_thresh=0.5),
])
self.inferer = SimpleInferer()
self.model = torch.jit.load(
f'{pathlib.Path(os.path.realpath(__file__)).parent}{os.path.sep}covid19_model.ts',
map_location=self.inference_device)
