def _define_model(self) -> None:
"""Instantiate the U-Net architecture."""
# Create U-Net
self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device '{self._device}'.", file=self._stdout)
# Try to free memory on the GPU
if self._device != "cpu":
torch.cuda.empty_cache()
# Instantiate the requested model
if self._option_architecture == SegmentationArchitectures.ResidualUNet2D:
# Monai's UNet
self._model = UNet(
dimensions=2,
in_channels=self._in_channels,
out_channels=self._out_channels,
channels=tuple((self._num_filters_in_first_layer * 2**i for i in range(0, 5))),
strides=(2, 2, 2, 2),
num_res_units=2
).to(self._device)
elif self._option_architecture == SegmentationArchitectures.AttentionUNet2D:
# Attention U-Net
self._model = AttentionUNet2D(
img_ch=self._in_channels,
output_ch=self._out_channels,
n1=self._num_filters_in_first_layer
).to(self._device)
else:
raise ValueError(f"Unexpected architecture {self._option_architecture}! Aborting.")
def setUp(self):
self.bundle_dir = tempfile.TemporaryDirectory()
self.dir_name = os.path.join(self.bundle_dir.name, "TestBundle")
self.configs_name = os.path.join(self.dir_name, "configs")
self.models_name = os.path.join(self.dir_name, "models")
self.metadata_name = os.path.join(self.configs_name, "metadata.json")
self.test_name = os.path.join(self.configs_name, "test.json")
self.modelpt_name = os.path.join(self.models_name, "model.pt")
self.zip_file = os.path.join(self.bundle_dir.name, "TestBundle.zip")
self.ts_file = os.path.join(self.bundle_dir.name, "TestBundle.ts")
# create the directories for the bundle
os.mkdir(self.dir_name)
os.mkdir(self.configs_name)
os.mkdir(self.models_name)
# fill bundle configs
with open(self.metadata_name, "w") as o:
o.write(metadata)
with open(self.test_name, "w") as o:
o.write(test_json)
# save network
net = UNet(2, 1, 1, [4, 8], [2])
torch.save(net.state_dict(), self.modelpt_name)
def run_test(batch_size=64, train_steps=100, device=torch.device("cuda:0")):
class _TestBatch(Dataset):
def __getitem__(self, _unused_id):
im, seg = create_test_image_2d(128,
128,
noise_max=1,
num_objs=4,
num_seg_classes=1)
return im[None], seg[None].astype(np.float32)
def __len__(self):
return train_steps
net = UNet(
dimensions=2,
in_channels=1,
out_channels=1,
channels=(4, 8, 16, 32),
strides=(2, 2, 2),
num_res_units=2,
).to(device)
loss = DiceLoss(do_sigmoid=True)
opt = torch.optim.Adam(net.parameters(), 1e-4)
src = DataLoader(_TestBatch(), batch_size=batch_size)
trainer = create_supervised_trainer(net, opt, loss, device, False)
trainer.run(src, 1)
loss = trainer.state.output
return loss
def test_epistemic_scoring(self):
input_size = (20, 20, 20)
device = "cuda" if torch.cuda.is_available() else "cpu"
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)
transforms = Compose([
AddChanneld(keys),
CropForegroundd(keys, source_key="image"),
DivisiblePadd(keys, 4),
])
infer_transforms = Compose([
AddChannel(),
CropForeground(),
DivisiblePad(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(3, 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)
entropy_score = EpistemicScoring(model=model,
transforms=infer_transforms,
roi_size=[20, 20, 20],
num_samples=10)
# Call Individual Infer from Epistemic Scoring
ip_stack = [test_data["image"], test_data["image"], test_data["image"]]
ip_stack = np.array(ip_stack)
score_3d = entropy_score.entropy_3d_volume(ip_stack)
score_3d_sum = np.sum(score_3d)
# Call Entropy Metric from Epistemic Scoring
self.assertEqual(score_3d.shape, input_size)
self.assertIsInstance(score_3d_sum, np.float32)
self.assertGreater(score_3d_sum, 3.0)
def run_test(batch_size=64, train_steps=200, device=torch.device("cuda:0")):
class _TestBatch(Dataset):
def __init__(self, transforms):
self.transforms = transforms
def __getitem__(self, _unused_id):
im, seg = create_test_image_2d(128,
128,
noise_max=1,
num_objs=4,
num_seg_classes=1)
seed = np.random.randint(2147483647)
self.transforms.set_random_state(seed=seed)
im = self.transforms(im)
self.transforms.set_random_state(seed=seed)
seg = self.transforms(seg)
return im, seg
def __len__(self):
return train_steps
net = UNet(
dimensions=2,
in_channels=1,
out_channels=1,
channels=(4, 8, 16, 32),
strides=(2, 2, 2),
num_res_units=2,
).to(device)
loss = DiceLoss(do_sigmoid=True)
opt = torch.optim.Adam(net.parameters(), 1e-2)
train_transforms = Compose([
AddChannel(),
ScaleIntensity(),
RandSpatialCrop((96, 96), random_size=False),
RandRotate90(),
ToTensor()
])
src = DataLoader(_TestBatch(train_transforms),
batch_size=batch_size,
shuffle=True)
net.train()
epoch_loss = 0
step = 0
for img, seg in src:
step += 1
opt.zero_grad()
output = net(img.to(device))
step_loss = loss(output, seg.to(device))
step_loss.backward()
opt.step()
epoch_loss += step_loss.item()
epoch_loss /= step
return epoch_loss, step
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"]),
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))
with eval_mode(model):
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 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_ill_input_shape(self):
net = UNet(spatial_dims=2,
in_channels=1,
out_channels=3,
channels=(16, 32, 64),
strides=(2, 2))
with eval_mode(net):
with self.assertRaisesRegex(RuntimeError,
"Sizes of tensors must match"):
net.forward(torch.randn(2, 1, 16, 5))
def test_bundle(self):
with tempfile.TemporaryDirectory() as tempdir:
net = UNet(2, 1, 1, [4, 8], [2])
torch.save(net.state_dict(), tempdir + "/test.pt")
bundle_root = tempdir + "/test_bundle"
cmd = ["coverage", "run", "-m", "monai.bundle", "init_bundle", bundle_root, tempdir + "/test.pt"]
subprocess.check_call(cmd)
self.assertTrue(os.path.exists(bundle_root + "/configs/metadata.json"))
self.assertTrue(os.path.exists(bundle_root + "/configs/inference.json"))
self.assertTrue(os.path.exists(bundle_root + "/models/model.pt"))
def __init__(self, hparams):
super().__init__()
self.hparams = hparams
self.unet = UNet(
dimensions=3,
in_channels=1,
out_channels=2,
channels=(64, 128, 258, 512, 1024),
strides=(2, 2, 2, 2),
norm=monai.networks.layers.Norm.BATCH,
dropout=0,
)
self.sample_masks = []
def scoring_method(
self) -> Union[None, ScoringMethod, Dict[str, ScoringMethod]]:
methods: Dict[str, ScoringMethod] = {
"dice": Dice(),
"sum": Sum(),
}
if self.epistemic_enabled:
methods[f"{self.name}_epistemic"] = EpistemicScoring(
model=self.path,
network=UNet(
spatial_dims=3,
in_channels=1,
out_channels=2,
channels=[16, 32, 64, 128, 256],
strides=[2, 2, 2, 2],
num_res_units=2,
norm="batch",
dropout=0.2,
),
transforms=lib.infers.SegmentationSpleen(
None).pre_transforms(),
num_samples=self.epistemic_samples,
)
if self.tta_enabled:
methods[f"{self.name}_tta"] = TTAScoring(
model=self.path,
network=self.network,
deepedit=False,
num_samples=self.tta_samples,
spatial_size=(128, 128, 64),
spacing=(1.0, 1.0, 1.0),
)
return methods
def get_network(network, channels, dimensions):
if network == 'unet':
if channels == 16:
features = (16, 32, 64, 128, 256)
elif channels == 32:
features = (32, 64, 128, 256, 512)
else:
features = (64, 128, 256, 512, 1024)
logging.info('Using Unet with features: {}'.format(features))
network = UNet(dimensions=dimensions,
in_channels=3,
out_channels=1,
channels=features,
strides=[2, 2, 2, 2],
norm=Norm.BATCH)
else:
if channels == 16:
features = (16, 32, 64, 128, 256, 16)
elif channels == 32:
features = (32, 64, 128, 256, 512, 32)
else:
features = (64, 128, 256, 512, 1024, 64)
logging.info('Using BasicUnet with features: {}'.format(features))
network = BasicUNet(dimensions=dimensions,
in_channels=3,
out_channels=1,
features=features)
return network
def run_test(batch_size, img_name, seg_name, output_dir, device="cuda:0"):
ds = NiftiDataset([img_name], [seg_name], transform=AddChannel(), seg_transform=AddChannel(), image_only=False)
loader = DataLoader(ds, batch_size=1, pin_memory=torch.cuda.is_available())
net = UNet(
dimensions=3, in_channels=1, out_channels=1, channels=(4, 8, 16, 32), strides=(2, 2, 2), num_res_units=2
).to(device)
roi_size = (16, 32, 48)
sw_batch_size = batch_size
def _sliding_window_processor(_engine, batch):
net.eval()
img, seg, meta_data = batch
with torch.no_grad():
seg_probs = sliding_window_inference(img.to(device), roi_size, sw_batch_size, net, device=device)
return predict_segmentation(seg_probs)
infer_engine = Engine(_sliding_window_processor)
SegmentationSaver(
output_dir=output_dir, output_ext=".nii.gz", output_postfix="seg", batch_transform=lambda x: x[2]
).attach(infer_engine)
infer_engine.run(loader)
basename = os.path.basename(img_name)[: -len(".nii.gz")]
saved_name = os.path.join(output_dir, basename, f"{basename}_seg.nii.gz")
return saved_name
def test_script(self):
net = UNet(dimensions=2,
in_channels=1,
out_channels=3,
channels=(16, 32, 64),
strides=(2, 2),
num_res_units=0)
test_data = torch.randn(16, 1, 32, 32)
test_script_save(net, test_data)
def test_inverse_inferred_seg(self):
test_data = []
for _ in range(20):
image, label = create_test_image_2d(100, 101)
test_data.append({
"image": image,
"label": label.astype(np.float32)
})
batch_size = 10
# num workers = 0 for mac
num_workers = 2 if sys.platform != "darwin" else 0
transforms = Compose([
AddChanneld(KEYS),
SpatialPadd(KEYS, (150, 153)),
CenterSpatialCropd(KEYS, (110, 99))
])
num_invertible_transforms = sum(1 for i in transforms.transforms
if isinstance(i, InvertibleTransform))
dataset = CacheDataset(test_data, transform=transforms, progress=False)
loader = DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = UNet(
dimensions=2,
in_channels=1,
out_channels=1,
channels=(2, 4),
strides=(2, ),
).to(device)
data = first(loader)
labels = data["label"].to(device)
segs = model(labels).detach().cpu()
label_transform_key = "label" + InverseKeys.KEY_SUFFIX.value
segs_dict = {
"label": segs,
label_transform_key: data[label_transform_key]
}
segs_dict_decollated = decollate_batch(segs_dict)
# inverse of individual segmentation
seg_dict = first(segs_dict_decollated)
with allow_missing_keys_mode(transforms):
inv_seg = transforms.inverse(seg_dict)["label"]
self.assertEqual(len(data["label_transforms"]),
num_invertible_transforms)
self.assertEqual(len(seg_dict["label_transforms"]),
num_invertible_transforms)
self.assertEqual(inv_seg.shape[1:], test_data[0]["label"].shape)
def test_inverse_inferred_seg(self, extra_transform):
test_data = []
for _ in range(20):
image, label = create_test_image_2d(100, 101)
test_data.append({
"image": image,
"label": label.astype(np.float32)
})
batch_size = 10
# num workers = 0 for mac
num_workers = 2 if sys.platform == "linux" else 0
transforms = Compose([
AddChanneld(KEYS),
SpatialPadd(KEYS, (150, 153)), extra_transform
])
dataset = CacheDataset(test_data, transform=transforms, progress=False)
loader = DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = UNet(spatial_dims=2,
in_channels=1,
out_channels=1,
channels=(2, 4),
strides=(1, )).to(device)
data = first(loader)
self.assertEqual(data["image"].shape[0], batch_size * NUM_SAMPLES)
labels = data["label"].to(device)
self.assertIsInstance(labels, MetaTensor)
segs = model(labels).detach().cpu()
segs_decollated = decollate_batch(segs)
self.assertIsInstance(segs_decollated[0], MetaTensor)
# inverse of individual segmentation
seg_metatensor = first(segs_decollated)
# test to convert interpolation mode for 1 data of model output batch
convert_applied_interp_mode(seg_metatensor.applied_operations,
mode="nearest",
align_corners=None)
# manually invert the last crop samples
xform = seg_metatensor.applied_operations.pop(-1)
shape_before_extra_xform = xform["orig_size"]
resizer = ResizeWithPadOrCrop(spatial_size=shape_before_extra_xform)
with resizer.trace_transform(False):
seg_metatensor = resizer(seg_metatensor)
with allow_missing_keys_mode(transforms):
inv_seg = transforms.inverse({"label": seg_metatensor})["label"]
self.assertEqual(inv_seg.shape[1:], test_data[0]["label"].shape)
def _define_model(self) -> None:
"""Instantiate the U-Net architecture."""
# Create U-Net
self._device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device '{self._device}'.", file=self._stdout)
# Try to free memory on the GPU
if self._device != "cpu":
torch.cuda.empty_cache()
# Monai's UNet
self._model = UNet(dimensions=2,
in_channels=self._in_channels,
out_channels=self._out_channels,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2).to(self._device)
def test_script(self):
net = UNet(dimensions=2,
in_channels=1,
out_channels=3,
channels=(16, 32, 64),
strides=(2, 2),
num_res_units=0)
test_data = torch.randn(16, 1, 32, 32)
out_orig, out_reloaded = test_script_save(net, test_data)
assert torch.allclose(out_orig, out_reloaded)
def UNet3D():
unet = UNet(
dimensions=3,
in_channels=1,
out_channels=2,
channels=(32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
norm=Norm.BATCH,
)
return unet
def test_shape(self, spatial_dim):
spatial_dim = int(spatial_dim)
model = UNet(
spatial_dims=2, in_channels=1, out_channels=1, channels=(4, 8, 16), strides=(2, 2), num_res_units=2
)
device = "cuda:0" if torch.cuda.is_available() else "cpu"
model.to(device)
model.eval()
# Initialize a dummy 3D tensor volume with shape (N,C,D,H,W)
input_volume = torch.ones(1, 1, 64, 256, 256, device=device)
# Remove spatial dim to slide across from the roi_size
roi_size = list(input_volume.shape[2:])
roi_size.pop(spatial_dim)
# Initialize and run inferer
inferer = SliceInferer(roi_size=roi_size, spatial_dim=spatial_dim, sw_batch_size=1, cval=-1)
result = inferer(input_volume, model)
self.assertTupleEqual(result.shape, input_volume.shape)
# test that the inferer can be run multiple times
result = inferer(input_volume, model)
def run_inference_test(root_dir, device=torch.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),
ScaleIntensityd(keys=["img", "seg"]),
ToTensord(keys=["img", "seg"]),
])
val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
# sliding window inferene need to input 1 image in every iteration
val_loader = DataLoader(val_ds,
batch_size=1,
num_workers=4,
collate_fn=list_data_collate,
pin_memory=torch.cuda.is_available())
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
saver = NiftiSaver(output_dir=os.path.join(root_dir, "output"),
dtype=int)
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)
value = compute_meandice(y_pred=val_outputs,
y=val_labels,
include_background=True,
to_onehot_y=False,
add_sigmoid=True)
metric_count += len(value)
metric_sum += value.sum().item()
val_outputs = (val_outputs.sigmoid() >= 0.5).float()
saver.save_batch(
val_outputs, {
"filename_or_obj": val_data["img.filename_or_obj"],
"affine": val_data["img.affine"]
})
metric = metric_sum / metric_count
return metric
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_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")))
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"]),
AddChanneld(keys=["img", "seg"]),
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, to_onehot_y=False, sigmoid=True, reduction="mean")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = 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)
model.load_state_dict(torch.load("best_metric_model_segmentation2d_dict.pth"))
model.eval()
with torch.no_grad():
metric_sum = 0.0
metric_count = 0
saver = PNGSaver(output_dir="./output")
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)
sw_batch_size = 4
val_outputs = sliding_window_inference(val_images, roi_size, sw_batch_size, model)
value = dice_metric(y_pred=val_outputs, y=val_labels)
metric_count += len(value)
metric_sum += value.item() * len(value)
val_outputs = val_outputs.sigmoid() >= 0.5
saver.save_batch(val_outputs)
metric = metric_sum / metric_count
print("evaluation metric:", metric)
def __init__(self, hparams):
super().__init__()
self.hparams = hparams
self.generator = UNet(
dimensions=3,
in_channels=1,
out_channels=2,
channels=(64, 128, 258, 512, 1024),
strides=(2, 2, 2, 2),
norm=monai.networks.layers.Norm.BATCH,
dropout=0,
)
self.discriminator = Discriminator(
in_shape=self.hparams.patch_size,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
norm=monai.networks.layers.Norm.BATCH,
)
self.generated_masks = None
self.sample_masks = []
def test_script_without_running_stats(self):
net = UNet(
spatial_dims=2,
in_channels=1,
out_channels=3,
channels=(16, 32, 64),
strides=(2, 2),
num_res_units=0,
norm=("batch", {
"track_running_stats": False
}),
)
test_data = torch.randn(16, 1, 16, 4)
test_script_save(net, test_data)
def main():
config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
tempdir = tempfile.mkdtemp()
print('generating synthetic data to {} (this may take a while)'.format(tempdir))
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, 'im%i.nii.gz' % i))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, 'seg%i.nii.gz' % i))
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 = NiftiDataset(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())
device = torch.device('cuda:0')
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.pth'))
model.eval()
with torch.no_grad():
metric_sum = 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)
value = compute_meandice(y_pred=val_outputs, y=val_labels, include_background=True,
to_onehot_y=False, add_sigmoid=True)
metric_count += len(value)
metric_sum += value.sum().item()
val_outputs = (val_outputs.sigmoid() >= 0.5).float()
saver.save_batch(val_outputs, val_data[2])
metric = metric_sum / metric_count
print('evaluation metric:', metric)
shutil.rmtree(tempdir)
def init(self, name: str, model_dir: str, conf: Dict[str, str],
planner: Any, **kwargs):
super().init(name, model_dir, conf, planner, **kwargs)
# Labels
self.labels = {
"spleen": 1,
"right kidney": 2,
"left kidney": 3,
"gallbladder": 4,
"esophagus": 5,
"liver": 6,
"stomach": 7,
"aorta": 8,
"inferior vena cava": 9,
"portal vein and splenic vein": 10,
"pancreas": 11,
"right adrenal gland": 12,
"left adrenal gland": 13,
}
# Model Files
self.path = [
os.path.join(self.model_dir,
f"pretrained_{name}.pt"), # pretrained
os.path.join(self.model_dir, f"{name}.pt"), # published
]
# Download PreTrained Model
if strtobool(self.conf.get("use_pretrained_model", "true")):
url = f"{self.PRE_TRAINED_PATH}/segmentation_unet_multilabel.pt"
download_file(url, self.path[0])
self.target_spacing = (1.0, 1.0, 1.0) # target space for image
self.spatial_size = (96, 96, 96) # train input size
self.roi_size = (128, 128, 128) # sliding window size for infer
# Network
self.network = UNet(
spatial_dims=3,
in_channels=1,
out_channels=len(self.labels.keys()) +
1, # All labels plus background
channels=[16, 32, 64, 128, 256],
strides=[2, 2, 2, 2],
num_res_units=2,
norm="batch",
)
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_inverse_inferred_seg(self, extra_transform):
test_data = []
for _ in range(20):
image, label = create_test_image_2d(100, 101)
test_data.append({"image": image, "label": label.astype(np.float32)})
batch_size = 10
# num workers = 0 for mac
num_workers = 2 if sys.platform == "linux" else 0
transforms = Compose([AddChanneld(KEYS), SpatialPadd(KEYS, (150, 153)), extra_transform])
num_invertible_transforms = sum(1 for i in transforms.transforms if isinstance(i, InvertibleTransform))
dataset = CacheDataset(test_data, transform=transforms, progress=False)
loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = UNet(spatial_dims=2, in_channels=1, out_channels=1, channels=(2, 4), strides=(2,)).to(device)
data = first(loader)
self.assertEqual(len(data["label_transforms"]), num_invertible_transforms)
self.assertEqual(data["image"].shape[0], batch_size * NUM_SAMPLES)
labels = data["label"].to(device)
segs = model(labels).detach().cpu()
label_transform_key = "label" + InverseKeys.KEY_SUFFIX
segs_dict = {"label": segs, label_transform_key: data[label_transform_key]}
segs_dict_decollated = decollate_batch(segs_dict)
# inverse of individual segmentation
seg_dict = first(segs_dict_decollated)
# test to convert interpolation mode for 1 data of model output batch
convert_inverse_interp_mode(seg_dict, mode="nearest", align_corners=None)
with allow_missing_keys_mode(transforms):
inv_seg = transforms.inverse(seg_dict)["label"]
self.assertEqual(len(data["label_transforms"]), num_invertible_transforms)
self.assertEqual(len(seg_dict["label_transforms"]), num_invertible_transforms)
self.assertEqual(inv_seg.shape[1:], test_data[0]["label"].shape)
# Inverse of batch
batch_inverter = BatchInverseTransform(transforms, loader, collate_fn=no_collation, detach=True)
with allow_missing_keys_mode(transforms):
inv_batch = batch_inverter(segs_dict)
self.assertEqual(inv_batch[0]["label"].shape[1:], test_data[0]["label"].shape)
def run_test(batch_size, img_name, seg_name, output_dir, device="cuda:0"):
ds = ImageDataset([img_name], [seg_name],
transform=AddChannel(),
seg_transform=AddChannel(),
image_only=True)
loader = DataLoader(ds, batch_size=1, pin_memory=torch.cuda.is_available())
net = UNet(spatial_dims=3,
in_channels=1,
out_channels=1,
channels=(4, 8, 16, 32),
strides=(2, 2, 2),
num_res_units=2).to(device)
roi_size = (16, 32, 48)
sw_batch_size = batch_size
saver = SaveImage(output_dir=output_dir,
output_ext=".nii.gz",
output_postfix="seg")
def _sliding_window_processor(_engine, batch):
img = batch[0] # first item from ImageDataset is the input image
with eval_mode(net):
seg_probs = sliding_window_inference(img.to(device),
roi_size,
sw_batch_size,
net,
device=device)
return predict_segmentation(seg_probs)
def save_func(engine):
if pytorch_after(1, 9, 1):
for m in engine.state.output:
saver(m)
else:
saver(engine.state.output[0])
infer_engine = Engine(_sliding_window_processor)
infer_engine.add_event_handler(Events.ITERATION_COMPLETED, save_func)
infer_engine.run(loader)
basename = os.path.basename(img_name)[:-len(".nii.gz")]
saved_name = os.path.join(output_dir, basename, f"{basename}_seg.nii.gz")
return saved_name
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()
