swcgeom 0.19.4__cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

swcgeom/core/tree_utils_impl.py

@@ -0,0 +1,58 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""SWC util wrapper for tree, split to avoid circle imports.
+
+NOTE: Do not import `Tree` and keep this file minimized.
+"""
+
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+
+from swcgeom.core.swc import SWCLike, SWCNames
+from swcgeom.core.swc_utils import Topology, to_sub_topology, traverse
+
+__all__ = ["get_subtree_impl", "to_subtree_impl"]
+
+Mapping = dict[int, int] | list[int]
+TreeArgs = tuple[int, dict[str, npt.NDArray[Any]], str, SWCNames]
+
+
+def get_subtree_impl(
+    swc_like: SWCLike, n: int, *, out_mapping: Mapping | None = None
+) -> TreeArgs:
+    ids = []
+    topo = (swc_like.id(), swc_like.pid())
+    traverse(topo, enter=lambda n, _: ids.append(n), root=n)
+
+    sub_ids = np.array(ids, dtype=np.int32)
+    sub_pid = swc_like.pid()[sub_ids]
+    sub_pid[0] = -1
+    return to_subtree_impl(swc_like, (sub_ids, sub_pid), out_mapping=out_mapping)
+
+
+def to_subtree_impl(
+    swc_like: SWCLike,
+    sub: Topology,
+    *,
+    out_mapping: Mapping | None = None,
+) -> TreeArgs:
+    (new_id, new_pid), mapping = to_sub_topology(sub)
+
+    n_nodes = new_id.shape[0]
+    ndata = {k: swc_like.get_ndata(k)[mapping].copy() for k in swc_like.keys()}
+    ndata.update(id=new_id, pid=new_pid)
+
+    if isinstance(out_mapping, list):
+        out_mapping.clear()
+        out_mapping.extend(mapping)
+    elif isinstance(out_mapping, dict):
+        out_mapping.clear()
+        for new_id, old_id in enumerate(mapping):
+            out_mapping[new_id] = old_id  # returning a dict may leads to bad perf
+
+    return n_nodes, ndata, swc_like.source, swc_like.names

swcgeom/images/__init__.py

@@ -0,0 +1,9 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Image Stack Related."""
+
+from swcgeom.images.folder import *  # noqa: F403
+from swcgeom.images.io import *  # noqa: F403

swcgeom/images/augmentation.py

@@ -0,0 +1,149 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Play augment in image stack.
+
+NOTE: This is expremental code, and the API is subject to change.
+"""
+
+import random
+from typing import Literal
+
+import numpy as np
+import numpy.typing as npt
+
+__all__ = ["play_augment", "random_augmentations"]
+
+NDArrf32 = npt.NDArray[np.float32]
+
+# Augmentation = Literal[
+#     "swap_xy",
+#     "swap_xz",
+#     "swap_yz",
+#     "flip_x",
+#     "flip_y",
+#     "flip_z",
+#     "rot90_xy",
+#     "rot90_xz",
+#     "rot90_yz",
+# ]
+
+IDENTITY = -1
+
+augs = {
+    # swaps
+    "swap_xy": lambda x: np.swapaxes(x, 0, 1),
+    "swap_xz": lambda x: np.swapaxes(x, 0, 2),
+    "swap_yz": lambda x: np.swapaxes(x, 1, 2),
+    # flips
+    "flip_x": lambda x: np.flip(x, axis=[0]),
+    "flip_y": lambda x: np.flip(x, axis=[1]),
+    "flip_z": lambda x: np.flip(x, axis=[2]),
+    # rotations
+    "rot90_xy": lambda x: np.rot90(x, k=1, axes=(0, 1)),
+    "rot90_xz": lambda x: np.rot90(x, k=1, axes=(0, 2)),
+    "rot90_yz": lambda x: np.rot90(x, k=1, axes=(1, 2)),
+}
+
+
+class Augmentation:
+    """Play augmentation."""
+
+    def __init__(self, *, seed: int | None) -> None:
+        self.seed = seed
+        self.rand = random.Random(seed)
+
+    def swapaxes(self, x, mode: Literal["xy", "xz", "yz"] | None = None) -> NDArrf32:
+        if mode is None:
+            modes: list[Literal["xy", "xz", "yz"]] = ["xy", "xz", "yz"]
+            mode = modes[self.rand.randint(0, 2)]
+
+        match mode:
+            case "xy":
+                return np.swapaxes(x, 0, 1)
+            case "xz":
+                return np.swapaxes(x, 0, 2)
+            case "yz":
+                return np.swapaxes(x, 1, 2)
+            case _:
+                raise ValueError(f"invalid mode: {mode}")
+
+    def flip(self, x, mode: Literal["xy", "xz", "yz"] | None = None) -> NDArrf32:
+        if mode is None:
+            modes: list[Literal["xy", "xz", "yz"]] = ["xy", "xz", "yz"]
+            mode = modes[random.randint(0, 2)]
+
+        match mode:
+            case "xy":
+                return np.flip(x, axis=0)
+            case "xz":
+                return np.flip(x, axis=1)
+            case "yz":
+                return np.flip(x, axis=2)
+            case _:
+                raise ValueError(f"invalid mode: {mode}")
+
+
+fns = list(augs.keys())
+
+
+def play_augment(x: NDArrf32, method: Augmentation | int | None = None) -> NDArrf32:
+    """Play augment in x.
+
+    Args
+        x: Array of shape (X, Y, Z, C)
+        method: Augmentation method index / name.
+            If not provided, a random augment will be apply.
+    """
+
+    if isinstance(method, str):
+        key = method
+    elif isinstance(method, int):
+        key = fns[method]
+    elif method is None:
+        key = fns[random.randint(0, len(augs))]
+    elif method == IDENTITY:
+        return x
+    else:
+        raise ValueError("invalid augment method")
+
+    return augs[key](x)
+
+
+def random_augmentations(
+    n: int, k: int, *, seed: int | None = None, include_identity: bool = True
+) -> npt.NDArray[np.int64]:
+    """Generate a sequence of augmentations.
+
+    >>> xs = os.listdir("path_to_imgs")  # doctest: +SKIP
+    >>> augs = generate_random_augmentations(len(xs), 5)  # doctest: +SKIP
+    >>> for i, j in range(augs):  # doctest: +SKIP
+    ...     x = play_augment(read_imgs(os.path.join("path_to_imgs", xs[i])), j)
+
+    Args:
+        n: Size of image stacks.
+        k: Each image stack augmented to K image stack.
+        seed: Random seed, forwarding to `random.Random`
+        include_identity: Include identity transform.
+
+    Returns:
+        augmentations: List of (int, int)
+            Sequence of length N * K, contains index of image and augmentation method.
+    """
+
+    rand = random.Random(seed)
+    seq = list(range(len(augs)))
+    if include_identity:
+        seq.append(IDENTITY)
+
+    assert 0 < k < len(seq), "too large augment specify."
+
+    augmentations = []
+    for _ in range(n):
+        rand.shuffle(seq)
+        augmentations.extend(seq[:k])
+
+    xs = np.stack([np.repeat(np.arange(n), k), augmentations])
+    return xs

swcgeom/images/contrast.py

@@ -0,0 +1,87 @@
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""The contrast of an image.
+
+NOTE: This is expremental code, and the API is subject to change.
+"""
+
+from typing import 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.
+
+    Args:
+        imgs: ndarray
+
+    Returns:
+        contrast
+    """
+    ...
+
+
+@overload
+def contrast_std(image: Array3D, contrast: float) -> Array3D:
+    """Adjust the contrast of an image stack.
+
+    Args:
+        imgs: ndarray
+        contrast: The contrast adjustment factor. 1.0 leaves the image unchanged.
+
+    Returns:
+        imgs: The adjusted image.
+    """
+    ...
+
+
+def contrast_std(image: Array3D, contrast: float | None = 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.
+
+    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.
+
+    Returns:
+        contrast
+    """
+    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.
+
+    Args:
+        imgs: ndarray
+        mask: The mask to segment the foreground and background.
+            1 for foreground, 0 for background.
+
+    Returns:
+        contrast
+    """
+    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

@@ -0,0 +1,217 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Image stack folder."""
+
+import math
+import os
+import re
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass
+from typing import Generic, Literal, TypeVar, overload
+
+import numpy as np
+import numpy.typing as npt
+from tqdm import tqdm
+from typing_extensions import Self, deprecated
+
+from swcgeom.images.io import ScalarType, read_imgs
+from swcgeom.transforms import Identity, Transform
+
+__all__ = ["ImageStackFolder", "LabeledImageStackFolder", "PathImageStackFolder"]
+
+T = TypeVar("T")
+
+
+class ImageStackFolderBase(Generic[ScalarType, T]):
+    """Image stack folder base."""
+
+    files: list[str]
+    transform: Transform[npt.NDArray[ScalarType], T]
+
+    @overload
+    def __init__(
+        self,
+        files: Iterable[str],
+        *,
+        transform: Transform[npt.NDArray[np.float32], T] | None = ...,
+    ) -> None: ...
+    @overload
+    def __init__(
+        self,
+        files: Iterable[str],
+        *,
+        dtype: ScalarType,
+        transform: Transform[npt.NDArray[ScalarType], T] | None = ...,
+    ) -> None: ...
+    def __init__(self, files: Iterable[str], *, dtype=None, transform=None) -> None:
+        super().__init__()
+        self.files = list(files)
+        self.dtype = dtype or np.float32
+        self.transform = transform or Identity()  # type: ignore
+
+    def __len__(self) -> int:
+        return len(self.files)
+
+    def _get(self, fname: str) -> T:
+        imgs = self._read(fname)
+        imgs = self.transform(imgs)
+        return imgs
+
+    def _read(self, fname: str) -> npt.NDArray[ScalarType]:
+        return read_imgs(fname, dtype=self.dtype).get_full()  # type: ignore
+
+    @staticmethod
+    def scan(root: str, *, pattern: str | None = None) -> list[str]:
+        if not os.path.isdir(root):
+            raise NotADirectoryError(f"not a directory: {root}")
+
+        is_valid = re.compile(pattern).match if pattern is not None else truthly
+
+        fs = []
+        for d, _, files in os.walk(root):
+            fs.extend(os.path.join(d, f) for f in files if is_valid(f))
+
+        return fs
+
+    @staticmethod
+    @deprecated("Use `~swcgeom.images.io.read_imgs(fname).get_full()` instead")
+    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.
+        """
+
+        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.
+
+        NOTE: We are asserting that the images are of the same shape.
+
+        Args:
+            transform: Apply transform to the images.
+                If True, you need to make sure the transformed data is a ndarray.
+            verbose: Show progress bar.
+        """
+
+        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
+            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: str | None = None, **kwargs) -> Self:
+        """
+        Args:
+            root: str
+            pattern: Filter files by pattern.
+            **kwargs: Pass to `cls.__init__`
+        """
+
+        return cls(cls.scan(root, pattern=pattern), **kwargs)
+
+
+class LabeledImageStackFolder(ImageStackFolderBase[ScalarType, T]):
+    """Image stack folder with label."""
+
+    labels: list[int]
+
+    def __init__(self, files: Iterable[str], labels: Iterable[int], **kwargs):
+        super().__init__(files, **kwargs)
+        self.labels = list(labels)
+
+    def __getitem__(self, idx: int) -> tuple[T, int]:
+        return self._get(self.files[idx]), self.labels[idx]
+
+    @classmethod
+    def from_dir(
+        cls,
+        root: str,
+        label: int | Callable[[str], int],
+        *,
+        pattern: str | None = None,
+        **kwargs,
+    ) -> Self:
+        files = cls.scan(root, pattern=pattern)
+        if callable(label):
+            labels = [label(f) for f in files]
+        elif isinstance(label, int):
+            labels = [label for _ in files]
+        else:
+            raise ValueError("invalid label")
+        return cls(files, labels, **kwargs)
+
+
+class PathImageStackFolder(ImageStackFolderBase[ScalarType, T]):
+    """Image stack folder with relpath."""
+
+    root: str
+
+    def __init__(self, files: Iterable[str], *, root: str, **kwargs):
+        super().__init__(files, **kwargs)
+        self.root = root
+
+    def __getitem__(self, idx: int) -> tuple[T, str]:
+        relpath = os.path.relpath(self.files[idx], self.root)
+        return self._get(self.files[idx]), relpath
+
+    @classmethod
+    def from_dir(cls, root: str, *, pattern: str | None = None, **kwargs) -> Self:
+        """
+        Args:
+            root: str
+            pattern: Filter files by pattern.
+            **kwargs: Pass to `cls.__init__`
+        """
+
+        return cls(cls.scan(root, pattern=pattern), root=root, **kwargs)
+
+
+def truthly(*args, **kwargs) -> Literal[True]:  # pylint: disable=unused-argument
+    return True