--- a +++ b/utils/preprocess.py @@ -0,0 +1,110 @@ +import os +import SimpleITK as sitk +import numpy as np +from skimage import exposure + +def anisotropic_diffusion_denoise_3d(input_volume, conductance_parameter, time_step, number_of_iterations): + ''' + Denoise the input volume using anisotropic diffusion. + + Args: + input_volume ('SimpleITK Image'): Input volume. + conductance_parameter ('float'): Conductance parameter. + time_step ('float'): Time step. + number_of_iterations ('int'): Number of iterations. + + Returns: + output_volume ('SimpleITK Image'): Denoised volume. + + ''' + # Convert the input volume to float + input_volume_float = sitk.Cast(input_volume, sitk.sitkFloat32) + + anisotropic_diffusion = sitk.GradientAnisotropicDiffusionImageFilter() + anisotropic_diffusion.SetConductanceParameter(conductance_parameter) + anisotropic_diffusion.SetTimeStep(time_step) + anisotropic_diffusion.SetNumberOfIterations(number_of_iterations) + + output_volume_float = anisotropic_diffusion.Execute(input_volume_float) + output_volume = sitk.Cast(output_volume_float, sitk.sitkInt16) + + return output_volume + +def bilateral_filter_3d(input_volume, domain_sigma, range_sigma): + ''' + Apply bilateral filter to the input volume. + + Args: + input_volume ('SimpleITK Image'): Input volume. + domain_sigma ('float'): Domain sigma. + range_sigma ('float'): Range sigma. + + Returns: + output_volume ('SimpleITK Image'): Filtered volume. + ''' + filtered_slices = [] + + for z in range(input_volume.GetDepth()): + # Get 2D slice + input_slice = sitk.Extract(input_volume, (input_volume.GetWidth(), input_volume.GetHeight(), 1), (0, 0, z)) + + # Apply bilateral filter to the 2D slice + bilateral_filter = sitk.BilateralImageFilter() + bilateral_filter.SetDomainSigma(domain_sigma) + bilateral_filter.SetRangeSigma(range_sigma) + output_slice = bilateral_filter.Execute(input_slice) + + # Append the filtered slice to the list + filtered_slices.append(sitk.GetArrayFromImage(output_slice)[0]) # [0] to get the slice from the 3D volume + + # Convert the list of slices to a NumPy array + filtered_slices_array = np.array(filtered_slices) + output_volume = sitk.GetImageFromArray(filtered_slices_array) + + # Copy the information from the input volume to the output volume + output_volume.CopyInformation(input_volume) + + return output_volume + + +def clahe_3d(input_volume, clip_limit=0.1): + """ + Apply Contrast Limited Adaptive Histogram Equalization (CLAHE) to a 3D image. + + Args: + input_volume (SimpleITK Image): The input image. + + Returns: + SimpleITK Image: The output image after applying CLAHE. + """ + # Get the minimum and maximum intensity values of the input image + min_max_filter = sitk.MinimumMaximumImageFilter() + min_max_filter.Execute(input_volume) + + original_min = min_max_filter.GetMinimum() + original_max = min_max_filter.GetMaximum() + + # Get the dimensions of the 3D volume + size_x, size_y, size_z = input_volume.GetSize() + + # Determine kernel sizes in each dim relative to image shape + kernel_size = (size_x // 5, size_y // 5, size_z // 2) + + # get the image from the input volume + input_volume_image = sitk.GetArrayFromImage(input_volume) + + # Normalize intensity values to the range [0, 1] + # input_volume_normalized = sitk.RescaleIntensity(input_volume_image, outputMinimum=0, outputMaximum=1) + input_volume_normalized = exposure.rescale_intensity(input_volume_image, in_range='image', out_range=(0, 1)) + + # Apply CLAHE to the normalized image + input_volume_enhanced = exposure.equalize_adapthist(input_volume_normalized, kernel_size=kernel_size, clip_limit=clip_limit) + + # Rescale intensity values to the original range + input_volume_rescaled = exposure.rescale_intensity(input_volume_enhanced, in_range='image', out_range=(original_min, original_max)) + + # Convert the NumPy array to a SimpleITK image + output_result = sitk.GetImageFromArray(input_volume_rescaled) + output_result.CopyInformation(input_volume) + + return output_result \ No newline at end of file