Repository: mlcommons/GaNDLF

Issue:
Update `FillHoles` implementation
**Is your feature request related to a problem? Please describe.**
The current implementation can be significantly simplified.

**Describe the solution you'd like**
Add pure morphological operations based on `torch` operations, so that they can take place on gpus, if needed [[ref](https://github.com/DIVA-DIA/Generating-Synthetic-Handwritten-Historical-Documents/blob/e6a798dc2b374f338804222747c56cb44869af5b/HTR_ctc/utils/auxilary_functions.py#L10)].

**Describe alternatives you've considered**
N.A.

**Additional context**
Would be great if this could be moved under `GANDLF.post_process` to start a single cohesive location to put post-processing functionality. 



Discussion:

Resolving patch:
diff --git a/GANDLF/FillHoles.py b/GANDLF/FillHoles.py
deleted file mode 100644
--- a/GANDLF/FillHoles.py
+++ /dev/null
@@ -1,38 +0,0 @@
-# ------------------------------------------------------------------------------
-# Name: Fill Holes
-# Date: February 22, 2021
-# Version: 1.0
-# Author: Jose L. Agraz, PhD
-#
-# Description: This program fills in holes in a tensor using a given structure.
-#              First, the input is eroded and dilated (closed). Then, holes
-#              filled in. Keep in mind structure shape and size is key to a
-#              successful application. Structure examples in the scipy library
-#              are square, rectangle, diamond, disk, cube, octahedron, ball,
-#              octagon, and star.
-#              See scipy structure documentation "Generate structuring elements"
-#
-# Input:  Tensor, Structure (optional, default cube, 4x4 pixels)
-# Output: Tensor
-#
-# ------------------------------------------------------------------------------
-import torch
-from skimage.morphology import cube
-from scipy.ndimage.morphology import binary_fill_holes
-from scipy.ndimage.morphology import binary_closing
-
-
-def FillHoles(InputTensor, FillingStructure=cube(4)):
-    # Convert tensor to numpy array
-    InputNumpyArray = InputTensor.numpy()
-    # The closing of an input image by erosion and dilation of the image by a structuring element.
-    # Closing fills holes smaller than the structuring element
-    ClosedArray = binary_closing(InputNumpyArray, structure=FillingStructure).astype(
-        int
-    )
-    # Fill the holes in binary objects
-    OutputArray = binary_fill_holes(ClosedArray, structure=FillingStructure).astype(int)
-    # convert numpy array to tensor
-    TensorOutput = torch.from_numpy(OutputArray)
-
-    return TensorOutput
diff --git a/GANDLF/post_process/__init__.py b/GANDLF/post_process/__init__.py
new file mode 100644
--- /dev/null
+++ b/GANDLF/post_process/__init__.py
@@ -0,0 +1 @@
+from .morphology import torch_morphological, fill_holes
diff --git a/GANDLF/post_process/morphology.py b/GANDLF/post_process/morphology.py
new file mode 100644
--- /dev/null
+++ b/GANDLF/post_process/morphology.py
@@ -0,0 +1,66 @@
+import torch
+import torch.nn.functional as F
+from scipy.ndimage.morphology import binary_fill_holes, binary_closing
+
+
+def torch_morphological(input_image, kernel_size=1, mode="dilation"):
+    """
+    This function enables morphological operations using torch. Adapted from https://github.com/DIVA-DIA/Generating-Synthetic-Handwritten-Historical-Documents/blob/e6a798dc2b374f338804222747c56cb44869af5b/HTR_ctc/utils/auxilary_functions.py#L10.
+
+    Args:
+        input_image (torch.Tensor): The input image.
+        kernel_size (list): The size of the window to take a max over.
+        mode (str): The type of morphological operation to perform.
+
+    Returns:
+        torch.Tensor: The output image after morphological operations.
+    """
+
+    if len(input_image.shape) == 4:
+        max_pool = F.max_pool2d
+    elif len(input_image.shape) == 5:
+        max_pool = F.max_pool3d
+    else:
+        raise ValueError("Input image has invalid shape for morphological operations.")
+
+    if mode == "dilation":
+        output_image = max_pool(
+            input_image, kernel_size=kernel_size, stride=1, padding=kernel_size // 2
+        )
+    elif mode == "erosion":
+        output_image = -max_pool(
+            -input_image, kernel_size=kernel_size, stride=1, padding=kernel_size // 2
+        )
+    elif mode == "closing":
+        output_image = max_pool(
+            input_image, kernel_size=kernel_size, stride=1, padding=kernel_size // 2
+        )
+        output_image = -max_pool(
+            -output_image, kernel_size=kernel_size, stride=1, padding=kernel_size // 2
+        )
+    elif mode == "opening":
+        output_image = -max_pool(
+            -input_image, kernel_size=kernel_size, stride=1, padding=kernel_size // 2
+        )
+        output_image = max_pool(
+            output_image, kernel_size=kernel_size, stride=1, padding=kernel_size // 2
+        )
+
+    return output_image
+
+
+def fill_holes(input_image):
+    """
+    This function fills holes in masks.
+
+    Args:
+        input_image (torch.Tensor): The input image.
+
+    Returns:
+        torch.Tensor: The output image after morphological operations.
+    """
+    input_image_array_closed = binary_closing(input_image.numpy())
+    # Fill the holes in binary objects
+    output_array = binary_fill_holes(input_image_array_closed)
+
+    return torch.from_numpy(output_array)


Test patch:
diff --git a/testing/test_full.py b/testing/test_full.py
--- a/testing/test_full.py
+++ b/testing/test_full.py
@@ -15,6 +15,7 @@
 from GANDLF.schedulers import global_schedulers_dict
 from GANDLF.optimizers import global_optimizer_dict
 from GANDLF.models import global_models_dict
+from GANDLF.post_process import torch_morphological, fill_holes
 
 device = "cpu"
 ## global defines
@@ -930,6 +931,8 @@ def test_preprocess_functions():
     non_zero_normalizer = global_preprocessing_dict["normalize_nonZero"]
     input_transformed = non_zero_normalizer(input_tensor)
 
+    input_transformed = fill_holes(input_tensor)
+
     input_image = sitk.GetImageFromArray(input_tensor[0].numpy())
     img_resized = resize_image(
         input_image,
@@ -952,6 +955,13 @@ def test_preprocess_functions():
     input_transformed = cropper(input_tensor)
     assert input_transformed.shape == (1, 128, 128, 128), "Resampling should work"
 
+    # test pure morphological operations
+    input_tensor_3d = torch.rand(1, 1, 256, 256, 256)
+    input_tensor_2d = torch.rand(1, 3, 256, 256)
+    for mode in ["dilation", "erosion", "opening", "closing"]:
+        input_transformed_3d = torch_morphological(input_tensor_3d, mode=mode)
+        input_transformed_2d = torch_morphological(input_tensor_2d, mode=mode)
+
     print("passed")