Update tempdir logic in examples and fix CI issue

examples/segmentation_3d/unet_evaluation_array.py

@@ -28,62 +28,62 @@
 from monai.transforms import AddChannel, Compose, ScaleIntensity, ToTensor
 
 
-def main():
+def main(tempdir):
     config.print_config()
     logging.basicConfig(stream=sys.stdout, level=logging.INFO)
 
-    with tempfile.TemporaryDirectory() as tempdir:
-        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)
+    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(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"))
+        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")))
+    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())
-        dice_metric = DiceMetric(include_background=True, to_onehot_y=False, sigmoid=True, reduction="mean")
+    # 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())
+    dice_metric = DiceMetric(include_background=True, to_onehot_y=False, sigmoid=True, reduction="mean")
 
-        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)
+    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.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 = 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).float()
-                saver.save_batch(val_outputs, val_data[2])
-            metric = metric_sum / metric_count
-            print("evaluation metric:", metric)
+    model.load_state_dict(torch.load("best_metric_model.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)
+            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).float()
+            saver.save_batch(val_outputs, val_data[2])
+        metric = metric_sum / metric_count
+        print("evaluation metric:", metric)
 
 
 if __name__ == "__main__":
-    main()
+    with tempfile.TemporaryDirectory() as tempdir:
+        main(tempdir)

examples/segmentation_3d/unet_evaluation_dict.py

@@ -29,75 +29,75 @@
 from monai.transforms import AsChannelFirstd, Compose, LoadNiftid, ScaleIntensityd, ToTensord
 
 
-def main():
+def main(tempdir):
     monai.config.print_config()
     logging.basicConfig(stream=sys.stdout, level=logging.INFO)
 
-    with tempfile.TemporaryDirectory() as tempdir:
-        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, to_onehot_y=False, sigmoid=True, reduction="mean")
-
-        # 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.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)
-                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).float()
-                saver.save_batch(val_outputs, val_data["img_meta_dict"])
-            metric = metric_sum / metric_count
-            print("evaluation metric:", metric)
+    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, to_onehot_y=False, sigmoid=True, reduction="mean")
+
+    # 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.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)
+            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).float()
+            saver.save_batch(val_outputs, val_data["img_meta_dict"])
+        metric = metric_sum / metric_count
+        print("evaluation metric:", metric)
 
 
 if __name__ == "__main__":
-    main()
+    with tempfile.TemporaryDirectory() as tempdir:
+        main(tempdir)