Borders Appear When Aggregating Patches Using GridAggregator

[isurusuranga]

I am new to TorchIO and am currently attempting to train a GAN model using overlapping patches extracted from 3D medical images. For my experiment, I used patch dimensions of (64, 64, 64) and an overlap of (16, 16, 16). To achieve this, I utilized the TorchIO package to prepare the dataset before feeding it into the deep learning model. Specifically, I used the GridSampler to extract patches from the entire volume during training. Below is the code snippet showing how I used the GridSampler and SubjectsLoader for the training process:

# Create a GridSampler for patch extraction
train_samplers = [tio.GridSampler(subject, patch_size, patch_overlap) for subject in train_dataset]
train_patches_loaders = [tio.SubjectsLoader(sampler, batch_size=4) for sampler in train_samplers]

for epoch in range(start_epoch, epochs):
    for loader in train_patches_loaders:
        for patches_batch in loader:
            x =  patches_batch['input'][tio.DATA].to(device) # Shape: [batch_size, 1, 64, 64, 64]
            real = patches_batch['target'][tio.DATA].to(device) # Shape: [batch_size, 1, 64, 64, 64]

Once the model was trained, I utilized it to predict the reconstructed image volumes by aggregating the patches using GridAggregator as shown below:

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# Initialize generator model
netG = Generator().to(device)

# Load checkpoint
checkpoint = torch.load(checkpoint_save_path)
netG.load_state_dict(checkpoint['netG_state_dict'])
netG.eval()

patch_size = (64, 64, 64) 
patch_overlap = (16, 16, 16)

for test_lf_img_path, test_hf_image_path in zip(test_lf_image_paths, test_hf_image_paths):
    # Create paired subject
    test_subject = tio.Subject(
        input=tio.ScalarImage(test_lf_img_path),
        target=tio.ScalarImage(test_hf_image_path)
    )

    original_min, original_max = test_subject['input'].data.min(), test_subject['input'].data.max()

    transformed_test_subject = validation_transform(test_subject)
    
    test_sampler = tio.inference.GridSampler(transformed_test_subject, patch_size, patch_overlap)
    test_patches_loader = tio.SubjectsLoader(test_sampler, batch_size=1)
    
    original_image = nib.load(test_ulf_img_path)
    original_affine = original_image.affine  # Affine matrix of the original image
    original_header = original_image.header  # Header of the original image
    
    aggregator = tio.inference.GridAggregator(test_sampler)

    with torch.no_grad():
        for patches_batch in test_patches_loader:
            # Run the model on input patches
            input_tensor = patches_batch['input'][tio.DATA].to(device)
            locations = patches_batch[tio.LOCATION]
            predicted_patch = netG(input_tensor)
            aggregator.add_batch(predicted_patch, locations)

    reconstructed_volume = aggregator.get_output_tensor()

    # Reverse transformations to the original intensity range
    original_intensity_volume = reverse_transforms(reconstructed_volume[0], original_min, original_max)

    reconstructed_nifti = nib.Nifti1Image(
        original_intensity_volume.numpy(), 
        affine=original_affine, 
        header=original_header
    )
    nib.save(reconstructed_nifti, "reconstructed_volume.nii.gz")

    break

The output for a given input MRI image volume, however, displays visible border lines corresponding to the patch boundaries. This issue is evident in the aggregated output, as the attached image shows.

How can I produce a smooth reconstructed image volume without these border artefacts? Any guidance or suggestions on resolving this issue would be greatly appreciated.

[isurusuranga]

Can I ask for a solution to my issue, please? Thanks in advance for your time and consideration.

[romainVala]

Hello
GridAggregator is by default with overlap_mode = 'crop', which is the worst strategie to reduce patch discontinuity
Try to use
overlap_mode = 'average' or 'hann'
it should improve the results