N4BiasFieldCorrectionImageFilter: GetLogBiasFieldAsImage, wrap BSplineControlPointImageFilter, alternatives

With SimpleITK I can do:

def bias_correct(
    input: sitk.Image,
    mask: sitk.Image = None,
    shrink_factor: int = 4,
    num_fitting_levels: int = 4,
    num_iterations: int = 50,
) -> sitk.Image:
    """Perform N4 bias correction on MRI

    Note:
    - if no mask is provided it will be generated using Otsu-thresholding
    """
    if not isinstance(mask, sitk.Image):
        mask = sitk.OtsuThreshold(input, 0, 1, 200)

    input = sitk.Cast(input, sitk.sitkFloat32)
    image = sitk.Shrink(input, [shrink_factor] * input.GetDimension())
    mask = sitk.Shrink(mask, [shrink_factor] * input.GetDimension())

    corrector = sitk.N4BiasFieldCorrectionImageFilter()
    corrector.SetMaximumNumberOfIterations([num_iterations] * num_fitting_levels)

    corrector.Execute(image, mask)
    log_bias_field = corrector.GetLogBiasFieldAsImage(input)
    corrected_image_full_resolution = input / sitk.Exp(log_bias_field)
    return corrected_image_full_resolution

But since ITK Python does not have the convenience function GetLogBiasFieldAsImage and does not wrap itk::BSplineControlPointImageFilter, it seems I cannot implement bias correction as in ANTs using ITK Python.

If (there are no workarounds) I might do a pull request, but I wonder if I should implement the GetLogBiasFieldAsImage like helper or wrap BSplineControlPointImageFilter to implement something like this.

Then again, I might just install ITK (for MONAI ITKReader) and SimpleITK (for image pre-processing) side-by-side, since dealing with image/filter types is much easier and auto-completion works (for simpleitk)

Looking at itkN4BiasFieldCorrectionImageFilter.h there is a function ReconstructBiasField. It looks like I should basically be able to do:

    image_small = itk.shrink_image_filter(image, shrink_factors=[shrink_factor] * image.GetImageDimension())
    mask_small = itk.shrink_image_filter(mask, shrink_factors=[shrink_factor] * image.GetImageDimension())

    corrector = itk.N4BiasFieldCorrectionImageFilter.New(image_small, mask_small)
    corrector.SetNumberOfFittingLevels(num_fitting_levels)
    corrector.SetMaximumNumberOfIterations([num_iterations] * num_fitting_levels)
    corrector.Update()

    full_res_corrector = itk.N4BiasFieldCorrectionImageFilter.New(image, mask)
    log_bias_field = full_res_corrector.ReconstructBiasField(corrector.GetLogBiasFieldControlPointLattice())

    print("image", itk.size(image), itk.origin(image), itk.spacing(image))
    print("log_bias_field", itk.size(log_bias_field), itk.origin(log_bias_field), itk.spacing(log_bias_field))

    corrected_image_np = itk.array_view_from_image(image) / np.exp(itk.array_view_from_image(log_bias_field))
    corrected_image_full_resolution = itk.image_from_array(corrected_image_np)

For some reason the reconstructed bias field has the wrong shape though:

image itkSize3 ([256, 256, 150])
log_bias_field itkSize3 ([253, 253, 145])

But ReconstructBiasField is setting the size from the input image RequestedRegion after the data is generated (?):

    biasField->SetRegions(inputImage->GetRequestedRegion());

So, my workaround for this bug (?) is to to first ensure the requested region is the same as the largest possible (and buffered) region:

    image.SetRequestedRegion(image.GetBufferedRegion())
    
    full_res_corrector = itk.N4BiasFieldCorrectionImageFilter.New(image, mask)
    log_bias_field = full_res_corrector.ReconstructBiasField(corrector.GetLogBiasFieldControlPointLattice())

This sounds better to me. But if engineering time is not an issue for you, doing both sounds good to me.