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 #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:

• MONAI/monai/transforms/post/array.py

  Line 517 in ef68f16

  class LabelToContour(Transform):

• MONAI/monai/transforms/post/array.py

  Line 812 in ef68f16

  class SobelGradients(Transform):

• MONAI/monai/transforms/intensity/array.py

  Line 1139 in ef68f16

  class GaussianSmooth(Transform):

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:

MONAI/monai/networks/layers/simplelayers.py

Lines 38 to 39 in 607123b

	"apply_filter",
	"separable_filtering",

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.