Antialiasing for labels

[Spenhouet]

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

[wyli]

Duplicate of https://github.com/Project-MONAI/MONAI/issues/178? Any idea how this affects the model training?

[wyli]

There is also a multiscale loss wrapper https://github.com/Project-MONAI/MONAI/blob/dev/monai/losses/multi_scale.py

[Spenhouet]

@wyli I did see #178 but I did not recognize it as duplicate in that my suggestion does not do any downsampling, right?

Also, could you elaborate on the multiscale loss wrapper. I'm unfamiliar with that.

Any idea how this affects the model training?

Not yet. I guess we will find out 🙂

We will need to make sure to tune the parameters so that there is no significant information loss but smoother shapes actually seem more natural and learning smooth shapes compared to jagged ones, does seem like a good idea. I'm positive that this will have a positive effect on the training.

I'm also planning on extending the implementation in that the labels can be provided as parameter. That would allow to set different parameters for different label classes. This will allow to better fine tune the parameters for every label class.

[Spenhouet]

Here the version with the option to define the labels the edge smoothing is applied to and an optional pre-dilation. The pre-dilation seems to help to preserve smaller shapes (just an experimental addition).

def isin_tensor(ar1: Tensor, ar2: Iterable[Any]):
    """
    This is a placeholder for torch.isin which will be released with torch 1.10.0.
    https://pytorch.org/docs/1.10.0/generated/torch.isin.html?highlight=isin
    https://github.com/pytorch/pytorch/issues/3025
    """
    return (ar1[..., None] == torch.tensor(list(ar2))).any(-1)


def dilation3d(binary_image: torch.Tensor) -> torch.Tensor:
    kernel = torch.ones(size=[1, 1, 3, 3, 3], dtype=torch.float32)
    return torch.clamp(conv3d(binary_image, kernel, padding=(1, 1, 1)), 0, 1)


class Antialiasing(MapTransform):

    def __init__(
        self,
        keys: KeysCollection,
        applied_labels: Optional[Union[Iterable[int], int]] = None,
        sigma: Union[Sequence[float], float] = 1.0,
        approx: str = "erf",
        threshold: float = 0.5,
        pre_dilation: bool = False,
        allow_missing_keys: bool = False,
    ) -> None:
        super().__init__(keys, allow_missing_keys)
        self.applied_labels = ensure_tuple(applied_labels) if applied_labels else None
        self.sigma = sigma
        self.approx = approx
        self.threshold = threshold
        self.pre_dilation = pre_dilation

    def __call__(self, data: Mapping[Hashable, NdarrayTensor]) -> Dict[Hashable, NdarrayTensor]:
        d = dict(data)
        for key in self.key_iterator(d):
            img: torch.Tensor
            img, *_ = convert_data_type(d[key], torch.Tensor)  # type: ignore

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

            # Get labels if not provided. Exclude background label.
            applied_labels = set(self.applied_labels or torch.unique(img))
            background_label = 0
            applied_labels.discard(background_label)

            # Ignore labels not in `applied_labels`
            new_img = img * (isin_tensor(img, applied_labels) == 0)

            for label in applied_labels:
                label_mask = (img == label).to(torch.float32).unsqueeze(0)
                if self.pre_dilation:
                    label_mask = dilation3d(label_mask)
                blurred = gaussian_filter(label_mask).squeeze(0)
                new_img[blurred > self.threshold] = label
            d[key] = new_img  # type: ignore
        return d

[Jingnan-Jia]

@Spenhouet Hi Spenhouet, I am very interested in this smooth method. I also implementated a similar method [#3333 ].

However I found that in my task, I can not get good results using Gaussian filger (I tried different sigma values from 1 to 10). So I applied Mean filter and then got very good results. My threshold is still 0.4.

I do not know how you choose the correct sigma value in your method? Have you tried Mean filger? Because if you use Mean filter, you do not need to set sigma value. You just need to set the filter size which should be the zoom factor I think.

[Spenhouet]

@Jingnan-Jia I just finetuned the parameters for our data (trial and error). What is the issue of "not good results"?

Could you share your transform code?

[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

[wyli]

that looks like mean filtering instead of median, ~perhaps you can use https://github.com/Project-MONAI/MONAI/blob/dev/monai/transforms/post/array.py~ (I reread, and realised no good transform currently in monai for it)...

and a few transforms should be unified to have a generic filtering API:

• https://github.com/Project-MONAI/MONAI/blob/ef68f16dde33172b2eeae5397ac75aa717808875/monai/transforms/post/array.py#L517
• https://github.com/Project-MONAI/MONAI/blob/ef68f16dde33172b2eeae5397ac75aa717808875/monai/transforms/post/array.py#L812
• https://github.com/Project-MONAI/MONAI/blob/ef68f16dde33172b2eeae5397ac75aa717808875/monai/transforms/intensity/array.py#L1139

cc @Nic-Ma

[kbressem]

I can open a PR and contribute the mean filtering transform if you want, also supporting multiple kernels (but I might need some help with creating the kernels).

[wyli]

thanks @kbressem, please go ahead, I haven't looked into the details, but in my opinion some MONAI transform to achieve similar functions like https://github.com/kornia/kornia/blob/master/kornia/filters/filter.py would be very useful, using the utilities: https://github.com/Project-MONAI/MONAI/blob/607123b6723f2c7d22783d9ac4b30a2c6391ac48/monai/networks/layers/simplelayers.py#L38-L39 cc @ericspod @rijobro @Nic-Ma @mingxin-zheng

[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.

[dzenanz]

I created #5264, as a related but separate feature request.

[kbressem]

Hi @dzenanz I'm currently working on a more general FilterKernelTransform, which provides mean filtering, laplacian filter, elliptical filter and sobel filter. Maybe we can work together on this?

[dzenanz]

Maybe we can collaborate. I see you have a fork, but I was unable to find any progress there. But overhead of combining our efforts might be high, as I am starting from existing working code which I only need to plug into MONAI form with accompanying tests, docs etc.