feat(transforms): add CLAHE intensity transform (Fixes #9049)

docs/requirements.txt

@@ -5,6 +5,7 @@ itk>=5.2
 nibabel
 parameterized
 scikit-image>=0.19.0
+SimpleITK>=2.1.1.1
 scipy>=1.12.0
 tensorboard
 commonmark==0.9.1

docs/source/installation.md

@@ -77,10 +77,10 @@ pip install monai
 MONAI supports the extras syntax such as `pip install 'monai[nibabel]'`. The options are
 
 ```text
-clearml, cucim, cupy, einops, fire, gdown, h5py, huggingface_hub, hyena, ignite, imagecodecs, itk, jsonschema, lmdb, lpips, matplotlib, metrics_reloaded, mlflow, nibabel, nni, nvimgcodec, onnx, openslide, optuna, pandas, pillow, polygraphy, psutil, pyamg, pybind11, pydicom, pynrrd, pynvml, pyyaml, requests, segment_anything, scipy, skimage, tensorboard, tensorboardX, tifffile, torchio, torchvision, tqdm, transformers, zarr
+clearml, cucim, cupy, einops, fire, gdown, h5py, huggingface_hub, hyena, ignite, imagecodecs, itk, jsonschema, lmdb, lpips, matplotlib, metrics_reloaded, mlflow, nibabel, nni, nvimgcodec, onnx, openslide, optuna, pandas, pillow, polygraphy, psutil, pyamg, pybind11, pydicom, pynrrd, pynvml, pyyaml, requests, segment_anything, scipy, simpleitk, skimage, tensorboard, tensorboardX, tifffile, torchio, torchvision, tqdm, transformers, zarr
 ```
 
-which correspond to the packages: `clearml`, `cucim` (`cucim-cu12` or `cucim-cu13`), `cupy-cuda13x`, `einops`, `fire`, `gdown`, `h5py`, `huggingface_hub`, `nvsubquadratic`, `omegaconf`, `pytorch-ignite`, `imagecodecs`, `itk`, `jsonschema`, `lmdb`, `lpips`, `matplotlib`, `MetricsReloaded`, `mlflow`, `nibabel`, `nni`, `filelock`, `nvidia-nvimgcodec-cu13`, `onnx`, `onnxruntime`, `onnx_graphsurgeon`, `onnxscript`, `openslide-python`, `openslide-bin`, `optuna`, `pandas`, `pillow`, `polygraphy`, `psutil`, `pyamg`, `pybind11`, `pydicom`, `pynrrd`, `nvidia-ml-py`, `pyyaml`, `requests`, `segment_anything`, `scipy`, `scikit-image`, `tensorboard`, `tensorboardX`, `tifffile`, `torchio`, `torchvision`, `tqdm`, `transformers`, `zarr`.
+which correspond to the packages: `clearml`, `cucim` (`cucim-cu12` or `cucim-cu13`), `cupy-cuda13x`, `einops`, `fire`, `gdown`, `h5py`, `huggingface_hub`, `nvsubquadratic`, `omegaconf`, `pytorch-ignite`, `imagecodecs`, `itk`, `jsonschema`, `lmdb`, `lpips`, `matplotlib`, `MetricsReloaded`, `mlflow`, `nibabel`, `nni`, `filelock`, `nvidia-nvimgcodec-cu13`, `onnx`, `onnxruntime`, `onnx_graphsurgeon`, `onnxscript`, `openslide-python`, `openslide-bin`, `optuna`, `pandas`, `pillow`, `polygraphy`, `psutil`, `pyamg`, `pybind11`, `pydicom`, `pynrrd`, `nvidia-ml-py`, `pyyaml`, `requests`, `segment_anything`, `scipy`, `SimpleITK`, `scikit-image`, `tensorboard`, `tensorboardX`, `tifffile`, `torchio`, `torchvision`, `tqdm`, `transformers`, `zarr`.
 
 Almost all of these can be installed together with the `all` option. For development on MONAI, this should be accompanied by `testing` which will install the testing static checking packages. Cupy and `nvimgcodec` are omitted from `all` since the choice between
 Cuda 12 and 13 versions of the libraries can't be resolved when installing and must be manually installed.

docs/source/transforms.rst

@@ -528,6 +528,12 @@ Intensity
     :members:
     :special-members: __call__
 
+`CLAHE`
+"""""""
+.. autoclass:: CLAHE
+    :members:
+    :special-members: __call__
+
 
 `ForegroundMask`
 """"""""""""""""
@@ -1660,6 +1666,13 @@ Intensity (Dict)
     :members:
     :special-members: __call__
 
+`CLAHEd`
+""""""""
+.. autoclass:: CLAHEd
+    :members:
+    :special-members: __call__
+
+
 `ForegroundMaskd`
 """""""""""""""""
 .. image:: https://raw.githubusercontent.com/Project-MONAI/DocImages/main/transforms/ForegroundMaskd.png

monai/transforms/__init__.py

@@ -91,6 +91,7 @@
 )
 from .croppad.functional import crop_func, crop_or_pad_nd, pad_func, pad_nd
 from .intensity.array import (
+    CLAHE,
     AdjustContrast,
     ClipIntensityPercentiles,
     ComputeHoVerMaps,
@@ -133,9 +134,12 @@
     UltrasoundConfidenceMapTransform,
 )
 from .intensity.dictionary import (
+    CLAHED,
     AdjustContrastd,
     AdjustContrastD,
     AdjustContrastDict,
+    CLAHEd,
+    CLAHEDict,
     ClipIntensityPercentilesd,
     ClipIntensityPercentilesD,
     ClipIntensityPercentilesDict,
@@ -695,6 +699,7 @@
     allow_missing_keys_mode,
     attach_hook,
     check_non_lazy_pending_ops,
+    clahe,
     compute_divisible_spatial_size,
     convert_applied_interp_mode,
     convert_pad_mode,

monai/transforms/intensity/array.py

@@ -32,7 +32,7 @@
 from monai.networks.layers import GaussianFilter, HilbertTransform, MedianFilter, SavitzkyGolayFilter
 from monai.transforms.inverse import InvertibleTransform
 from monai.transforms.transform import RandomizableTransform, Transform
-from monai.transforms.utils import Fourier, equalize_hist, is_positive, rescale_array, soft_clip
+from monai.transforms.utils import Fourier, clahe, equalize_hist, is_positive, rescale_array, soft_clip
 from monai.transforms.utils_pytorch_numpy_unification import clip, nonzero, percentile, ravel, where
 from monai.utils.enums import TraceKeys, TransformBackends
 from monai.utils.misc import ensure_tuple, ensure_tuple_rep, ensure_tuple_size, fall_back_tuple
@@ -77,6 +77,7 @@
     "RandCoarseDropout",
     "RandCoarseShuffle",
     "HistogramNormalize",
+    "CLAHE",
     "IntensityRemap",
     "RandIntensityRemap",
     "ForegroundMask",
@@ -2694,6 +2695,112 @@ def __call__(self, img: NdarrayOrTensor, mask: NdarrayOrTensor | None = None) ->
         return out
 
 
+class CLAHE(Transform):
+    """
+    Apply Contrast Limited Adaptive Histogram Equalization (CLAHE) to an image.
+
+    Backed by `SimpleITK.AdaptiveHistogramEqualization
+    <https://simpleitk.org/doxygen/latest/html/classitk_1_1simple_1_1AdaptiveHistogramEqualizationImageFilter.html>`_,
+    ITK's power-law variant, described in `Adaptive Image Contrast Enhancement using
+    Generalizations of Histogram Equalization <https://ieeexplore.ieee.org/document/841534/>`_.
+    Expects a channel-first 2D or 3D image, equalizes each channel independently, and keeps
+    the intensity range of the input.
+
+    Args:
+        radius: radius of the contextual window in voxels, used for every spatial axis if a
+            single number. default to `5`, ITK's own default, which suits small images: scale
+            it to the image, roughly a sixteenth of the shorter side.
+        alpha: `0` gives classical adaptive histogram equalization, `1` an unsharp mask,
+            in `[0, 1]`. default to `0.3`.
+        beta: `0` gives an unsharp mask, `1` together with `alpha=1` returns the input
+            unchanged, in `[0, 1]`. default to `0.3`.
+        nonzero: whether to leave zero voxels out of the result, keeping their original value.
+        mask: if provided, must be ndarray of bools or 0s and 1s, and same shape as the image
+            or its spatial shape, in which case it is shared across channels. voxels at which
+            `mask==False` keep their original value. can also provide the mask along with img
+            at runtime.
+        num_threads: number of threads SimpleITK may use. `None` leaves its default, which
+            is every core and will oversubscribe inside dataloader workers.
+        dtype: data type of the output, if None, same as input image. default to `float32`.
+
+    """
+
+    backend = [TransformBackends.NUMPY]
+
+    def __init__(
+        self,
+        radius: Sequence[int] | int = 5,
+        alpha: float = 0.3,
+        beta: float = 0.3,
+        nonzero: bool = False,
+        mask: NdarrayOrTensor | None = None,
+        num_threads: int | None = None,
+        dtype: DtypeLike = np.float32,
+    ) -> None:
+        if not 0.0 <= alpha <= 1.0:
+            raise ValueError(f"alpha must be in [0, 1], got {alpha}.")
+        if not 0.0 <= beta <= 1.0:
+            raise ValueError(f"beta must be in [0, 1], got {beta}.")
+        if num_threads is not None and not (isinstance(num_threads, (int, np.integer)) and num_threads >= 1):
+            raise ValueError(f"num_threads must be a whole number of at least 1, got {num_threads}.")
+        self.radius = radius
+        self.alpha = alpha
+        self.beta = beta
+        self.nonzero = nonzero
+        self.mask = mask
+        self.num_threads = num_threads
+        self.dtype = dtype
+
+    @staticmethod
+    def _resolve_channel_mask(mask: NdarrayOrTensor | None, img: np.ndarray) -> np.ndarray | None:
+        """
+        Turn `mask` into a boolean array shaped like the channel-first `img`.
+
+        A mask matching only the spatial shape is broadcast across channels, so one brain mask
+        can serve every modality of a multi-channel volume.
+        """
+        if mask is None:
+            return None
+        converted, *_ = convert_data_type(mask, np.ndarray)
+        mask_np = np.asarray(converted, dtype=bool)
+        if mask_np.shape == img.shape[1:]:
+            return np.broadcast_to(mask_np, img.shape)
+        if mask_np.shape != img.shape:
+            raise ValueError(
+                f"mask shape {mask_np.shape} must match either the image shape {img.shape}"
+                f" or its spatial shape {img.shape[1:]}."
+            )
+        return mask_np
+
+    def __call__(self, img: NdarrayOrTensor, mask: NdarrayOrTensor | None = None) -> NdarrayOrTensor:
+        img = convert_to_tensor(img, track_meta=get_track_meta())
+        img_np, *_ = convert_data_type(img, np.ndarray)
+        if img_np.ndim < 3:
+            raise ValueError(
+                "CLAHE expects a channel-first image with at least two spatial"
+                f" dimensions, got shape {img_np.shape}. Use `EnsureChannelFirst` to add the channel dimension."
+            )
+        mask_np = self._resolve_channel_mask(mask if mask is not None else self.mask, img_np)
+
+        ret = np.stack(
+            [
+                clahe(
+                    img=c,
+                    radius=self.radius,
+                    alpha=self.alpha,
+                    beta=self.beta,
+                    nonzero=self.nonzero,
+                    mask=None if mask_np is None else mask_np[i],
+                    num_threads=self.num_threads,
+                )
+                for i, c in enumerate(img_np)
+            ]
+        )
+        out, *_ = convert_to_dst_type(src=ret, dst=img, dtype=self.dtype or img.dtype)
+
+        return out
+
+
 class IntensityRemap(RandomizableTransform):
     """
     Transform for intensity remapping of images. The intensity at each

monai/transforms/intensity/dictionary.py

@@ -26,6 +26,7 @@
 from monai.config.type_definitions import NdarrayOrTensor
 from monai.data.meta_obj import get_track_meta
 from monai.transforms.intensity.array import (
+    CLAHE,
     AdjustContrast,
     ClipIntensityPercentiles,
     ComputeHoVerMaps,
@@ -99,6 +100,7 @@
     "RandCoarseDropoutd",
     "RandCoarseShuffled",
     "HistogramNormalized",
+    "CLAHEd",
     "ForegroundMaskd",
     "ComputeHoVerMapsd",
     "RandGaussianNoiseD",
@@ -164,6 +166,8 @@
     "RandCoarseShuffleDict",
     "HistogramNormalizeD",
     "HistogramNormalizeDict",
+    "CLAHED",
+    "CLAHEDict",
     "RandKSpaceSpikeNoiseD",
     "RandKSpaceSpikeNoiseDict",
     "ForegroundMaskD",
@@ -1906,6 +1910,60 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, N
         return d
 
 
+class CLAHEd(MapTransform):
+    """
+    Dictionary-based wrapper of :py:class:`monai.transforms.CLAHE`.
+
+    Args:
+        keys: keys of the corresponding items to be transformed.
+            See also: :py:class:`monai.transforms.compose.MapTransform`
+        radius: radius of the contextual window in voxels, used for every spatial axis if a
+            single number. default to `5`, ITK's own default, which suits small images: scale
+            it to the image, roughly a sixteenth of the shorter side.
+        alpha: `0` gives classical adaptive histogram equalization, `1` an unsharp mask,
+            in `[0, 1]`. default to `0.3`.
+        beta: `0` gives an unsharp mask, `1` together with `alpha=1` returns the input
+            unchanged, in `[0, 1]`. default to `0.3`.
+        nonzero: whether to leave zero voxels out of the result, keeping their original value.
+        mask: if provided, must be ndarray of bools or 0s and 1s, and same shape as the image
+            or its spatial shape. voxels at which `mask==False` keep their original value.
+        mask_key: if mask is None, will try to get the mask with `mask_key`.
+        num_threads: number of threads SimpleITK may use. `None` leaves its default, which
+            is every core and will oversubscribe inside dataloader workers.
+        dtype: output data type, if None, same as input image. defaults to float32.
+        allow_missing_keys: do not raise exception if key is missing.
+
+    """
+
+    backend = CLAHE.backend
+
+    def __init__(
+        self,
+        keys: KeysCollection,
+        radius: Sequence[int] | int = 5,
+        alpha: float = 0.3,
+        beta: float = 0.3,
+        nonzero: bool = False,
+        mask: NdarrayOrTensor | None = None,
+        mask_key: str | None = None,
+        num_threads: int | None = None,
+        dtype: DtypeLike = np.float32,
+        allow_missing_keys: bool = False,
+    ) -> None:
+        super().__init__(keys, allow_missing_keys)
+        self.transform = CLAHE(
+            radius=radius, alpha=alpha, beta=beta, nonzero=nonzero, mask=mask, num_threads=num_threads, dtype=dtype
+        )
+        self.mask_key = mask_key if mask is None else None
+
+    def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]:
+        d = dict(data)
+        for key in self.key_iterator(d):
+            d[key] = self.transform(d[key], d[self.mask_key]) if self.mask_key is not None else self.transform(d[key])
+
+        return d
+
+
 class ForegroundMaskd(MapTransform):
     """
     Creates a binary mask that defines the foreground based on thresholds in RGB or HSV color space.
@@ -2016,6 +2074,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, N
 RandKSpaceSpikeNoiseD = RandKSpaceSpikeNoiseDict = RandKSpaceSpikeNoised
 RandCoarseDropoutD = RandCoarseDropoutDict = RandCoarseDropoutd
 HistogramNormalizeD = HistogramNormalizeDict = HistogramNormalized
+CLAHED = CLAHEDict = CLAHEd
 RandCoarseShuffleD = RandCoarseShuffleDict = RandCoarseShuffled
 ForegroundMaskD = ForegroundMaskDict = ForegroundMaskd
 ComputeHoVerMapsD = ComputeHoVerMapsDict = ComputeHoVerMapsd