swcgeom 0.19.4__cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

swcgeom/transforms/images.py

@@ -0,0 +1,206 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Image stack related transform."""
+
+import numpy as np
+import numpy.typing as npt
+from typing_extensions import deprecated, override
+
+from swcgeom.transforms.base import Identity, Transform
+
+__all__ = [
+    "ImagesCenterCrop",
+    "ImagesScale",
+    "ImagesClip",
+    "ImagesFlip",
+    "ImagesFlipY",
+    "ImagesNormalizer",
+    "ImagesMeanVarianceAdjustment",
+    "ImagesScaleToUnitRange",
+    "ImagesHistogramEqualization",
+    "Center",  # legacy
+]
+
+
+NDArrayf32 = npt.NDArray[np.float32]
+
+
+class ImagesCenterCrop(Transform[NDArrayf32, NDArrayf32]):
+    """Get image stack center."""
+
+    def __init__(self, shape_out: int | tuple[int, int, int]):
+        super().__init__()
+        self.shape_out = (
+            shape_out
+            if isinstance(shape_out, tuple)
+            else (shape_out, shape_out, shape_out)
+        )
+
+    @override
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        diff = np.subtract(x.shape[:3], self.shape_out)
+        s = diff // 2
+        e = np.add(s, self.shape_out)
+        return x[s[0] : e[0], s[1] : e[1], s[2] : e[2], :]
+
+    @override
+    def extra_repr(self) -> str:
+        return f"shape_out=({','.join(str(a) for a in self.shape_out)})"
+
+
+@deprecated("use `ImagesCenterCrop` instead", stacklevel=2)
+class Center(ImagesCenterCrop):
+    """Get image stack center.
+
+    .. deprecated:: 0.16.0
+        Use :class:`ImagesCenterCrop` instead.
+    """
+
+
+class ImagesScale(Transform[NDArrayf32, NDArrayf32]):
+    def __init__(self, scaler: float) -> None:
+        super().__init__()
+        self.scaler = scaler
+
+    @override
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        return self.scaler * x
+
+    @override
+    def extra_repr(self) -> str:
+        return f"scaler={self.scaler}"
+
+
+class ImagesClip(Transform[NDArrayf32, NDArrayf32]):
+    def __init__(self, vmin: float = 0, vmax: float = 1, /) -> None:
+        super().__init__()
+        self.vmin, self.vmax = vmin, vmax
+
+    @override
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        return np.clip(x, self.vmin, self.vmax)
+
+    @override
+    def extra_repr(self) -> str:
+        return f"vmin={self.vmin}, vmax={self.vmax}"
+
+
+class ImagesFlip(Transform[NDArrayf32, NDArrayf32]):
+    """Flip image stack along axis."""
+
+    def __init__(self, axis: int, /) -> None:
+        super().__init__()
+        self.axis = axis
+
+    @override
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        return np.flip(x, axis=self.axis)
+
+    @override
+    def extra_repr(self) -> str:
+        return f"axis={self.axis}"
+
+
+class ImagesFlipY(ImagesFlip):
+    """Flip image stack along Y-axis.
+
+    See Also:
+        ~.images.io.TeraflyImageStack:
+            Terafly and Vaa3d use a especial right-handed coordinate system (with
+            origin point in the left-top and z-axis points front), but we flip y-axis
+            to makes it a left-handed coordinate system (with origin point in the
+            left-bottom and z-axis points front). If you need to use its coordinate
+            system, remember to FLIP Y-AXIS BACK.
+    """
+
+    def __init__(self, axis: int = 1, /) -> None:
+        super().__init__(axis)  # (X, Y, Z, C)
+
+
+class ImagesNormalizer(Transform[NDArrayf32, NDArrayf32]):
+    """Normalize image stack."""
+
+    @override
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        mean = np.mean(x)
+        variance = np.var(x)
+        return (x - mean) / variance
+
+
+class ImagesMeanVarianceAdjustment(Transform[NDArrayf32, NDArrayf32]):
+    """Adjust image stack mean and variance.
+
+    See Also:
+        ~swcgeom.images.ImageStackFolder.stat
+    """
+
+    def __init__(self, mean: float, variance: float) -> None:
+        super().__init__()
+        self.mean = mean
+        self.variance = variance
+
+    @override
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        return (x - self.mean) / self.variance
+
+    @override
+    def extra_repr(self) -> str:
+        return f"mean={self.mean}, variance={self.variance}"
+
+
+class ImagesScaleToUnitRange(Transform[NDArrayf32, NDArrayf32]):
+    """Scale image stack to unit range."""
+
+    def __init__(self, vmin: float, vmax: float, *, clip: bool = True) -> None:
+        """Scale image stack to unit range.
+
+        Args:
+            vmin: Minimum value.
+            vmax: Maximum value.
+            clip: Clip values to [0, 1] to avoid numerical issues.
+        """
+
+        super().__init__()
+        self.vmin = vmin
+        self.vmax = vmax
+        self.diff = vmax - vmin
+        self.clip = clip
+        self.post = ImagesClip(0, 1) if self.clip else Identity()
+
+    @override
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        return self.post((x - self.vmin) / self.diff)
+
+    @override
+    def extra_repr(self) -> str:
+        return f"vmin={self.vmin}, vmax={self.vmax}, clip={self.clip}"
+
+
+class ImagesHistogramEqualization(Transform[NDArrayf32, NDArrayf32]):
+    """Image histogram equalization.
+
+    References:
+        http://www.janeriksolem.net/histogram-equalization-with-python-and.html
+    """
+
+    def __init__(self, bins: int = 256) -> None:
+        super().__init__()
+        self.bins = bins
+
+    @override
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        # get image histogram
+        hist, bin_edges = np.histogram(x.flatten(), self.bins, density=True)
+        cdf = hist.cumsum()  # cumulative distribution function
+        cdf = cdf / cdf[-1]  # normalize
+
+        # use linear interpolation of cdf to find new pixel values
+        equalized = np.interp(x.flatten(), bin_edges[:-1], cdf)
+        return equalized.reshape(x.shape).astype(np.float32)
+
+    @override
+    def extra_repr(self) -> str:
+        return f"bins={self.bins}"

swcgeom/transforms/mst.py

@@ -0,0 +1,183 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Minimum spanning tree."""
+
+import warnings
+
+import numpy as np
+import pandas as pd
+from numpy import ma
+from numpy import typing as npt
+from typing_extensions import override
+
+from swcgeom.core import Tree, sort_tree
+from swcgeom.core.swc_utils import SWCNames, SWCTypes, get_names, get_types
+from swcgeom.transforms.base import Transform
+
+__all__ = ["PointsToCuntzMST", "PointsToMST"]
+
+
+class PointsToCuntzMST(Transform[npt.NDArray[np.float32], Tree]):
+    """Create tree from points.
+
+    Creates trees corresponding to the minimum spanning tree keeping
+    the path length to the root small (with balancing factor bf).
+
+    References:
+    .. [1] Cuntz, H., Forstner, F., Borst, A. & Häusser, M. One Rule to
+       Grow Them Al: A General Theory of Neuronal Branching and Its
+       Practical Application. PLOS Comput Biol 6, e1000877 (2010).
+    .. [2] Cuntz, H., Borst, A. & Segev, I. Optimization principles of
+       dendritic structure. Theor Biol Med Model 4, 21 (2007).
+    """
+
+    def __init__(
+        self,
+        *,
+        bf: float = 0.4,
+        furcations: int = 2,
+        exclude_soma: bool = True,
+        sort: bool = True,
+        names: SWCNames | None = None,
+        types: SWCTypes | None = None,
+    ) -> None:
+        """
+        Args:
+            bf: Balancing factor between 0~1.
+            furcations: Suppress multi-furcations which more than k.
+                If set to -1, no suppression.
+            exclude_soma: Suppress multi-furcations exclude soma.
+        """
+        self.bf = np.clip(bf, 0, 1)
+        self.furcations = furcations
+        self.exclude_soma = exclude_soma
+        self.sort = sort
+        self.names = get_names(names)
+        self.types = get_types(types)
+
+    @override
+    def __call__(  # pylint: disable=too-many-locals
+        self,
+        points: npt.NDArray[np.floating],
+        soma: npt.ArrayLike | None = None,
+        *,
+        names: SWCNames | None = None,
+    ) -> Tree:
+        """
+        Args:
+            points: Positions of points cloud.
+            soma: Position of soma. If none, use the first point as soma.
+        """
+        if names is None:
+            names = self.names
+        else:
+            warnings.warn(
+                "`PointsToCuntzMST(...)(names=...)` has been replaced by "
+                "`PointsToCuntzMST(...,names=...)` since v0.12.0, and will be removed "
+                "in next version",
+                DeprecationWarning,
+            )
+            names = get_names(names)  # TODO: remove it
+
+        if soma is not None:
+            soma = np.array(soma)
+            assert soma.shape == (3,)
+            points = np.concatenate([[soma], points])
+
+        n = points.shape[0]
+        dis = np.linalg.norm(
+            points.reshape((-1, 1, 3)) - points.reshape((1, -1, 3)), axis=2
+        )
+
+        pid = np.full(n, fill_value=-1)
+        acc = np.zeros(n)
+        furcations = np.zeros(n, dtype=np.int32)
+        conn = np.zeros(n, dtype=np.bool_)  # connect points
+        conn[0] = True
+        mask = np.ones((n, n), dtype=np.bool_)
+        mask[0, :] = False
+        mask[0, 0] = True
+        for _ in range(n - 1):  # for tree: e = n-1
+            cost = ma.array(dis + self.bf * acc, mask=mask)
+            (i, j) = np.unravel_index(cost.argmin(), cost.shape)
+
+            furcations[i] += 1
+            if (
+                self.furcations != -1
+                and furcations[i] >= self.furcations
+                and (not self.exclude_soma or i != 0)
+            ):
+                mask[i, :] = True  # avoid it link to any point
+                mask[:, i] = True  # avoid any point lint to it
+
+            pid[j] = i
+            acc[j] = acc[i] + dis[i, j]
+            conn[j] = True
+            mask[j, :] = conn  # enable link to other points
+            mask[:, j] = True  # avoid any point lint to it
+
+        dic = {
+            names.id: np.arange(n),
+            names.type: np.full(n, fill_value=self.types.glia_processes),
+            names.x: points[:, 0],
+            names.y: points[:, 1],
+            names.z: points[:, 2],
+            names.r: 1,
+            names.pid: pid,
+        }
+        dic[names.type][0] = self.types.soma
+        df = pd.DataFrame.from_dict(dic)
+        t = Tree.from_data_frame(df, names=names)
+        if self.sort:
+            t = sort_tree(t)
+        return t
+
+    @override
+    def extra_repr(self) -> str:  # TODO: names, types
+        return f"bf={self.bf:.4f}, furcations={self.furcations}, exclude_soma={self.exclude_soma}, sort={self.sort}"
+
+
+class PointsToMST(PointsToCuntzMST):  # pylint: disable=too-few-public-methods
+    """Create minimum spanning tree from points."""
+
+    def __init__(
+        self,
+        furcations: int = 2,
+        *,
+        k_furcations: int | None = None,
+        exclude_soma: bool = True,
+        names: SWCNames | None = None,
+        types: SWCTypes | None = None,
+        **kwargs,
+    ) -> None:
+        """
+        Args:
+            furcations: Suppress multifurcations which more than k.
+                If set to -1, no suppression.
+            exclude_soma: Suppress multi-furcations exclude soma.
+        """
+
+        if k_furcations is not None:
+            warnings.warn(
+                "`PointsToMST(k_furcations=...)` has been renamed to "
+                "`PointsToMST(furcations=...)` since v0.12.0, and will "
+                "be removed in next version",
+                DeprecationWarning,
+            )
+            furcations = k_furcations
+
+        super().__init__(
+            bf=0,
+            furcations=furcations,
+            exclude_soma=exclude_soma,
+            names=names,
+            types=types,
+            **kwargs,
+        )
+
+    @override
+    def extra_repr(self) -> str:
+        return f"furcations-{self.furcations}, exclude-soma={self.exclude_soma}"