Local thickness wrapper

qim3d/processing/local_thickness.py

 """Wrapper for the local thickness function from the localthickness package including visualization functions."""
 import localthickness as lt
 import numpy as np
-import matplotlib.pyplot as plt
+from typing import Optional
 from skimage.filters import threshold_otsu
 from qim3d.io.logger import log
-from typing import Optional
+from qim3d.viz import local_thickness as viz_local_thickness
 
-def local_thickness(image: np.ndarray, scale: float = 1, mask: Optional[np.ndarray] = None) -> np.ndarray:
+def local_thickness(image: np.ndarray, scale: float = 1, mask: Optional[np.ndarray] = None, visualize = False, **viz_kwargs) -> np.ndarray:
     """Wrapper for the local thickness function from the localthickness package* (https://github.com/vedranaa/local-thickness)
 
     Args:
         image (np.ndarray): 2D or 3D NumPy array representing the image/volume. 
             If binary, it will be passed directly to the local thickness function. 
             If grayscale, it will be binarized using Otsu's method.
-        scale (float, optional): Downscaling factor, e.g. 0.5 for halving each dim of the image. Default is 1.
+        scale (float, optional): Downscaling factor, e.g. 0.5 for halving each dim of the image. 
+            Default is 1.
         mask (np.ndarray, optional): binary mask of the same size of the image defining parts of the
             image to be included in the computation of the local thickness. Default is None.
+        visualize (bool, optional): Whether to visualize the local thickness. Default is False.
+        **viz_kwargs: Additional keyword arguments for the visualization function. Only used if visualize=True.
     
     Returns:
         local_thickness (np.ndarray): 2D or 3D NumPy array representing the local thickness of the input image/volume.
@@ -27,9 +30,16 @@ def local_thickness(image: np.ndarray, scale: float = 1, mask: Optional[np.ndarr
 
     # Check if input is binary
     if np.unique(image).size > 2:
-        # If not, binarize it using Otsu's method
+        # If not, binarize it using Otsu's method, log the threshold and compute the local thickness
         threshold = threshold_otsu(image=image)
-        image = image > threshold
         log.warning("Input image is not binary. It will be binarized using Otsu's method with threshold: {}".format(threshold))
+        local_thickness = lt.local_thickness(image > threshold, scale=scale, mask=mask)
+    else:
+        # If it is binary, compute the local thickness directly
+        local_thickness = lt.local_thickness(image, scale=scale, mask=mask)
+
+    # Visualize the local thickness if requested
+    if visualize:
+        display(viz_local_thickness(image, local_thickness, **viz_kwargs))
 
-    return lt.local_thickness(image, scale=scale, mask=mask)
\ No newline at end of file
+    return local_thickness
\ No newline at end of file