Add histogram equalization procedure

SimpleITK/utilities/__init__.py

@@ -22,6 +22,7 @@
 from .fft import fft_based_translation_initialization
 from .overlay_bounding_boxes import overlay_bounding_boxes
 from .resize import resize
+from .histogram_equalization import histogram_equalization
 
 try:
     from ._version import version as __version__
@@ -35,6 +36,7 @@
     "make_isotropic",
     "fft_based_translation_initialization",
     "overlay_bounding_boxes",
+    "histogram_equalization",
     "resize",
     "__version__",
 ]

SimpleITK/utilities/histogram_equalization.py

@@ -0,0 +1,60 @@
+# ========================================================================
+#
+#  Copyright NumFOCUS
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0.txt
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+# ========================================================================
+
+import SimpleITK as sitk
+import numpy as np
+from SimpleITK.extra import _get_numpy_dtype
+
+
+def histogram_equalization(
+    image: sitk.Image,
+    number_of_histogram_levels: int = 256,
+    number_of_match_points: int = 32,
+    threshold_at_mean_intensity: bool = False,
+) -> sitk.Image:
+    """
+    Histogram equalization is a method in image processing of contrast adjustment using the image's histogram.
+
+    The output intensity range will be [0, number_of_histogram_levels - 1].
+
+    This implementation uses the itk::HistogramMatchingImageFilter to perform the histogram matching with a flat
+    histogram. The flat histogram is generated from a reference image created with the specified number of levels
+    and then the filter performs histogram matching with the specified number of match points.
+
+    :param image:
+    :param number_of_histogram_levels:
+    :param number_of_match_points:
+    :param threshold_at_mean_intensity:
+    :return:
+    """
+
+    ramp_values = np.arange(number_of_histogram_levels, dtype=_get_numpy_dtype(image))
+
+    # reshape ramp_values to match the number of dimensions in the image
+    ramp_values = ramp_values.reshape((1,) * (image.GetDimension() - 1) + (-1,))
+
+    ramp_image = sitk.GetImageFromArray(ramp_values)
+
+    histigram_matching_filter = sitk.HistogramMatchingImageFilter()
+    histigram_matching_filter.SetNumberOfHistogramLevels(number_of_histogram_levels)
+    histigram_matching_filter.SetNumberOfMatchPoints(number_of_match_points)
+    histigram_matching_filter.SetThresholdAtMeanIntensity(threshold_at_mean_intensity)
+
+    output_image = histigram_matching_filter.Execute(image, ramp_image)
+
+    return output_image

test/test_utilities.py

@@ -2,6 +2,7 @@
 import SimpleITK as sitk
 import SimpleITK.utilities as sitkutils
 from numpy.testing import assert_allclose
+import numpy as np
 
 
 def test_Logger():
@@ -322,3 +323,80 @@ def test_resize_anti_aliasing():
     assert resized_image[0, 0] == resized_image[2, 2]
     assert resized_image[1, 0] == 0.0
     assert resized_image[1, 2] == 0.0
+
+
+def test_histogram_equalization():
+    # Test with 0 image corner case
+    image = sitk.Image([32, 32], sitk.sitkFloat32)
+    image.SetSpacing([2.1, 3.2])
+    image.SetOrigin([0.5, 0.5])
+    image[:] = 1.0
+
+    equalized_image = sitkutils.histogram_equalization(
+        image,
+        number_of_histogram_levels=256,
+        number_of_match_points=32,
+        threshold_at_mean_intensity=False,
+    )
+    assert equalized_image.GetSize() == image.GetSize()
+    assert equalized_image.GetSpacing() == image.GetSpacing()
+    assert equalized_image.GetOrigin() == image.GetOrigin()
+    assert equalized_image.GetPixelID() == image.GetPixelID()
+    assert equalized_image.GetSize() == (32, 32)
+
+    # check that all pixels are equal
+    pixel_value = equalized_image[0, 0]
+    assert all([pv == pixel_value for pv in equalized_image])
+
+    # test with ramp image
+    image = sitk.Image([256, 1], sitk.sitkFloat32)
+    for i in range(image.GetSize()[0]):
+        image[i, 0] = i
+
+    equalized_image = sitkutils.histogram_equalization(
+        image,
+        number_of_histogram_levels=256,
+        number_of_match_points=32,
+        threshold_at_mean_intensity=False,
+    )
+    assert equalized_image.GetSize() == (256, 1)
+    assert equalized_image[0, 0] == 0.0
+    assert equalized_image[127, 0] == 127.0
+    assert equalized_image[255, 0] == 255.0
+
+    for i in range(image.GetSize()[0]):
+        image[i, 0] = -32.0 + i / 8.0
+
+    equalized_image = sitkutils.histogram_equalization(
+        image,
+        number_of_histogram_levels=256,
+        number_of_match_points=32,
+        threshold_at_mean_intensity=False,
+    )
+    assert equalized_image.GetSize() == (256, 1)
+    assert equalized_image[0, 0] == 0.0
+    assert equalized_image[127, 0] == 127.0
+    assert equalized_image[255, 0] == 255.0
+
+    for i in range(image.GetSize()[0]):
+        image[i, 0] = -32.0 + i**2 / 8.0
+
+    equalized_image = sitkutils.histogram_equalization(
+        image,
+        number_of_histogram_levels=256,
+        number_of_match_points=32,
+        threshold_at_mean_intensity=False,
+    )
+    assert np.isclose(equalized_image[0, 0], 0.0, atol=1e-6)
+    assert np.isclose(equalized_image[127, 0], 127.0, atol=1e-6)
+    assert np.isclose(equalized_image[255, 0], 255.0, atol=1e-6)
+
+    equalized_image = sitkutils.histogram_equalization(
+        image,
+        number_of_histogram_levels=128,
+        number_of_match_points=32,
+        threshold_at_mean_intensity=False,
+    )
+    assert np.isclose(equalized_image[0, 0], 0.0, atol=1e-6)
+    assert np.isclose(equalized_image[127, 0], 63.0, atol=0.5)
+    assert np.isclose(equalized_image[255, 0], 127.0, atol=1e-6)