Antialiasing for labels

Is your feature request related to a problem? Please describe. Low resolution labels are often too jagged.

Describe the solution you’d like Originally I hoped that the Spacingd transform would just have some antialiasing. But a separate transform for this might be a better solution anyway.

I implemented a simple transform using a gaussian filter combined with a threshold on every label.

class Antialiasingd(MapTransform):

    def __init__(
        self,
        keys: KeysCollection,
        sigma: Union[Sequence[float], float] = 1.0,
        approx: str = "erf",
        threshold: float = 0.5,
        allow_missing_keys: bool = False,
    ) -> None:
        super().__init__(keys, allow_missing_keys)
        self.sigma = sigma
        self.approx = approx
        self.threshold = threshold

    def __call__(self, data: Mapping[Hashable, NdarrayTensor]) -> Dict[Hashable, NdarrayTensor]:
        d = dict(data)
        for key in self.key_iterator(d):
            img = d[key]

            gaussian_filter = GaussianFilter(img.ndim - 1, self.sigma, approx=self.approx)

            labels = torch.unique(img)[1:]
            new_img = torch.zeros_like(img)
            for label in labels:
                label_mask = (img == label).to(torch.float)
                blurred = gaussian_filter(label_mask.unsqueeze(0)).squeeze(0)
                new_img[blurred > self.threshold] = label
            d[key] = new_img
        return d

What do you think?

Here some results

Parameters	Result
Original
Sigma=1.0, Threshold=0.5
Sigma=4.0, Threshold=0.5
Sigma=4.0, Threshold=0.4

[jlvahldiek]

Another option could be median smoothing. We achieved very good and fast results with the below transform.

from typing import Optional, Union, Iterable, Hashable, Mapping, Dict

from monai.data import MetaTensor
from monai.config import KeysCollection
from monai.config.type_definitions import NdarrayTensor
from monai.transforms.transform import MapTransform
from monai.utils import ensure_tuple
from monai.utils.type_conversion import convert_data_type
import torch
import torch.nn.functional as F


class MedianSmoothd(MapTransform):

    """Apply channel-wise median smoothing to a channel first image"""
    def __init__(
        self, 
        keys: KeysCollection,
        applied_labels: Optional[Union[Iterable[int], int]] = None,
        kernel_size: int = 5,
        allow_missing_keys: bool = False,
    ) -> None:
        super().__init__(keys, allow_missing_keys)
        self.applied_labels = ensure_tuple(applied_labels) if applied_labels else None  # ToDo: filter out non-applied labels
        assert kernel_size % 2 != 0, "Kernel size should be uneven."
        self.kernel_size = kernel_size

    def __call__(self, data: Mapping[Hashable, NdarrayTensor]) -> Dict[Hashable, NdarrayTensor]:
        d = dict(data)
        for key in self.key_iterator(d):
            image, prev_type, device = convert_data_type(d[key], MetaTensor)
            smooth_items = [self._smooth(im, self.kernel_size) for im in torch.unbind(image, 0)]
            smooth_image = torch.stack(smooth_items, 0)
            smooth_image.meta = image.meta  # keep metadata
            image, *_ = convert_data_type(smooth_image, prev_type, device)
            d[key] = image
        return d
    
    def _smooth(self, image: MetaTensor, kernel_size: int) -> MetaTensor:
        """Smoothes a 2D or 3D array using median smoothing."""
        assert image.ndim in [2,3], f"Only 2D or 3D images are supported, got {image.ndim}D"
        kernel_dim = [1, 1] + [kernel_size]*image.ndim
        kernel = torch.ones(*kernel_dim, device=image.device) / kernel_size**image.ndim
        conv = getattr(F, f"conv{image.ndim}d")
        # add channel and batch dim, so image.shape is now 1, 1, [W, H, [D]]
        image = image.unsqueeze(0).unsqueeze(0) 
        smooth_image = conv(image, kernel, padding=kernel_size // 2)
        return smooth_image.squeeze() # remove batch dim

cc @kbressem

[dzenanz]

Is there interest in adding a median filter based on @ebrahimebrahim’s medianblur_with_gpu? If so, I would be interested in contributing it with a PR.