(root_dir, device="cuda:0", amp=False, num_workers=4)
| 62 | |
| 63 | |
| 64 | def run_training_test(root_dir, device="cuda:0", amp=False, num_workers=4): |
| 65 | images = sorted(glob(os.path.join(root_dir, "img*.nii.gz"))) |
| 66 | segs = sorted(glob(os.path.join(root_dir, "seg*.nii.gz"))) |
| 67 | train_files = [{"image": img, "label": seg} for img, seg in zip(images[:20], segs[:20])] |
| 68 | val_files = [{"image": img, "label": seg} for img, seg in zip(images[-20:], segs[-20:])] |
| 69 | |
| 70 | # define transforms for image and segmentation |
| 71 | train_transforms = Compose( |
| 72 | [ |
| 73 | LoadImaged(keys=["image", "label"]), |
| 74 | EnsureChannelFirstd(keys=["image", "label"], channel_dim=-1), |
| 75 | ScaleIntensityd(keys=["image", "label"]), |
| 76 | RandCropByPosNegLabeld( |
| 77 | keys=["image", "label"], label_key="label", spatial_size=[96, 96, 96], pos=1, neg=1, num_samples=4 |
| 78 | ), |
| 79 | RandRotate90d(keys=["image", "label"], prob=0.5, spatial_axes=[0, 2]), |
| 80 | ] |
| 81 | ) |
| 82 | val_transforms = Compose( |
| 83 | [ |
| 84 | LoadImaged(keys=["image", "label"]), |
| 85 | EnsureChannelFirstd(keys=["image", "label"], channel_dim=-1), |
| 86 | ScaleIntensityd(keys=["image", "label"]), |
| 87 | ] |
| 88 | ) |
| 89 | |
| 90 | # create a training data loader |
| 91 | train_ds = monai.data.CacheDataset(data=train_files, transform=train_transforms, cache_rate=0.5) |
| 92 | # use batch_size=2 to load images and use RandCropByPosNegLabeld to generate 2 x 4 images for network training |
| 93 | train_loader = monai.data.DataLoader(train_ds, batch_size=2, shuffle=True, num_workers=num_workers) |
| 94 | # create a validation data loader |
| 95 | val_ds = monai.data.CacheDataset(data=val_files, transform=val_transforms, cache_rate=1.0) |
| 96 | val_loader = monai.data.DataLoader(val_ds, batch_size=1, num_workers=num_workers) |
| 97 | |
| 98 | # create UNet, DiceLoss and Adam optimizer |
| 99 | net = monai.networks.nets.UNet( |
| 100 | spatial_dims=3, |
| 101 | in_channels=1, |
| 102 | out_channels=1, |
| 103 | channels=(16, 32, 64, 128, 256), |
| 104 | strides=(2, 2, 2, 2), |
| 105 | num_res_units=2, |
| 106 | ).to(device) |
| 107 | loss = monai.losses.DiceLoss(sigmoid=True) |
| 108 | opt = torch.optim.Adam(net.parameters(), 1e-3) |
| 109 | lr_scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=2, gamma=0.1) |
| 110 | summary_writer = SummaryWriter(log_dir=root_dir) |
| 111 | |
| 112 | val_postprocessing = Compose( |
| 113 | [ |
| 114 | Activationsd(keys="pred", sigmoid=True), |
| 115 | AsDiscreted(keys="pred", threshold=0.5), |
| 116 | KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]), |
| 117 | ] |
| 118 | ) |
| 119 | |
| 120 | class _TestEvalIterEvents: |
| 121 | def attach(self, engine): |
no test coverage detected
searching dependent graphs…