swcgeom 0.20.0__cp312-cp312-macosx_15_0_arm64.whl

swcgeom/images/loaders/raw.pyx

@@ -0,0 +1,183 @@
+# -----------------------------------------------------------------------------
+# This file is adapted from the v3d-py-helper project:
+# https://github.com/SEU-ALLEN-codebase/v3d-py-helper
+#
+# Original license: MIT License
+# Copyright (c) Zuohan Zhao
+#
+# Vaa3D in Python Made Easy
+# Python library for Vaa3D functions.
+#
+# The original project is distributed via PyPI under the name `v3d-py-helper`,
+# with its latest release (v0.4.1) not supporting Python 3.13.
+#
+# As of Python 3.13 (released on October 7, 2024), this package fails to build
+# from source due to missing dependencies (e.g., libtiff), and no prebuilt wheels
+# are available on PyPI for Python 3.13. An issue has been raised, but the project
+# appears to be unmaintained at this time, and the author has not responded.
+#
+# To ensure continued compatibility and usability of Vaa3D features under Python 3.13+,
+# we have copied and minimally adapted necessary source files into this project,
+# preserving license and attribution in accordance with the MIT License.
+#
+# Please consult the original repository for full documentation:
+# https://SEU-ALLEN-codebase.github.io/v3d-py-helper
+#
+# If the upstream project resumes maintenance and releases official support
+# for Python 3.13+, this bundled version may be deprecated in favor of the
+# canonical package.
+# -----------------------------------------------------------------------------
+
+import io
+import os
+import struct
+import sys
+import cython
+cimport cython
+import numpy as np
+cimport numpy as np
+
+
+DEF FORMAT_KEY_4 = b"raw_image_stack_by_hpeng"
+DEF FORMAT_KEY_5 = b"raw5image_stack_by_hpeng"
+DEF FORMAT_LEN = 24
+
+
+cdef class Raw:
+    """
+    For vaa3d raw formats, allowing uint8, uint16 and float32 data types. The image dimension is presumed to be 4.
+
+    As a raw format, it not only stores the image buffer, but also the image dimension sizes and its data type in the
+    header. It also stores the endian of the buffer.
+
+    This interface saves and loads a multi-dimension numpy array of either of the 3 data types. It also compares the endian of
+    the file/numpy array and the machine to make it work properly. Based on the v3draw format key in the header, the
+    array can be 4D or 5D.
+
+    modified from v3d_external/v3d_main/basic_c_fun/stackutils.cpp
+
+    by Zuohan Zhao
+
+    2022/5/8
+    """
+    cdef:
+        bint sz2byte
+
+    def __init__(self, sz2byte = False):
+        """
+
+        :param sz2byte: set size array to be 2 byte (short int), for compatibility. Default as False.
+        """
+        self.sz2byte = sz2byte
+
+    cpdef np.ndarray load(self, path: str | os.PathLike, int choose = -1):
+        """
+        :param path: input image path of v3draw.
+        :param choose: choose a channel(4D) or stack(5D) to load, starting from 0, default as -1, meaning all.
+        :return: a numpy array of 4D or 5D based on the format key.
+        """
+        cdef:
+            short datatype
+            list sz
+            bytes endian_code_data
+            str endian, dt
+            char dim, header_sz, i
+            bytes format_key
+            long long bulk_sz, filesize
+
+        filesize = os.path.getsize(path)
+        assert filesize >= FORMAT_LEN, "File size too small, file might be corrupted"
+
+        with open(path, "rb") as f:
+            format_key = f.read(FORMAT_LEN)
+            if format_key == FORMAT_KEY_4:
+                dim = 4
+            elif format_key == FORMAT_KEY_5:
+                dim = 5
+            else:
+                raise RuntimeError("Format key isn't for v3draw")
+            if self.sz2byte:
+                header_sz = FORMAT_LEN + dim * 2 + 2 + 1
+            else:
+                header_sz = FORMAT_LEN + dim * 4 + 2 + 1
+            assert filesize >= header_sz, "File size too small, file might be corrupted"
+            endian_code_data = f.read(1)
+            if endian_code_data == b'B':
+                endian = '>'
+            elif endian_code_data == b'L':
+                endian = '<'
+            else:
+                raise RuntimeError('Endian code should be either B/L')
+            datatype = struct.unpack(f'{endian}h', f.read(2))[0]
+            if datatype == 1:
+                dt = 'u1'
+            elif datatype == 2:
+                dt = 'u2'
+            elif datatype == 4:
+                dt = 'f4'
+            else:
+                raise RuntimeError('v3draw data type can only be 1/2/4')
+            if self.sz2byte:
+                sz = list(struct.unpack(f'{endian}{dim}h', f.read(dim * 2)))
+            else:
+                sz = list(struct.unpack(f'{endian}{dim}i', f.read(dim * 4)))
+            bulk_sz = 1
+            for i in range(dim - 1):
+                bulk_sz *= sz[i]
+            assert bulk_sz * sz[-1] * datatype + header_sz == filesize, "file size doesn't match with the image"
+            if choose < 0:
+                return np.frombuffer(f.read(), endian + dt).reshape(sz[::-1])
+            else:
+                assert choose < sz[-1], "Choose index exceeding the range"
+                f.seek(bulk_sz * datatype * choose, 1)
+                img = np.frombuffer(f.read(bulk_sz * datatype), endian + dt)
+                return img.reshape(sz[-2::-1])
+
+    cpdef void save(self, path: str | os.PathLike, np.ndarray img):
+        """
+        :param path: output image path of v3draw.
+        :param img: the image array to save. 
+        """
+        assert img.ndim == 4, "The image has to be 4D"
+        assert img.dtype in [np.uint8, np.uint16], "The pixel type has to be uint8 or uint16"
+        cdef:
+            short datatype
+            list sz
+            bytes endian_code_data, format
+            str endian, bo = img.dtype.byteorder
+            bytes header
+            char dim = img.ndim
+
+        if dim == 4:
+            format = FORMAT_KEY_4
+        elif dim == 5:
+            format = FORMAT_KEY_5
+        else:
+            raise RuntimeError("Dimension not supported by v3draw")
+
+        sz = [img.shape[i] for i in range(dim)]
+        sz.extend([1] * (dim - len(sz)))
+        sz.reverse()
+
+        if bo == '>' or bo == '=' and sys.byteorder == 'big':
+            endian_code_data = b'B'
+            endian = '>'
+        else:
+            endian_code_data = b'L'
+            endian = '<'
+
+        if img.dtype == np.uint8:
+            datatype = 1
+        elif img.dtype == np.uint16:
+            datatype = 2
+        elif img.dtype == np.float32:
+            datatype = 4
+        else:
+            raise RuntimeError("numpy data type not supported by v3draw")
+
+        header = struct.pack(f'{endian}{FORMAT_LEN}sch4{"h" if self.sz2byte else "i"}',
+                             format, endian_code_data, datatype, *sz)
+
+        with open(path, 'wb') as f:
+            f.write(header)
+            f.write(img.tobytes())

swcgeom/transforms/__init__.py

@@ -0,0 +1,20 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""A series of transformations to compose codes."""
+
+from swcgeom.transforms.base import *  # noqa: F403
+from swcgeom.transforms.branch import *  # noqa: F403
+from swcgeom.transforms.branch_tree import *  # noqa: F403
+from swcgeom.transforms.geometry import *  # noqa: F403
+from swcgeom.transforms.image_preprocess import *  # noqa: F403
+from swcgeom.transforms.image_stack import *  # noqa: F403
+from swcgeom.transforms.images import *  # noqa: F403
+from swcgeom.transforms.mst import *  # noqa: F403
+from swcgeom.transforms.neurolucida_asc import *  # noqa: F403
+from swcgeom.transforms.path import *  # noqa: F403
+from swcgeom.transforms.population import *  # noqa: F403
+from swcgeom.transforms.tree import *  # noqa: F403
+from swcgeom.transforms.tree_assembler import *  # noqa: F403

swcgeom/transforms/base.py

@@ -0,0 +1,136 @@
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Transformation in tree."""
+
+from abc import ABC, abstractmethod
+from typing import Any, Generic, TypeVar, overload
+
+from typing_extensions import override
+
+__all__ = ["Transform", "Transforms", "Identity"]
+
+T = TypeVar("T")
+K = TypeVar("K")
+
+T1 = TypeVar("T1")
+T2 = TypeVar("T2")
+T3 = TypeVar("T3")
+T4 = TypeVar("T4")
+T5 = TypeVar("T5")
+T6 = TypeVar("T6")
+
+
+class Transform(ABC, Generic[T, K]):
+    r"""An abstract class representing a :class:`Transform`.
+
+    All transforms that represent a map from `T` to `K`.
+    """
+
+    @abstractmethod
+    def __call__(self, x: T) -> K:
+        """Apply transform.
+
+        NOTE: All subclasses should overwrite :meth:`__call__`, supporting
+        applying transform in `x`.
+        """
+        raise NotImplementedError()
+
+    def __repr__(self) -> str:
+        classname = self.__class__.__name__
+        repr_ = self.extra_repr()
+        return f"{classname}({repr_})"
+
+    def extra_repr(self) -> str:
+        """Provides a human-friendly representation of the module.
+
+        This method extends the basic string representation provided by
+        `__repr__` method. It is designed to display additional details
+        about the module's parameters or its specific configuration,
+        which can be particularly useful for debugging and model
+        architecture introspection.
+
+        >>> class Foo(Transform[T, K]):
+        ...     def __init__(self, my_parameter: int = 1):
+        ...         self.my_parameter = my_parameter
+        ...
+        ...     def extra_repr(self) -> str:
+        ...         return f"my_parameter={self.my_parameter}"
+
+        NOTE: This method should be overridden in custom modules to provide
+        specific details relevant to the module's functionality and
+        configuration.
+        """
+        return ""
+
+
+class Transforms(Transform[T, K]):
+    """A simple typed wrapper for transforms."""
+
+    transforms: list[Transform[Any, Any]]
+
+    # fmt: off
+    @overload
+    def __init__(self, t1: Transform[T, K], /) -> None: ...
+    @overload
+    def __init__(self, t1: Transform[T, T1],  t2: Transform[T1, K], /) -> None: ...
+    @overload
+    def __init__(self, t1: Transform[T, T1],  t2: Transform[T1, T2],
+                       t3: Transform[T2, K], /) -> None: ...
+    @overload
+    def __init__(self, t1: Transform[T, T1],  t2: Transform[T1, T2],
+                       t3: Transform[T2, T3], t4: Transform[T3, K], /) -> None: ...
+    @overload
+    def __init__(self, t1: Transform[T, T1],  t2: Transform[T1, T2],
+                       t3: Transform[T2, T3], t4: Transform[T3, T4],
+                       t5: Transform[T4, K], /) -> None: ...
+    @overload
+    def __init__(self, t1: Transform[T, T1],  t2: Transform[T1, T2],
+                       t3: Transform[T2, T3], t4: Transform[T3, T4],
+                       t5: Transform[T4, T5], t6: Transform[T5, K], /) -> None: ...
+    @overload
+    def __init__(self, t1: Transform[T, T1],  t2: Transform[T1, T2],
+                       t3: Transform[T2, T3], t4: Transform[T3, T4],
+                       t5: Transform[T4, T5], t6: Transform[T5, T6],
+                       t7: Transform[T6, K], /) -> None: ...
+    @overload
+    def __init__(self, t1: Transform[T, T1],  t2: Transform[T1, T2],
+                       t3: Transform[T2, T3], t4: Transform[T3, T4],
+                       t5: Transform[T4, T5], t6: Transform[T5, T6],
+                       t7: Transform[T6, Any], /, *transforms: Transform[Any, K]) -> None: ...
+    # fmt: on
+    def __init__(self, *transforms: Transform[Any, Any]) -> None:
+        trans = []
+        for t in transforms:
+            if isinstance(t, Transforms):
+                trans.extend(t.transforms)
+            else:
+                trans.append(t)
+        self.transforms = trans
+
+    @override
+    def __call__(self, x: T) -> K:
+        """Apply transforms."""
+        for transform in self.transforms:
+            x = transform(x)
+
+        return x  # type: ignore
+
+    def __getitem__(self, idx: int) -> Transform[Any, Any]:
+        return self.transforms[idx]
+
+    def __len__(self) -> int:
+        return len(self.transforms)
+
+    @override
+    def extra_repr(self) -> str:
+        return ", ".join([str(transform) for transform in self])
+
+
+class Identity(Transform[T, T]):
+    """Resurn input as-is."""
+
+    @override
+    def __call__(self, x: T) -> T:
+        return x

swcgeom/transforms/branch.py

@@ -0,0 +1,223 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Transformation in branch."""
+
+from abc import ABC, abstractmethod
+from typing import cast
+
+import numpy as np
+import numpy.typing as npt
+from scipy import signal
+from typing_extensions import override
+
+from swcgeom.core import Branch, DictSWC
+from swcgeom.transforms.base import Transform
+from swcgeom.utils import (
+    angle,
+    rotate3d_x,
+    rotate3d_y,
+    rotate3d_z,
+    scale3d,
+    to_homogeneous,
+    translate3d,
+)
+
+__all__ = ["BranchLinearResampler", "BranchConvSmoother", "BranchStandardizer"]
+
+
+class _BranchResampler(Transform[Branch, Branch], ABC):
+    r"""Resample branch."""
+
+    @override
+    def __call__(self, x: Branch) -> Branch:
+        xyzr = self.resample(x.xyzr())
+        return Branch.from_xyzr(xyzr)
+
+    @abstractmethod
+    def resample(self, xyzr: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]: ...
+
+
+class BranchLinearResampler(_BranchResampler):
+    r"""Resampling by linear interpolation, DO NOT keep original node."""
+
+    def __init__(self, n_nodes: int) -> None:
+        """Resample branch to special num of nodes.
+
+        Args:
+            n_nodes: Number of nodes after resample.
+        """
+        super().__init__()
+        self.n_nodes = n_nodes
+
+    @override
+    def resample(self, xyzr: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
+        """Resampling by linear interpolation, DO NOT keep original node.
+
+        Args:
+            xyzr: The array of shape (N, 4).
+
+        Returns:
+            coordinates: An array of shape (n_nodes, 4).
+        """
+
+        xp = np.cumsum(np.linalg.norm(xyzr[1:, :3] - xyzr[:-1, :3], axis=1))
+        xp = np.insert(xp, 0, 0)
+        xvals = np.linspace(0, xp[-1], self.n_nodes)
+
+        x = np.interp(xvals, xp, xyzr[:, 0])
+        y = np.interp(xvals, xp, xyzr[:, 1])
+        z = np.interp(xvals, xp, xyzr[:, 2])
+        r = np.interp(xvals, xp, xyzr[:, 3])
+        return cast(npt.NDArray[np.float32], np.stack([x, y, z, r], axis=1))
+
+    @override
+    def extra_repr(self) -> str:
+        return f"n_nodes={self.n_nodes}"
+
+
+class BranchIsometricResampler(_BranchResampler):
+    def __init__(self, distance: float, *, adjust_last_gap: bool = True) -> None:
+        super().__init__()
+        self.distance = distance
+        self.adjust_last_gap = adjust_last_gap
+
+    @override
+    def resample(self, xyzr: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
+        """Resampling by isometric interpolation, DO NOT keep original node.
+
+        Args:
+            xyzr: The array of shape (N, 4).
+
+        Returns:
+            new_xyzr: An array of shape (n_nodes, 4).
+        """
+
+        # Compute the cumulative distances between consecutive points
+        diffs = np.diff(xyzr[:, :3], axis=0)
+        distances = np.sqrt((diffs**2).sum(axis=1))
+        cumulative_distances = np.concatenate([[0], np.cumsum(distances)])
+
+        total_length = cumulative_distances[-1]
+        n_nodes = int(np.ceil(total_length / self.distance)) + 1
+
+        # Determine the new distances
+        if self.adjust_last_gap and n_nodes > 1:
+            new_distances = np.linspace(0, total_length, n_nodes)
+        else:
+            new_distances = np.arange(0, total_length, self.distance)
+            # keep endpoint
+            new_distances = np.concatenate([new_distances, total_length])
+
+        # Interpolate the new points
+        new_xyzr = np.zeros((n_nodes, 4), dtype=np.float32)
+        new_xyzr[:, :3] = np.array(
+            [
+                np.interp(new_distances, cumulative_distances, xyzr[:, i])
+                for i in range(3)
+            ]
+        ).T
+        new_xyzr[:, 3] = np.interp(new_distances, cumulative_distances, xyzr[:, 3])
+        return new_xyzr
+
+    @override
+    def extra_repr(self) -> str:
+        return f"distance={self.distance},adjust_last_gap={self.adjust_last_gap}"
+
+
+class BranchConvSmoother(Transform[Branch, Branch[DictSWC]]):
+    r"""Smooth the branch by sliding window."""
+
+    def __init__(self, n_nodes: int = 5) -> None:
+        """
+        Args:
+            n_nodes: Window size.
+        """
+        super().__init__()
+        self.n_nodes = n_nodes
+        self.kernel = np.ones(n_nodes)
+
+    @override
+    def __call__(self, x: Branch) -> Branch[DictSWC]:
+        x = x.detach()
+        c = signal.convolve(np.ones(x.number_of_nodes()), self.kernel, mode="same")
+        for k in ["x", "y", "z"]:
+            v = x.get_ndata(k)
+            s = signal.convolve(v, self.kernel, mode="same")
+            x.attach.ndata[k][1:-1] = (s / c)[1:-1]
+
+        return x
+
+    def extra_repr(self) -> str:
+        return f"n_nodes={self.n_nodes}"
+
+
+class BranchStandardizer(Transform[Branch, Branch[DictSWC]]):
+    r"""Standardize branch.
+
+    Standardized branch starts at (0, 0, 0), ends at (1, 0, 0), up at y, and scale max
+    radius to 1.
+    """
+
+    @override
+    def __call__(self, x: Branch) -> Branch:
+        xyzr = x.xyzr()
+        xyz, r = xyzr[:, 0:3], xyzr[:, 3:4]
+        T = self.get_matrix(xyz)
+
+        xyz4 = to_homogeneous(xyz, 1).transpose()  # (4, N)
+        new_xyz = np.dot(T, xyz4)[:3].transpose()
+        new_xyzr = np.concatenate([new_xyz, r / r.max()], axis=1)
+        return Branch.from_xyzr(new_xyzr)
+
+    @staticmethod
+    def get_matrix(xyz: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
+        r"""Get standardize transformation matrix.
+
+        Standardized branch starts at (0, 0, 0), ends at (1, 0, 0), up at y.
+
+        Args:
+            xyz: The `x`, `y`, `z` matrix of shape (N, 3) of branch.
+
+        Returns:
+            T: An homogeneous transformation matrix of shape (4, 4).
+        """
+
+        assert xyz.ndim == 2 and xyz.shape[1] == 3, (
+            f"xyz should be of shape (N, 3), got {xyz.shape}"
+        )
+
+        xyz = xyz[:, :3]
+        T = np.identity(4)
+        v = np.concatenate([xyz[-1] - xyz[0], np.zeros((1))])[:, None]
+
+        # translate to the origin
+        T = translate3d(-xyz[0, 0], -xyz[0, 1], -xyz[0, 2]).dot(T)
+
+        # scale to unit vector
+        s = (1 / np.linalg.norm(v[:3, 0])).item()
+        T = scale3d(s, s, s).dot(T)
+
+        # rotate v to the xz-plane, v should be (x, 0, z) now
+        vy = np.dot(T, v)[:, 0]
+        # when looking at the xz-plane along the positive y-axis, the
+        # coordinates should be (z, x)
+        T = rotate3d_y(angle([vy[2], vy[0]], [0, 1])).dot(T)
+
+        # rotate v to the x-axis, v should be (1, 0, 0) now
+        vx = np.dot(T, v)[:, 0]
+        T = rotate3d_z(angle([vx[0], vx[1]], [1, 0])).dot(T)
+
+        # rotate the farthest point to the xy-plane
+        if xyz.shape[0] > 2:
+            xyz4 = to_homogeneous(xyz, 1).transpose()  # (4, N)
+            new_xyz4 = np.dot(T, xyz4)  # (4, N)
+            max_index = np.argmax(np.linalg.norm(new_xyz4[1:3, :], axis=0)[1:-1]) + 1
+            max_xyz4 = xyz4[:, max_index].reshape(4, 1)
+            max_xyz4_t = np.dot(T, max_xyz4)  # (4, 1)
+            angle_x = angle(max_xyz4_t[1:3, 0], [1, 0])
+            T = rotate3d_x(angle_x).dot(T)
+
+        return T

swcgeom/transforms/branch_tree.py

@@ -0,0 +1,74 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Iterable
+
+import numpy as np
+from typing_extensions import override
+
+from swcgeom.core import Branch, BranchTree, Node, Tree
+from swcgeom.transforms.base import Transform
+
+__all__ = ["BranchTreeAssembler"]
+
+
+class BranchTreeAssembler(Transform[BranchTree, Tree]):
+    EPS = 1e-6
+
+    @override
+    def __call__(self, x: BranchTree) -> Tree:
+        nodes = [x.soma().detach()]
+        stack = [(x.soma(), 0)]  # n_orig, id_new
+        while len(stack):
+            n_orig, pid_new = stack.pop()
+            children = n_orig.children()
+
+            for br, c in self.pair(x.branches.get(n_orig.id, []), children):
+                s = 1 if np.linalg.norm(br[0].xyz() - n_orig.xyz()) < self.EPS else 0
+                e = -2 if np.linalg.norm(br[-1].xyz() - c.xyz()) < self.EPS else -1
+
+                br_nodes = [n.detach() for n in br[s:e]] + [c.detach()]
+                for i, n in enumerate(br_nodes):
+                    # reindex
+                    n.id = len(nodes) + i
+                    n.pid = len(nodes) + i - 1
+
+                br_nodes[0].pid = pid_new
+                nodes.extend(br_nodes)
+                stack.append((c, br_nodes[-1].id))
+
+        return Tree(
+            len(nodes),
+            source=x.source,
+            comments=x.comments,
+            names=x.names,
+            **{
+                k: np.array([n.__getattribute__(k) for n in nodes])
+                for k in x.names.cols()
+            },
+        )
+
+    def pair(
+        self, branches: list[Branch], endpoints: list[Node]
+    ) -> Iterable[tuple[Branch, Node]]:
+        assert len(branches) == len(endpoints)
+        xyz1 = [br[-1].xyz() for br in branches]
+        xyz2 = [n.xyz() for n in endpoints]
+        v = np.reshape(xyz1, (-1, 1, 3)) - np.reshape(xyz2, (1, -1, 3))
+        dis = np.linalg.norm(v, axis=-1)
+
+        # greedy algorithm
+        pairs = []
+        for _ in range(len(branches)):
+            # find minimal
+            min_idx = np.argmin(dis)
+            min_branch_idx, min_endpoint_idx = np.unravel_index(min_idx, dis.shape)
+            pairs.append((branches[min_branch_idx], endpoints[min_endpoint_idx]))
+
+            # remove current node
+            dis[min_branch_idx, :] = np.inf
+            dis[:, min_endpoint_idx] = np.inf
+
+        return pairs