Function run_test

tests/integration/test_integration_sliding_window.py:33–67

(batch_size, img_name, seg_name, output_dir, device="cuda:0")

def run_test(batch_size, img_name, seg_name, output_dir, device="cuda:0"):
    ds = ImageDataset(
        [img_name],
        [seg_name],
        transform=EnsureChannelFirst(channel_dim="no_channel"),
        seg_transform=EnsureChannelFirst(channel_dim="no_channel"),
        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):
        for m in engine.state.output:
            saver(m)

    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


@skip_if_quick