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

Code before:
        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 or ndarray.

    Example:

        .. code-block:: python

            # define a naive mask
            mask = torch.zeros(3,2,3,3,3)
            mask[:,:,1,1,1] = 1.0
            filter_size = 3
            erode_result = erode(mask,filter_size) # expect torch.zeros(3,2,3,3,3)
            dilate_result = dilate(mask,filter_size) # expect torch.ones(3,2,3,3,3)
    """
    mask_t, *_ = convert_data_type(mask, torch.Tensor)
    res_mask_t = dilate_t(mask_t, filter_size=filter_size, pad_value=pad_value)
    res_mask: NdarrayOrTensor
    res_mask, *_ = convert_to_dst_type(src=res_mask_t, dst=mask)
    return res_mask


def get_morphological_filter_result_t(mask_t: Tensor, filter_size: int | Sequence[int], pad_value: float) -> Tensor:
    """
    Get morphological filter result for 2D/3D 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: morphological filter size, has to be odd numbers.
        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.

    Return:
        Morphological filter result mask, [N,C,M,N] or [N,C,M,N,P] torch tensor
    """
    spatial_dims = len(mask_t.shape) - 2
    if spatial_dims not in [2, 3]:
        raise ValueError(
            f"spatial_dims must be either 2 or 3, yet got spatial_dims={spatial_dims} for mask tensor with shape of {mask_t.shape}."
        )

    # Define the structuring element
    filter_size = ensure_tuple_rep(filter_size, spatial_dims)
    if any(size % 2 == 0 for size in filter_size):
        raise ValueError(f"All dimensions in filter_size must be odd numbers, yet got {filter_size}.")

    filter_shape = [mask_t.shape[1], mask_t.shape[1]] + list(filter_size)
    structuring_element = torch.ones(filter_shape).to(mask_t.device)

    # Pad the input tensor to handle border pixels
    # Calculate padding size

Review comment (suggestion):
```suggestion
    Apply a morphological filter to a 2D/3D binary mask tensor.
```

Code after:
        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, same shape and data type as input.

    Example:

        .. code-block:: python

            # define a naive mask
            mask = torch.zeros(3,2,3,3,3)
            mask[:,:,1,1,1] = 1.0
            filter_size = 3
            erode_result = erode(mask,filter_size) # expect torch.zeros(3,2,3,3,3)
            dilate_result = dilate(mask,filter_size) # expect torch.ones(3,2,3,3,3)
    """
    mask_t, *_ = convert_data_type(mask, torch.Tensor)
    res_mask_t = dilate_t(mask_t, filter_size=filter_size, pad_value=pad_value)
    res_mask: NdarrayOrTensor
    res_mask, *_ = convert_to_dst_type(src=res_mask_t, dst=mask)
    return res_mask


def get_morphological_filter_result_t(mask_t: Tensor, filter_size: int | Sequence[int], pad_value: float) -> Tensor:
    """
    Apply a morphological filter to a 2D/3D binary mask tensor.

    Args:
        mask_t: input 2D/3D binary mask, [N,C,M,N] or [N,C,M,N,P] torch tensor.
        filter_size: morphological filter size, has to be odd numbers.
        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.

    Return:
        Tensor: Morphological filter result mask, same shape as input.
    """
    spatial_dims = len(mask_t.shape) - 2
    if spatial_dims not in [2, 3]:
        raise ValueError(
            f"spatial_dims must be either 2 or 3, "
            f"got spatial_dims={spatial_dims} for mask tensor with shape of {mask_t.shape}."
        )

    # Define the structuring element
    filter_size = ensure_tuple_rep(filter_size, spatial_dims)
    if any(size % 2 == 0 for size in filter_size):
        raise ValueError(f"All dimensions in filter_size must be odd numbers, got {filter_size}.")