swcgeom 0.15.0__py3-none-any.whl → 0.17.0__py3-none-any.whl

swcgeom/images/contrast.py

@@ -0,0 +1,107 @@
+"""The contrast of an image.
+
+Notes
+-----
+This is expremental code, and the API is subject to change.
+"""
+
+from typing import Optional, overload
+
+import numpy as np
+import numpy.typing as npt
+
+__all__ = ["contrast_std", "contrast_michelson", "contrast_rms", "contrast_weber"]
+
+Array3D = npt.NDArray[np.float32]
+
+
+@overload
+def contrast_std(image: Array3D) -> float:
+    """Get the std contrast of an image stack.
+
+    Parameters
+    ----------
+    imgs : ndarray
+
+    Returns
+    -------
+    contrast : float
+    """
+    ...
+
+
+@overload
+def contrast_std(image: Array3D, contrast: float) -> Array3D:
+    """Adjust the contrast of an image stack.
+
+    Parameters
+    ----------
+    imgs : ndarray
+    constrast : float
+        The contrast adjustment factor. 1.0 leaves the image unchanged.
+
+    Returns
+    -------
+    imgs : ndarray
+        The adjusted image.
+    """
+    ...
+
+
+def contrast_std(image: Array3D, contrast: Optional[float] = None):
+    if contrast is None:
+        return np.std(image).item()
+    else:
+        return np.clip(contrast * image, 0, 1)
+
+
+def contrast_michelson(image: Array3D) -> float:
+    """Get the Michelson contrast of an image stack.
+
+    Parameters
+    ----------
+    imgs : ndarray
+
+    Returns
+    -------
+    contrast : float
+    """
+
+    vmax = np.max(image)
+    vmin = np.min(image)
+    return ((vmax - vmin) / (vmax + vmin)).item()
+
+
+def contrast_rms(imgs: npt.NDArray[np.float32]) -> float:
+    """Get the RMS contrast of an image stack.
+
+    Parameters
+    ----------
+    imgs : ndarray
+
+    Returns
+    -------
+    contrast : float
+    """
+
+    return np.sqrt(np.mean(imgs**2)).item()
+
+
+def contrast_weber(imgs: Array3D, mask: npt.NDArray[np.bool_]) -> float:
+    """Get the Weber contrast of an image stack.
+
+    Parameters
+    ----------
+    imgs : ndarray
+    mask : ndarray of bool
+        The mask to segment the foreground and background. 1 for
+        foreground, 0 for background.
+
+    Returns
+    -------
+    contrast : float
+    """
+
+    l_foreground = np.mean(imgs, where=mask)
+    l_background = np.mean(imgs, where=np.logical_not(mask))
+    return ((l_foreground - l_background) / l_background).item()

swcgeom/images/folder.py

@@ -1,9 +1,10 @@
 """Image stack folder."""
 
+import math
 import os
 import re
 import warnings
-from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from typing import (
     Callable,
     Generic,
@@ -18,21 +19,18 @@ from typing import (
 
 import numpy as np
 import numpy.typing as npt
+from tqdm import tqdm
 from typing_extensions import Self
 
 from swcgeom.images.io import ScalarType, read_imgs
 from swcgeom.transforms import Identity, Transform
 
-__all__ = [
-    "ImageStackFolder",
-    "LabeledImageStackFolder",
-    "PathImageStackFolder",
-]
+__all__ = ["ImageStackFolder", "LabeledImageStackFolder", "PathImageStackFolder"]
 
 T = TypeVar("T")
 
 
-class ImageStackFolderBase(Generic[ScalarType, T], ABC):
+class ImageStackFolderBase(Generic[ScalarType, T]):
     """Image stack folder base."""
 
     files: List[str]
@@ -51,10 +49,6 @@ class ImageStackFolderBase(Generic[ScalarType, T], ABC):
         self.dtype = dtype or np.float32
         self.transform = transform or Identity()  # type: ignore
 
-    @abstractmethod
-    def __getitem__(self, key: str, /) -> T:
-        raise NotImplementedError()
-
     def __len__(self) -> int:
         return len(self.files)
 
@@ -78,6 +72,12 @@ class ImageStackFolderBase(Generic[ScalarType, T], ABC):
 
     @staticmethod
     def read_imgs(fname: str) -> npt.NDArray[np.float32]:
+        """Read images.
+
+        .. deprecated:: 0.16.0
+            Use :meth:`~swcgeom.images.io.read_imgs(fname).get_full()` instead.
+        """
+
         warnings.warn(
             "`ImageStackFolderBase.read_imgs` serves as a "
             "straightforward wrapper for `~swcgeom.images.io.read_imgs(fname).get_full()`. "
@@ -89,12 +89,67 @@ class ImageStackFolderBase(Generic[ScalarType, T], ABC):
         return read_imgs(fname).get_full()
 
 
+@dataclass(frozen=True)
+class Statistics:
+    count: int = 0
+    minimum: float = math.nan
+    maximum: float = math.nan
+    mean: float = 0
+    variance: float = 0
+
+
 class ImageStackFolder(ImageStackFolderBase[ScalarType, T]):
     """Image stack folder."""
 
     def __getitem__(self, idx: int, /) -> T:
         return self._get(self.files[idx])
 
+    def stat(self, *, transform: bool = False, verbose: bool = False) -> Statistics:
+        """Statistics of folder.
+
+        Parameters
+        ----------
+        transform : bool, default to False
+            Apply transform to the images. If True, you need to make
+            sure the transformed data is a ndarray.
+        verbose : bool, optional
+
+        Notes
+        -----
+        We are asserting that the images are of the same shape.
+        """
+
+        vmin, vmax = math.inf, -math.inf
+        n, mean, M2 = 0, None, None
+
+        for idx in tqdm(range(len(self))) if verbose else range(len(self)):
+            imgs = self[idx] if transform else self._read(self.files[idx])
+
+            vmin = min(vmin, np.min(imgs))  # type: ignore
+            vmax = max(vmax, np.max(imgs))  # type: ignore
+            # Welford algorithm to calculate mean and variance
+            if mean is None:
+                mean = np.zeros_like(imgs)
+                M2 = np.zeros_like(imgs)
+
+            n += 1
+            delta = imgs - mean  # type: ignore
+            mean += delta / n
+            delta2 = imgs - mean
+            M2 += delta * delta2
+
+        if mean is None or M2 is None:  # n = 0
+            raise ValueError("empty folder")
+
+        variance = M2 / (n - 1) if n > 1 else np.zeros_like(mean)
+        return Statistics(
+            count=len(self),
+            maximum=vmax,
+            minimum=vmin,
+            mean=np.mean(mean).item(),
+            variance=np.mean(variance).item(),
+        )
+
     @classmethod
     def from_dir(cls, root: str, *, pattern: Optional[str] = None, **kwargs) -> Self:
         """
@@ -106,6 +161,7 @@ class ImageStackFolder(ImageStackFolderBase[ScalarType, T]):
         **kwargs
             Pass to `cls.__init__`
         """
+
         return cls(cls.scan(root, pattern=pattern), **kwargs)
 
 
@@ -118,8 +174,8 @@ class LabeledImageStackFolder(ImageStackFolderBase[ScalarType, T]):
         super().__init__(files, **kwargs)
         self.labels = list(labels)
 
-    def __getitem__(self, idx: int) -> Tuple[npt.NDArray[np.float32], int]:
-        return self.read_imgs(self.files[idx]), self.labels[idx]
+    def __getitem__(self, idx: int) -> Tuple[T, int]:
+        return self._get(self.files[idx]), self.labels[idx]
 
     @classmethod
     def from_dir(
@@ -140,7 +196,7 @@ class LabeledImageStackFolder(ImageStackFolderBase[ScalarType, T]):
         return cls(files, labels, **kwargs)
 
 
-class PathImageStackFolder(ImageStackFolder[ScalarType, T]):
+class PathImageStackFolder(ImageStackFolderBase[ScalarType, T]):
     """Image stack folder with relpath."""
 
     root: str
@@ -164,6 +220,7 @@ class PathImageStackFolder(ImageStackFolder[ScalarType, T]):
         **kwargs
             Pass to `cls.__init__`
         """
+
         return cls(cls.scan(root, pattern=pattern), root=root, **kwargs)
 
 

swcgeom/images/io.py

@@ -107,26 +107,42 @@ def read_imgs(fname: str, *, dtype: None =..., **kwargs) -> ImageStack[np.float3
 # fmt:on
 
 
-def read_imgs(fname: str, *, dtype=None, **kwargs):  # type: ignore
-    """Read image stack."""
-
-    kwargs["dtype"] = dtype or np.float32
-
-    ext = os.path.splitext(fname)[-1]
-    if ext in [".tif", ".tiff"]:
-        return TiffImageStack(fname, **kwargs)
-    if ext in [".nrrd"]:
-        return NrrdImageStack(fname, **kwargs)
-    if ext in [".v3dpbd"]:
-        return V3dpbdImageStack(fname, **kwargs)
-    if ext in [".v3draw"]:
-        return V3drawImageStack(fname, **kwargs)
-    if ext in [".npy", ".npz"]:
-        return NDArrayImageStack(np.load(fname), **kwargs)
+def read_imgs(fname: str, **kwargs):  # type: ignore
+    """Read image stack.
+
+    Parameters
+    ----------
+    fname : str
+        The path of image stack.
+    dtype : np.dtype, default to `np.float32`
+        Casting data to specified dtype. If integer and float
+        conversions occur, they will be scaled (assuming floats are
+        between 0 and 1).
+    **kwargs : Dict[str, Any]
+        Forwarding to the corresponding reader.
+    """
+
+    kwargs.setdefault("dtype", np.float32)
+    if not os.path.exists(fname):
+        raise ValueError(f"image stack not exists: {fname}")
+
+    # match file extension
+    match os.path.splitext(fname)[-1]:
+        case ".tif" | ".tiff":
+            return TiffImageStack(fname, **kwargs)
+        case ".nrrd":
+            return NrrdImageStack(fname, **kwargs)
+        case ".v3dpbd":
+            return V3dpbdImageStack(fname, **kwargs)
+        case ".v3draw":
+            return V3drawImageStack(fname, **kwargs)
+        case ".npy":
+            return NDArrayImageStack(np.load(fname), **kwargs)
+
+    # try to read as terafly
     if TeraflyImageStack.is_root(fname):
         return TeraflyImageStack(fname, **kwargs)
-    if not os.path.exists(fname):
-        raise ValueError("image stack not exists")
+
     raise ValueError("unsupported image stack")
 
 
@@ -135,7 +151,6 @@ def save_tiff(
     fname: str,
     *,
     dtype: Optional[np.unsignedinteger | np.floating] = None,
-    swap_xy: Optional[bool] = None,
     compression: str | Literal[False] = "zlib",
     **kwargs,
 ) -> None:
@@ -164,17 +179,6 @@ def save_tiff(
         data = np.expand_dims(data, -1)  # (_, _, _)  -> (_, _, _, C), C === 1
 
     axes = "ZXYC"
-    if swap_xy is not None:
-        warnings.warn(
-            "flag `swap_xy` is easy to implement in user space and "
-            "is more flexiable. Since this flag is rarely used, we "
-            "decided to remove it in the next version",
-            DeprecationWarning,
-        )
-        if swap_xy is True:
-            axes = "ZYXC"
-            data = data.swapaxes(0, 1)  # (X, Y, _, _) -> (Y, X, _, _)
-
     assert data.ndim == 4, "should be an array of shape (X, Y, Z, C)"
     assert data.shape[-1] in [1, 3], "support 'miniblack' or 'rgb'"
 
@@ -209,12 +213,7 @@ class NDArrayImageStack(ImageStack[ScalarType]):
     """NDArray image stack."""
 
     def __init__(
-        self,
-        imgs: npt.NDArray[Any],
-        swap_xy: Optional[bool] = None,
-        filp_xy: Optional[bool] = None,
-        *,
-        dtype: ScalarType,
+        self, imgs: npt.NDArray[Any], *, dtype: Optional[ScalarType] = None
     ) -> None:
         super().__init__()
 
@@ -222,34 +221,22 @@ class NDArrayImageStack(ImageStack[ScalarType]):
             imgs = np.expand_dims(imgs, -1)
         assert imgs.ndim == 4, "Should be shape of (X, Y, Z, C)"
 
-        if swap_xy is not None:
-            warnings.warn(
-                "flag `swap_xy` now is unnecessary, tifffile will "
-                "automatically adjust dimensions according to "
-                "`tags.axes`, so this flag will be removed in the next "
-                " version",
-                DeprecationWarning,
-            )
-            if swap_xy is True:
-                imgs = imgs.swapaxes(0, 1)  # (Y, X, _, _) -> (X, Y, _, _)
-
-        if filp_xy is not None:
-            warnings.warn(
-                "flag `filp_xy` is easy to implement in user space and "
-                "is more flexiable. Since this flag is rarely used, we "
-                "decided to remove it in the next version",
-                DeprecationWarning,
-            )
-            if filp_xy is True:
-                imgs = np.flip(imgs, (0, 1))  # (X, Y, Z, C)
+        if dtype is not None:
+            dtype_raw = imgs.dtype
+            if np.issubdtype(dtype, np.floating) and np.issubdtype(
+                dtype_raw, np.unsignedinteger
+            ):
+                sclar_factor = 1.0 / UINT_MAX[dtype_raw]
+                imgs = sclar_factor * imgs.astype(dtype)
+            elif np.issubdtype(dtype, np.unsignedinteger) and np.issubdtype(
+                dtype_raw, np.floating
+            ):
+                sclar_factor = UINT_MAX[dtype]  # type: ignore
+                imgs *= (sclar_factor * imgs).astype(dtype)
+            else:
+                imgs = imgs.astype(dtype)
 
-        dtype_raw = imgs.dtype
-        self.imgs = imgs.astype(dtype)
-        if np.issubdtype(dtype, np.floating) and np.issubdtype(
-            dtype_raw, np.unsignedinteger
-        ):  # TODO: add a option to disable this
-            sclar_factor = 1.0 / UINT_MAX[imgs.dtype]
-            self.imgs *= sclar_factor
+        self.imgs = imgs
 
     def __getitem__(self, key):
         return self.imgs.__getitem__(key)
@@ -265,15 +252,7 @@ class NDArrayImageStack(ImageStack[ScalarType]):
 class TiffImageStack(NDArrayImageStack[ScalarType]):
     """Tiff image stack."""
 
-    def __init__(
-        self,
-        fname: str,
-        swap_xy: Optional[bool] = None,
-        filp_xy: Optional[bool] = None,
-        *,
-        dtype: ScalarType,
-        **kwargs,
-    ) -> None:
+    def __init__(self, fname: str, *, dtype: ScalarType, **kwargs) -> None:
         with tifffile.TiffFile(fname, **kwargs) as f:
             s = f.series[0]
             imgs, axes = s.asarray(), s.axes
@@ -285,23 +264,15 @@ class TiffImageStack(NDArrayImageStack[ScalarType]):
 
         orders = [AXES_ORDER[c] for c in axes]
         imgs = imgs.transpose(np.argsort(orders))
-        super().__init__(imgs, swap_xy=swap_xy, filp_xy=filp_xy, dtype=dtype)
+        super().__init__(imgs, dtype=dtype)
 
 
 class NrrdImageStack(NDArrayImageStack[ScalarType]):
     """Nrrd image stack."""
 
-    def __init__(
-        self,
-        fname: str,
-        swap_xy: Optional[bool] = None,
-        filp_xy: Optional[bool] = None,
-        *,
-        dtype: ScalarType,
-        **kwargs,
-    ) -> None:
+    def __init__(self, fname: str, *, dtype: ScalarType, **kwargs) -> None:
         imgs, header = nrrd.read(fname, **kwargs)
-        super().__init__(imgs, swap_xy=swap_xy, filp_xy=filp_xy, dtype=dtype)
+        super().__init__(imgs, dtype=dtype)
         self.header = header
 
 
@@ -333,12 +304,16 @@ class V3dpbdImageStack(V3dImageStack[ScalarType]):
 class TeraflyImageStack(ImageStack[ScalarType]):
     """TeraFly image stack.
 
+    TeraFly is a terabytes of multidimensional volumetric images file
+    format as described in [1]_.
+
     References
     ----------
-    [1] Bria, Alessandro, Giulio Iannello, Leonardo Onofri, and
-    Hanchuan Peng. “TeraFly: Real-Time Three-Dimensional Visualization
-    and Annotation of Terabytes of Multidimensional Volumetric Images.”
-    Nature Methods 13, no. 3 (March 2016): 192-94. https://doi.org/10.1038/nmeth.3767.
+    .. [1] Bria, Alessandro, Giulio Iannello, Leonardo Onofri, and
+       Hanchuan Peng. “TeraFly: Real-Time Three-Dimensional
+       Visualization and Annotation of Terabytes of Multidimensional
+       Volumetric Images.” Nature Methods 13,
+       no. 3 (March 2016): 192-94. https://doi.org/10.1038/nmeth.3767.
 
     Notes
     -----
@@ -380,7 +355,12 @@ class TeraflyImageStack(ImageStack[ScalarType]):
 
         @lru_cache(maxsize=lru_maxsize)
         def read_patch(path: str) -> npt.NDArray[ScalarType]:
-            return read_imgs(path, dtype=dtype).get_full()
+            match os.path.splitext(path)[-1]:
+                case "raw":
+                    # Treat it as a v3draw file
+                    return V3drawImageStack(path, dtype=dtype).get_full()
+                case _:
+                    return read_imgs(path, dtype=dtype).get_full()
 
         self._listdir, self._read_patch = listdir, read_patch
 
@@ -633,6 +613,12 @@ class GrayImageStack:
 
 
 def read_images(*args, **kwargs) -> GrayImageStack:
+    """Read images.
+
+    .. deprecated:: 0.16.0
+        Use :meth:`read_imgs` instead.
+    """
+
     warnings.warn(
         "`read_images` has been replaced by `read_imgs` because it"
         "provide rgb support, and this will be removed in next version",

swcgeom/transforms/__init__.py

@@ -3,9 +3,11 @@
 from swcgeom.transforms.base import *
 from swcgeom.transforms.branch import *
 from swcgeom.transforms.geometry import *
+from swcgeom.transforms.image_preprocess import *
 from swcgeom.transforms.image_stack import *
 from swcgeom.transforms.images import *
 from swcgeom.transforms.mst import *
+from swcgeom.transforms.neurolucida_asc import *
 from swcgeom.transforms.path import *
 from swcgeom.transforms.population import *
 from swcgeom.transforms.tree import *

swcgeom/transforms/image_preprocess.py

@@ -0,0 +1,100 @@
+"""Image stack pre-processing."""
+
+import numpy as np
+import numpy.typing as npt
+from scipy.fftpack import fftn, fftshift, ifftn
+from scipy.ndimage import gaussian_filter, minimum_filter
+
+from swcgeom.transforms.base import Transform
+
+__all__ = ["SGuoImPreProcess"]
+
+
+class SGuoImPreProcess(Transform[npt.NDArray[np.uint8], npt.NDArray[np.uint8]]):
+    """Single-Neuron Image Enhancement.
+
+    Implementation of the image enhancement method described in the paper:
+
+    Shuxia Guo, Xuan Zhao, Shengdian Jiang, Liya Ding, Hanchuan Peng,
+    Image enhancement to leverage the 3D morphological reconstruction
+    of single-cell neurons, Bioinformatics, Volume 38, Issue 2,
+    January 2022, Pages 503–512, https://doi.org/10.1093/bioinformatics/btab638
+    """
+
+    def __call__(self, x: npt.NDArray[np.uint8]) -> npt.NDArray[np.uint8]:
+        # TODO: support np.float32
+        assert x.dtype == np.uint8, "Image must be in uint8 format"
+        x = self.sigmoid_adjustment(x)
+        x = self.subtract_min_along_z(x)
+        x = self.bilateral_filter_3d(x)
+        x = self.high_pass_fft(x)
+        return x
+
+    @staticmethod
+    def sigmoid_adjustment(
+        image: npt.NDArray[np.uint8], sigma: float = 3, percentile: float = 25
+    ) -> npt.NDArray[np.uint8]:
+        image_normalized = image / 255.0
+        u = np.percentile(image_normalized, percentile)
+        adjusted = 1 / (1 + np.exp(-sigma * (image_normalized - u)))
+        adjusted_rescaled = (adjusted * 255).astype(np.uint8)
+        return adjusted_rescaled
+
+    @staticmethod
+    def subtract_min_along_z(image: npt.NDArray[np.uint8]) -> npt.NDArray[np.uint8]:
+        min_along_z = minimum_filter(
+            image,
+            size=(1, 1, image.shape[2], 1),
+            mode="constant",
+            cval=np.max(image).item(),
+        )
+        subtracted = image - min_along_z
+        return subtracted
+
+    @staticmethod
+    def bilateral_filter_3d(
+        image: npt.NDArray[np.uint8], spatial_sigma=(1, 1, 0.33), range_sigma=35
+    ) -> npt.NDArray[np.uint8]:
+        # initialize the output image
+        filtered_image = np.zeros_like(image)
+
+        spatial_gaussian = gaussian_filter(image, spatial_sigma)
+
+        # traverse each pixel to perform bilateral filtering
+        # TODO: optimization is needed
+        for z in range(image.shape[2]):
+            for y in range(image.shape[1]):
+                for x in range(image.shape[0]):
+                    value = image[x, y, z]
+                    range_weight = np.exp(
+                        -((image - value) ** 2) / (2 * range_sigma**2)
+                    )
+                    weights = spatial_gaussian * range_weight
+                    filtered_image[x, y, z] = np.sum(image * weights) / np.sum(weights)
+
+        return filtered_image
+
+    @staticmethod
+    def high_pass_fft(image: npt.NDArray[np.uint8]) -> npt.NDArray[np.uint8]:
+        # fft
+        fft_image = fftn(image)
+        fft_shifted = fftshift(fft_image)
+
+        # create a high-pass filter
+        h, w, d, _ = image.shape
+        x, y, z = np.ogrid[:h, :w, :d]
+        center = (h / 2, w / 2, d / 2)
+        distance = np.sqrt(
+            (x - center[0]) ** 2 + (y - center[1]) ** 2 + (z - center[2]) ** 2
+        )
+        # adjust this threshold to control the filtering strength
+        high_pass_mask = distance > (d // 4)
+        # apply the high-pass filter
+        fft_shifted *= high_pass_mask
+
+        # inverse fft
+        fft_unshifted = np.fft.ifftshift(fft_shifted)
+        filtered_image = np.real(ifftn(fft_unshifted))
+
+        filtered_rescaled = np.clip(filtered_image, 0, 255).astype(np.uint8)
+        return filtered_rescaled

swcgeom/transforms/image_stack.py

@@ -25,6 +25,7 @@ from sdflit import (
     Scene,
     SDFObject,
 )
+from tqdm import tqdm
 
 from swcgeom.core import Population, Tree
 from swcgeom.transforms.base import Transform
@@ -89,8 +90,6 @@ class ToImageStack(Transform[Tree, npt.NDArray[np.uint8]]):
         samplers = self._get_samplers(coord_min, coord_max)
 
         if verbose:
-            from tqdm import tqdm
-
             total = (coord_max[2] - coord_min[2]) / self.resolution[2]
             samplers = tqdm(samplers, total=total.astype(np.int64).item())
 
@@ -117,8 +116,6 @@ class ToImageStack(Transform[Tree, npt.NDArray[np.uint8]]):
         )
 
         if verbose:
-            from tqdm import tqdm
-
             trees = tqdm(trees)
 
         # TODO: multiprocess