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

Code before:
            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
    pad_size = []
    for size in filter_size:
        pad_size.extend([size // 2, size // 2])

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

    # 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

Review comment (suggestion):
```suggestion
            f"spatial_dims must be either 2 or 3, got spatial_dims={spatial_dims} for mask tensor with shape of {mask_t.shape}."
```

Code after:
            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}.")

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

    # Pad the input tensor to handle border pixels
    # Calculate padding size
    pad_size = [size // 2 for size in filter_size for _ in range(2)]

    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