Repository: Project-MONAI/MONAI
File: monai/apps/generation/maisi/utils/morphological_ops.py (Python)

Code before:
    # Apply filter operation
    if spatial_dims == 2:
        output = F.conv2d(input_padded, structuring_element, padding=0) / torch.sum(structuring_element[0, ...])
    if spatial_dims == 3:
        output = F.conv3d(input_padded, structuring_element, padding=0) / torch.sum(structuring_element[0, ...])

    return output


def erode_t(mask_t: Tensor, filter_size: int | Sequence[int] = 3, pad_value: float = 1.0) -> Tensor:
    """
    Erode 2D/3D binary mask with data type as torch tensor.

    Args:
        mask_t: input 2D/3D binary mask, [N,C,M,N] or [N,C,M,N,P] torch tensor.
        filter_size: erosion filter size, has to be odd numbers, default to be 3.
        pad_value: the filled value for padding. We need to pad the input before filtering to keep the output with the same size as input. Usually use default value and not changed.

    Return:
        eroded mask, [N,C,M,N] or [N,C,M,N,P] torch tensor
    """

    output = get_morphological_filter_result_t(mask_t, filter_size, pad_value)

    # Set output values based on the minimum value within the structuring element
    output = torch.where(output == 1.0, 1.0, 0.0)

    return output


def dilate_t(mask_t: Tensor, filter_size: int | Sequence[int] = 3, pad_value: float = 0.0) -> Tensor:
    """
    Dilate 2D/3D binary mask with data type as torch tensor.

    Args:
        mask_t: input 2D/3D binary mask, [N,C,M,N] or [N,C,M,N,P] torch tensor.
        filter_size: dilation filter size, has to be odd numbers, default to be 3.
        pad_value: the filled value for padding. We need to pad the input before filtering to keep the output with the same size as input. Usually use default value and not changed.

    Return:
        dilated mask, [N,C,M,N] or [N,C,M,N,P] torch tensor
    """
    output = get_morphological_filter_result_t(mask_t, filter_size, pad_value)

    # Set output values based on the minimum value within the structuring element
    output = torch.where(output > 0, 1.0, 0.0)

    return output


Review comment (suggestion):
Maybe it's safer to add `eps`  when we compare with float numbers:

```
output = torch.where(torch.abs(output - 1.0) < 1e-7, 1.0, 0.0)  # or another small value as deemed appropriate
```

Code after:
    input_padded = F.pad(mask_t.float(), pad_size, mode="constant", value=pad_value)

    # Apply filter operation
    conv_fn = F.conv2d if spatial_dims == 2 else F.conv3d
    output = conv_fn(input_padded, structuring_element, padding=0) / torch.sum(structuring_element[0, ...])

    return output


def erode_t(mask_t: Tensor, filter_size: int | Sequence[int] = 3, pad_value: float = 1.0) -> Tensor:
    """
    Erode 2D/3D binary mask with data type as torch tensor.

    Args:
        mask_t: input 2D/3D binary mask, [N,C,M,N] or [N,C,M,N,P] torch tensor.
        filter_size: erosion filter size, has to be odd numbers, default to be 3.
        pad_value: the filled value for padding. We need to pad the input before filtering
                   to keep the output with the same size as input. Usually use default value
                   and not changed.

    Return:
        Tensor: eroded mask, same shape as input.
    """

    output = get_morphological_filter_result_t(mask_t, filter_size, pad_value)

    # Set output values based on the minimum value within the structuring element
    output = torch.where(torch.abs(output - 1.0) < 1e-7, 1.0, 0.0)

    return output


def dilate_t(mask_t: Tensor, filter_size: int | Sequence[int] = 3, pad_value: float = 0.0) -> Tensor:
    """
    Dilate 2D/3D binary mask with data type as torch tensor.

    Args:
        mask_t: input 2D/3D binary mask, [N,C,M,N] or [N,C,M,N,P] torch tensor.
        filter_size: dilation filter size, has to be odd numbers, default to be 3.
        pad_value: the filled value for padding. We need to pad the input before filtering
                   to keep the output with the same size as input. Usually use default value
                   and not changed.

    Return:
        Tensor: dilated mask, same shape as input.
    """
    output = get_morphological_filter_result_t(mask_t, filter_size, pad_value)

    # Set output values based on the minimum value within the structuring element
    output = torch.where(output > 0, 1.0, 0.0)