swcgeom 0.14.0__py3-none-any.whl → 0.16.0__py3-none-any.whl

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
@@ -37,7 +38,7 @@ class ToImageStack(Transform[Tree, npt.NDArray[np.uint8]]):
 
     resolution: npt.NDArray[np.float32]
 
-    def __init__(self, resolution: float | npt.ArrayLike = 1) -> None:
+    def __init__(self, resolution: int | float | npt.ArrayLike = 1) -> None:
         """Transform tree to image stack.
 
         Parameters
@@ -46,11 +47,11 @@ class ToImageStack(Transform[Tree, npt.NDArray[np.uint8]]):
             Resolution of image stack.
         """
 
-        if isinstance(resolution, float):
-            resolution = [resolution, resolution, resolution]
+        if isinstance(resolution, (int, float, np.integer, np.floating)):
+            resolution = [resolution, resolution, resolution]  # type: ignore
 
-        self.resolution = (resolution := np.array(resolution, dtype=np.float32))
-        assert tuple(resolution.shape) == (3,), "resolution shoule be vector of 3d."
+        self.resolution = np.array(resolution, dtype=np.float32)
+        assert len(self.resolution) == 3, "resolution shoule be vector of 3d."
 
     def __call__(self, x: Tree) -> npt.NDArray[np.uint8]:
         """Transform tree to image stack.
@@ -58,57 +59,74 @@ class ToImageStack(Transform[Tree, npt.NDArray[np.uint8]]):
         Notes
         -----
         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.
+        ONLY works for small image stacks, use :meth`transform_and_save`
+        for big image stack.
         """
-        return np.concatenate(list(self.transfrom(x, verbose=False)), axis=0)
-
-    def __repr__(self) -> str:
-        return "ToImageStack" + f"-resolution-{'-'.join(self.resolution)}"
+        return np.stack(list(self.transfrom(x, verbose=False)), axis=0)
 
     def transfrom(
-        self, x: Tree, verbose: bool = True
+        self,
+        x: Tree,
+        verbose: bool = True,
+        *,
+        ranges: Optional[Tuple[npt.ArrayLike, npt.ArrayLike]] = None,
     ) -> Iterable[npt.NDArray[np.uint8]]:
-        from tqdm import tqdm
-
         if verbose:
             print("To image stack: " + x.source)
             time_start = time.time()
 
         scene = self._get_scene(x)
 
-        xyz, r = x.xyz(), x.r().reshape(-1, 1)
-        coord_min = np.floor(np.min(xyz - r, axis=0))
-        coord_max = np.ceil(np.max(xyz + r, axis=0))
+        if ranges is None:
+            xyz, r = x.xyz(), x.r().reshape(-1, 1)
+            coord_min = np.floor(np.min(xyz - r, axis=0))
+            coord_max = np.ceil(np.max(xyz + 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)
-        total = (
-            ((coord_max[2] - coord_min[2]) / self.resolution[2]).astype(np.int64).item()
-        )
 
         if verbose:
+            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")  # type: ignore
 
-        for sampler in tqdm(samplers, total=total) if verbose else samplers:
-            voxel = sampler.sample(scene)  # shoule be shape of (x, y, z, 3) and z = 1
+        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
 
-    def transform_and_save(self, fname: str, x: Tree, verbose: bool = True) -> None:
-        self.save_tif(fname, self.transfrom(x, verbose=verbose))
+    def transform_and_save(
+        self, fname: str, x: Tree, verbose: bool = True, **kwargs
+    ) -> None:
+        self.save_tif(fname, self.transfrom(x, verbose=verbose, **kwargs))
 
     def transform_population(
         self, population: Population | str, verbose: bool = True
     ) -> None:
-        if isinstance(population, str):
-            population = Population.from_swc(population)
+        trees = (
+            Population.from_swc(population)
+            if isinstance(population, str)
+            else population
+        )
+
+        if verbose:
+            trees = tqdm(trees)
 
         # TODO: multiprocess
-        for tree in population:
+        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=verbose)
+                self.transform_and_save(tif, tree, verbose=False)
+
+    def extra_repr(self):
+        res = ",".join(f"{a:.4f}" for a in self.resolution)
+        return f"resolution=({res})"
 
     def _get_scene(self, x: Tree) -> Scene:
         material = ColoredMaterial((1, 0, 0)).into()

swcgeom/transforms/images.py

@@ -1,6 +1,6 @@
 """Image stack related transform."""
 
-
+import warnings
 from typing import Tuple
 
 import numpy as np
@@ -8,10 +8,20 @@ import numpy.typing as npt
 
 from swcgeom.transforms.base import Transform
 
-__all__ = ["Center"]
+__all__ = [
+    "ImagesCenterCrop",
+    "ImagesScale",
+    "ImagesClip",
+    "ImagesNormalizer",
+    "ImagesMeanVarianceAdjustment",
+    "Center",  # legacy
+]
+
 
+NDArrayf32 = npt.NDArray[np.float32]
 
-class Center(Transform[npt.NDArray[np.float32], npt.NDArray[np.float32]]):
+
+class ImagesCenterCrop(Transform[NDArrayf32, NDArrayf32]):
     """Get image stack center."""
 
     def __init__(self, shape_out: int | Tuple[int, int, int]):
@@ -22,11 +32,71 @@ class Center(Transform[npt.NDArray[np.float32], npt.NDArray[np.float32]]):
             else (shape_out, shape_out, shape_out)
         )
 
-    def __call__(self, x: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
+    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], :]
 
-    def __repr__(self) -> str:
-        return f"Center-{self.shape_out[0]}-{self.shape_out[1]}-{self.shape_out[2]}"
+    def extra_repr(self) -> str:
+        return f"shape_out=({','.join(str(a) for a in self.shape_out)})"
+
+
+class Center(ImagesCenterCrop):
+    """Get image stack center.
+
+    .. deprecated:: 0.5.0
+        Use :class:`ImagesCenterCrop` instead.
+    """
+
+    def __init__(self, shape_out: int | Tuple[int, int, int]):
+        warnings.warn(
+            "`Center` is deprecated, use `ImagesCenterCrop` instead",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        super().__init__(shape_out)
+
+
+class ImagesScale(Transform[NDArrayf32, NDArrayf32]):
+    def __init__(self, scaler: float) -> None:
+        super().__init__()
+        self.scaler = scaler
+
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        return self.scaler * x
+
+
+class ImagesClip(Transform[NDArrayf32, NDArrayf32]):
+    def __init__(self, vmin: float = 0, vmax: float = 1, /) -> None:
+        super().__init__()
+        self.vmin, self.vmax = vmin, vmax
+
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        return np.clip(x, self.vmin, self.vmax)
+
+
+class ImagesNormalizer(Transform[NDArrayf32, NDArrayf32]):
+    """Normalize image stack."""
+
+    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
+
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        return (x - self.mean) / self.variance

swcgeom/transforms/mst.py

@@ -23,11 +23,11 @@ class PointsToCuntzMST(Transform[npt.NDArray[np.float32], Tree]):
 
     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).
+    .. [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__(
@@ -140,13 +140,8 @@ class PointsToCuntzMST(Transform[npt.NDArray[np.float32], Tree]):
             t = sort_tree(t)
         return t
 
-    def __repr__(self) -> str:
-        return (
-            f"PointsToCuntzMST"
-            f"-bf-{self.bf}"
-            f"-furcations-{self.furcations}"
-            f"-{'exclude-soma' if self.exclude_soma else 'include-soma'}"
-        )  # TODO: names, types
+    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
@@ -173,6 +168,7 @@ class PointsToMST(PointsToCuntzMST):  # pylint: disable=too-few-public-methods
         names : SWCNames, optional
         types : SWCTypes, optional
         """
+
         if k_furcations is not None:
             warnings.warn(
                 "`PointsToMST(k_furcations=...)` has been renamed to "
@@ -191,9 +187,5 @@ class PointsToMST(PointsToCuntzMST):  # pylint: disable=too-few-public-methods
             **kwargs,
         )
 
-    def __repr__(self) -> str:
-        return (
-            f"PointsToMST"
-            f"-furcations-{self.furcations}"
-            f"-{'exclude-soma' if self.exclude_soma else 'include-soma'}"
-        )
+    def extra_repr(self) -> str:
+        return f"furcations-{self.furcations}, exclude-soma={self.exclude_soma}"