swcgeom 0.19.4__cp311-cp311-macosx_10_9_universal2.whl

swcgeom/transforms/geometry.py

@@ -0,0 +1,270 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""SWC geometry operations."""
+
+import warnings
+from typing import Generic, Literal, TypeVar
+
+import numpy as np
+import numpy.typing as npt
+from typing_extensions import override
+
+from swcgeom.core import DictSWC
+from swcgeom.core.swc_utils import SWCNames
+from swcgeom.transforms.base import Transform
+from swcgeom.utils import (
+    rotate3d,
+    rotate3d_x,
+    rotate3d_y,
+    rotate3d_z,
+    scale3d,
+    translate3d,
+)
+
+__all__ = [
+    "Normalizer",
+    "RadiusReseter",
+    "AffineTransform",
+    "Translate",
+    "TranslateOrigin",
+    "Scale",
+    "Rotate",
+    "RotateX",
+    "RotateY",
+    "RotateZ",
+]
+
+T = TypeVar("T", bound=DictSWC)
+Center = Literal["root", "soma", "origin"]
+
+
+# pylint: disable=too-few-public-methods
+class Normalizer(Generic[T], Transform[T, T]):
+    """Noramlize coordinates and radius to 0-1."""
+
+    def __init__(self, *, names: SWCNames | None = None) -> None:
+        super().__init__()
+        if names is not None:
+            warnings.warn(
+                "`name` parameter is no longer needed, now use the "
+                "built-in names table, you can directly remove it.",
+                DeprecationWarning,
+            )
+
+    @override
+    def __call__(self, x: T) -> T:
+        """Scale the `x`, `y`, `z`, `r` of nodes to 0-1."""
+        new_tree = x.copy()
+        xyzr = [x.names.x, x.names.y, x.names.z, x.names.r]
+        for key in xyzr:  # TODO: does r is the same?
+            vs = new_tree.ndata[key]
+            new_tree.ndata[key] = (vs - np.min(vs)) / np.max(vs)
+
+        return new_tree
+
+
+class RadiusReseter(Generic[T], Transform[T, T]):
+    """Reset radius to fixed value."""
+
+    def __init__(self, r: float) -> None:
+        super().__init__()
+        self.r = r
+
+    def __call__(self, x: T) -> T:
+        r = np.full_like(x.r(), fill_value=self.r)
+        new_tree = x.copy()
+        new_tree.ndata[new_tree.names.r] = r
+        return new_tree
+
+    @override
+    def extra_repr(self) -> str:
+        return f"r={self.r:.4f}"
+
+
+class AffineTransform(Generic[T], Transform[T, T]):
+    """Apply affine matrix."""
+
+    tm: npt.NDArray[np.float32]
+    center: Center
+    fmt: str
+
+    def __init__(
+        self,
+        tm: npt.NDArray[np.float32],
+        center: Center = "origin",
+        *,
+        fmt: str | None = None,
+        names: SWCNames | None = None,
+    ) -> None:
+        self.tm, self.center = tm, center
+
+        if fmt is not None:
+            warnings.warn(
+                "`fmt` parameter is no longer needed, now use the "
+                "extra_repr(), you can directly remove it.",
+                DeprecationWarning,
+            )
+
+        if names is not None:
+            warnings.warn(
+                "`name` parameter is no longer needed, now use the "
+                "built-in names table, you can directly remove it.",
+                DeprecationWarning,
+            )
+
+    @override
+    def __call__(self, x: T) -> T:
+        match self.center:
+            case "root" | "soma":
+                idx = np.nonzero(x.ndata[x.names.pid] == -1)[0][0].item()
+                xyz = x.xyz()[idx]
+                tm = (
+                    translate3d(-xyz[0], -xyz[1], -xyz[2])
+                    .dot(self.tm)
+                    .dot(translate3d(xyz[0], xyz[1], xyz[2]))
+                )
+            case _:
+                tm = self.tm
+
+        return self.apply(x, tm)
+
+    @staticmethod
+    def apply(x: T, tm: npt.NDArray[np.float32]) -> T:
+        xyzw = x.xyzw().dot(tm.T).T
+        xyzw /= xyzw[3]
+
+        y = x.copy()
+        y.ndata[x.names.x] = xyzw[0]
+        y.ndata[x.names.y] = xyzw[1]
+        y.ndata[x.names.z] = xyzw[2]
+        return y
+
+
+class Translate(Generic[T], AffineTransform[T]):
+    """Translate SWC."""
+
+    def __init__(self, tx: float, ty: float, tz: float, **kwargs) -> None:
+        super().__init__(translate3d(tx, ty, tz), **kwargs)
+        self.tx, self.ty, self.tz = tx, ty, tz
+
+    @override
+    def extra_repr(self) -> str:
+        return f"tx={self.tx:.4f}, ty={self.ty:.4f}, tz={self.tz:.4f}"
+
+    @classmethod
+    def transform(cls, x: T, tx: float, ty: float, tz: float, **kwargs) -> T:
+        return cls(tx, ty, tz, **kwargs)(x)
+
+
+class TranslateOrigin(Generic[T], Transform[T, T]):
+    """Translate root of SWC to origin point."""
+
+    def __call__(self, x: T) -> T:
+        return self.transform(x)
+
+    @classmethod
+    def transform(cls, x: T) -> T:
+        pid = np.nonzero(x.ndata[x.names.pid] == -1)[0][0].item()
+        xyzw = x.xyzw()
+        tm = translate3d(-xyzw[pid, 0], -xyzw[pid, 1], -xyzw[pid, 2])
+        return AffineTransform.apply(x, tm)
+
+
+class Scale(Generic[T], AffineTransform[T]):
+    """Scale SWC."""
+
+    def __init__(
+        self, sx: float, sy: float, sz: float, center: Center = "root", **kwargs
+    ) -> None:
+        super().__init__(scale3d(sx, sy, sz), center=center, **kwargs)
+
+    @classmethod
+    def transform(  # pylint: disable=too-many-arguments
+        cls, x: T, sx: float, sy: float, sz: float, center: Center = "root", **kwargs
+    ) -> T:
+        return cls(sx, sy, sz, center=center, **kwargs)(x)
+
+
+class Rotate(Generic[T], AffineTransform[T]):
+    """Rotate SWC."""
+
+    def __init__(
+        self,
+        n: npt.NDArray[np.float32],
+        theta: float,
+        center: Center = "root",
+        **kwargs,
+    ) -> None:
+        fmt = f"Rotate-{n[0]}-{n[1]}-{n[2]}-{theta:.4f}"
+        super().__init__(rotate3d(n, theta), center=center, fmt=fmt, **kwargs)
+        self.n = n
+        self.theta = theta
+        self.center = center
+
+    @override
+    def extra_repr(self) -> str:
+        return f"n={self.n}, theta={self.theta:.4f}, center={self.center}"  # TODO: improve format of n
+
+    @classmethod
+    def transform(
+        cls,
+        x: T,
+        n: npt.NDArray[np.float32],
+        theta: float,
+        center: Center = "root",
+        **kwargs,
+    ) -> T:
+        return cls(n, theta, center=center, **kwargs)(x)
+
+
+class RotateX(Generic[T], AffineTransform[T]):
+    """Rotate SWC with x-axis."""
+
+    def __init__(self, theta: float, center: Center = "root", **kwargs) -> None:
+        super().__init__(rotate3d_x(theta), center=center, **kwargs)
+        self.theta = theta
+
+    @override
+    def extra_repr(self) -> str:
+        return f"center={self.center}, theta={self.theta:.4f}"
+
+    @classmethod
+    def transform(cls, x: T, theta: float, center: Center = "root", **kwargs) -> T:
+        return cls(theta, center=center, **kwargs)(x)
+
+
+class RotateY(Generic[T], AffineTransform[T]):
+    """Rotate SWC with y-axis."""
+
+    def __init__(self, theta: float, center: Center = "root", **kwargs) -> None:
+        super().__init__(rotate3d_y(theta), center=center, **kwargs)
+        self.theta = theta
+        self.center = center
+
+    @override
+    def extra_repr(self) -> str:
+        return f"theta={self.theta:.4f}, center={self.center}"
+
+    @classmethod
+    def transform(cls, x: T, theta: float, center: Center = "root", **kwargs) -> T:
+        return cls(theta, center=center, **kwargs)(x)
+
+
+class RotateZ(Generic[T], AffineTransform[T]):
+    """Rotate SWC with z-axis."""
+
+    def __init__(self, theta: float, center: Center = "root", **kwargs) -> None:
+        super().__init__(rotate3d_z(theta), center=center, **kwargs)
+        self.theta = theta
+        self.center = center
+
+    @override
+    def extra_repr(self) -> str:
+        return f"theta={self.theta:.4f}, center={self.center}"
+
+    @classmethod
+    def transform(cls, x: T, theta: float, center: Center = "root", **kwargs) -> T:
+        return cls(theta, center=center, **kwargs)(x)

swcgeom/transforms/image_preprocess.py

@@ -0,0 +1,107 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""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 typing_extensions import override
+
+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
+    """
+
+    @override
+    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

@@ -0,0 +1,219 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Create image stack from morphology.
+
+NOTE: All denpendencies need to be installed, try:
+
+```sh
+pip install swcgeom[all]
+```
+"""
+
+import os
+import re
+import time
+from collections.abc import Iterable
+from typing import Sequence
+
+import numpy as np
+import numpy.typing as npt
+import tifffile
+from sdflit import (
+    ColoredMaterial,
+    ObjectsScene,
+    RangeSampler,
+    RoundCone,
+    Scene,
+    SDFObject,
+)
+from tqdm import tqdm
+from typing_extensions import deprecated, override
+
+from swcgeom.core import Population, Tree
+from swcgeom.transforms.base import Transform
+
+__all__ = ["ToImageStack"]
+
+
+class ToImageStack(Transform[Tree, npt.NDArray[np.uint8]]):
+    r"""Transform tree to image stack."""
+
+    resolution: npt.NDArray[np.float32]
+
+    def __init__(self, resolution: int | float | npt.ArrayLike = 1) -> None:
+        """Transform tree to image stack.
+
+        Args:
+            resolution: Resolution of image stack.
+                If a scalar, it will be broadcasted to a vector of 3d.
+        """
+        if isinstance(resolution, (int, float, np.integer, np.floating)):
+            resolution = [resolution, resolution, resolution]  # type: ignore
+
+        self.resolution = np.array(resolution, dtype=np.float32)
+        assert len(self.resolution) == 3, "resolution should be vector of 3d."
+
+    @override
+    def __call__(self, x: Tree) -> npt.NDArray[np.uint8]:
+        """Transform tree to image stack.
+
+        NOTE: This method loads the entire image stack into memory, so it ONLY works
+        for small image stacks, use :meth`transform_and_save` for big image stack.
+        """
+        return np.stack(list(self.transform(x, verbose=False)), axis=0)
+
+    def transform(
+        self,
+        x: Tree | Sequence[Tree],
+        verbose: bool = True,
+        *,
+        ranges: tuple[npt.ArrayLike, npt.ArrayLike] | None = None,
+    ) -> Iterable[npt.NDArray[np.uint8]]:
+        trees = [x] if isinstance(x, Tree) else x
+        if not trees:
+            return iter([])  # Return empty iterator if sequence is empty
+
+        time_start = None
+        if verbose:
+            sources = ", ".join(t.source for t in trees if t.source)
+            print(f"To image stack: {sources if sources else 'unnamed trees'}")
+            time_start = time.time()
+
+        scene = self._get_scene(trees)
+
+        if ranges is None:
+            all_xyz = np.concatenate([t.xyz() for t in trees], axis=0)
+            all_r = np.concatenate([t.r() for t in trees], axis=0).reshape(-1, 1)
+            if all_xyz.size == 0:  # Handle empty trees
+                coord_min = np.zeros(3, dtype=np.float32)
+                coord_max = np.zeros(3, dtype=np.float32)
+            else:
+                coord_min = np.floor(np.min(all_xyz - all_r, axis=0))
+                coord_max = np.ceil(np.max(all_xyz + all_r, axis=0))
+        else:
+            assert len(ranges) == 2
+            coord_min = np.array(ranges[0])
+            coord_max = np.array(ranges[1])
+            assert len(coord_min) == len(coord_max) == 3
+
+        samplers = self._get_samplers(coord_min, coord_max)
+
+        if verbose and time_start is not None:
+            total = (coord_max[2] - coord_min[2]) / self.resolution[2]
+            samplers = tqdm(samplers, total=total.astype(np.int64).item())
+
+            time_end = time.time()
+            print("Prepare in: ", time_end - time_start, "s")
+
+        for sampler in samplers:
+            voxel = sampler.sample(scene)  # should be shape of (x, y, z, 3) and z = 1
+            frame = (255 * voxel[..., 0, 0]).astype(np.uint8)
+            yield frame
+
+    @deprecated("Use transform instead")
+    def transfrom(
+        self,
+        x: Tree,
+        verbose: bool = True,
+        *,
+        ranges: tuple[npt.ArrayLike, npt.ArrayLike] | None = None,
+    ) -> Iterable[npt.NDArray[np.uint8]]:
+        return self.transform(x, verbose, ranges=ranges)
+
+    def transform_and_save(
+        self, fname: str, x: Tree | Sequence[Tree], verbose: bool = True, **kwargs
+    ) -> None:
+        self.save_tif(fname, self.transform(x, verbose=verbose, **kwargs))
+
+    def transform_population(
+        self, population: Population | str, verbose: bool = True
+    ) -> None:
+        trees = (
+            Population.from_swc(population)
+            if isinstance(population, str)
+            else population
+        )
+
+        if verbose:
+            trees = tqdm(trees)
+
+        # TODO: multiprocess
+        for tree in trees:
+            tif = re.sub(r".swc$", ".tif", tree.source)
+            if not os.path.isfile(tif):
+                self.transform_and_save(tif, tree, verbose=False)
+
+    @override
+    def extra_repr(self) -> str:
+        res = ",".join(f"{a:.4f}" for a in self.resolution)
+        return f"resolution=({res})"
+
+    def _get_scene(self, trees: Sequence[Tree]) -> Scene:
+        material = ColoredMaterial((1, 0, 0)).into()
+        scene = ObjectsScene()
+        scene.set_background((0, 0, 0))
+
+        def leave(n: Tree.Node, children: list[Tree.Node]) -> Tree.Node:
+            for c in children:
+                sdf = RoundCone(_tp3f(n.xyz()), _tp3f(c.xyz()), n.r, c.r).into()
+                scene.add_object(SDFObject(sdf, material).into())
+            return n
+
+        for tree in trees:
+            tree.traverse(leave=leave)
+
+        scene.build_bvh()
+        return scene.into()
+
+    def _get_samplers(
+        self,
+        coord_min: npt.NDArray,
+        coord_max: npt.NDArray,
+        offset: npt.NDArray | None = None,
+    ) -> Iterable[RangeSampler]:
+        """Get Samplers.
+
+        Args:
+            coord_min, coord_max: Coordinates array of shape (3,).
+        """
+
+        eps = 1e-6
+        stride = self.resolution
+        offset = offset or (stride / 2)
+
+        xmin, ymin, zmin = _tp3f(coord_min + offset)
+        xmax, ymax, zmax = _tp3f(coord_max)
+        z = zmin
+        while z < zmax:
+            yield RangeSampler(
+                (xmin, ymin, z), (xmax, ymax, z + stride[2] - eps), _tp3f(stride)
+            )
+            z += stride[2]
+
+    @staticmethod
+    def save_tif(
+        fname: str,
+        frames: Iterable[npt.NDArray[np.uint8]],
+        resolution: tuple[float, float] = (1, 1),
+    ) -> None:
+        with tifffile.TiffWriter(fname) as tif:
+            for frame in frames:
+                tif.write(
+                    frame,
+                    contiguous=True,
+                    photometric="minisblack",
+                    resolution=resolution,
+                    metadata={
+                        "unit": "um",
+                        "axes": "ZXY",
+                    },
+                )
+
+
+def _tp3f(x: npt.NDArray) -> tuple[float, float, float]:
+    """Convert to tuple of 3 floats."""
+    assert len(x) == 3
+    return (float(x[0]), float(x[1]), float(x[2]))