swcgeom 0.13.2__py3-none-any.whl → 0.15.0__py3-none-any.whl

swcgeom/images/folder.py

@@ -2,14 +2,25 @@
 
 import os
 import re
+import warnings
 from abc import ABC, abstractmethod
-from typing import Callable, Generic, Iterable, List, Literal, Optional, Tuple, TypeVar
+from typing import (
+    Callable,
+    Generic,
+    Iterable,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    TypeVar,
+    overload,
+)
 
 import numpy as np
 import numpy.typing as npt
 from typing_extensions import Self
 
-from swcgeom.images.io import read_imgs
+from swcgeom.images.io import ScalarType, read_imgs
 from swcgeom.transforms import Identity, Transform
 
 __all__ = [
@@ -21,20 +32,23 @@ __all__ = [
 T = TypeVar("T")
 
 
-class ImageStackFolderBase(Generic[T], ABC):
+class ImageStackFolderBase(Generic[ScalarType, T], ABC):
     """Image stack folder base."""
 
     files: List[str]
-    transform: Transform[npt.NDArray[np.float32], T]
+    transform: Transform[npt.NDArray[ScalarType], T]
 
-    def __init__(
-        self,
-        files: Iterable[str],
-        *,
-        transform: Optional[Transform[npt.NDArray[np.float32], T]] = None,
-    ) -> None:
+    # fmt: off
+    @overload
+    def __init__(self, files: Iterable[str], *, dtype: None = ..., transform: Optional[Transform[npt.NDArray[np.float32], T]] = None) -> None: ...
+    @overload
+    def __init__(self, files: Iterable[str], *, dtype: ScalarType, transform: Optional[Transform[npt.NDArray[ScalarType], T]] = None) -> None: ...
+    # fmt: on
+
+    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
 
     @abstractmethod
@@ -45,13 +59,12 @@ class ImageStackFolderBase(Generic[T], ABC):
         return len(self.files)
 
     def _get(self, fname: str) -> T:
-        imgs = self.read_imgs(fname)
+        imgs = self._read(fname)
         imgs = self.transform(imgs)
         return imgs
 
-    @staticmethod
-    def read_imgs(fname: str) -> npt.NDArray[np.float32]:
-        return read_imgs(fname).get_full()
+    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: Optional[str] = None) -> List[str]:
@@ -63,8 +76,20 @@ class ImageStackFolderBase(Generic[T], ABC):
 
         return fs
 
+    @staticmethod
+    def read_imgs(fname: str) -> npt.NDArray[np.float32]:
+        warnings.warn(
+            "`ImageStackFolderBase.read_imgs` serves as a "
+            "straightforward wrapper for `~swcgeom.images.io.read_imgs(fname).get_full()`. "
+            "However, as it is not utilized within our internal "
+            "processes, it is scheduled for removal in the "
+            "forthcoming version.",
+            DeprecationWarning,
+        )
+        return read_imgs(fname).get_full()
+
 
-class ImageStackFolder(Generic[T], ImageStackFolderBase[T]):
+class ImageStackFolder(ImageStackFolderBase[ScalarType, T]):
     """Image stack folder."""
 
     def __getitem__(self, idx: int, /) -> T:
@@ -84,7 +109,7 @@ class ImageStackFolder(Generic[T], ImageStackFolderBase[T]):
         return cls(cls.scan(root, pattern=pattern), **kwargs)
 
 
-class LabeledImageStackFolder(Generic[T], ImageStackFolderBase[T]):
+class LabeledImageStackFolder(ImageStackFolderBase[ScalarType, T]):
     """Image stack folder with label."""
 
     labels: List[int]
@@ -115,7 +140,7 @@ class LabeledImageStackFolder(Generic[T], ImageStackFolderBase[T]):
         return cls(files, labels, **kwargs)
 
 
-class PathImageStackFolder(Generic[T], ImageStackFolder[T]):
+class PathImageStackFolder(ImageStackFolder[ScalarType, T]):
     """Image stack folder with relpath."""
 
     root: str

swcgeom/images/io.py

@@ -1,6 +1,5 @@
 """Read and write image stack."""
 
-
 import os
 import re
 import warnings
@@ -9,11 +8,13 @@ from functools import cache, lru_cache
 from typing import (
     Any,
     Callable,
+    Generic,
     Iterable,
     List,
     Literal,
     Optional,
     Tuple,
+    TypeVar,
     cast,
     overload,
 )
@@ -27,6 +28,8 @@ from v3dpy.loaders import PBD, Raw
 __all__ = ["read_imgs", "save_tiff", "read_images"]
 
 Vec3i = Tuple[int, int, int]
+ScalarType = TypeVar("ScalarType", bound=np.generic, covariant=True)
+
 RE_TERAFLY_ROOT = re.compile(r"^RES\((\d+)x(\d+)x(\d+)\)$")
 RE_TERAFLY_NAME = re.compile(r"^\d+(_\d+)?(_\d+)?")
 
@@ -46,30 +49,30 @@ AXES_ORDER = {
 }
 
 
-class ImageStack(ABC):
+class ImageStack(ABC, Generic[ScalarType]):
     """Image stack."""
 
     # fmt: off
     @overload
     @abstractmethod
-    def __getitem__(self, key: int) -> npt.NDArray[np.float32]: ...                     # array of shape (Y, Z, C)
+    def __getitem__(self, key: int) -> npt.NDArray[ScalarType]: ...                     # array of shape (Y, Z, C)
     @overload
     @abstractmethod
-    def __getitem__(self, key: Tuple[int, int]) -> npt.NDArray[np.float32]: ...         # array of shape (Z, C)
+    def __getitem__(self, key: Tuple[int, int]) -> npt.NDArray[ScalarType]: ...         # array of shape (Z, C)
     @overload
     @abstractmethod
-    def __getitem__(self, key: Tuple[int, int, int]) -> npt.NDArray[np.float32]: ...    # array of shape (C,)
+    def __getitem__(self, key: Tuple[int, int, int]) -> npt.NDArray[ScalarType]: ...    # array of shape (C,)
     @overload
     @abstractmethod
-    def __getitem__(self, key: Tuple[int, int, int, int]) -> np.float32: ...            # value
+    def __getitem__(self, key: Tuple[int, int, int, int]) -> ScalarType: ...            # value
     @overload
     @abstractmethod
     def __getitem__(
         self, key: slice | Tuple[slice, slice] | Tuple[slice, slice, slice] |  Tuple[slice, slice, slice, slice],
-    ) -> npt.NDArray[np.float32]: ... # array of shape (X, Y, Z, C)
+    ) -> npt.NDArray[ScalarType]: ... # array of shape (X, Y, Z, C)
     @overload
     @abstractmethod
-    def __getitem__(self, key: npt.NDArray[np.integer[Any]]) -> npt.NDArray[np.float32]: ...
+    def __getitem__(self, key: npt.NDArray[np.integer[Any]]) -> npt.NDArray[ScalarType]: ...
     # fmt: on
     @abstractmethod
     def __getitem__(self, key):
@@ -82,7 +85,7 @@ class ImageStack(ABC):
         """
         raise NotImplementedError()
 
-    def get_full(self) -> npt.NDArray[np.float32]:
+    def get_full(self) -> npt.NDArray[ScalarType]:
         """Get full image stack.
 
         Notes
@@ -96,8 +99,19 @@ class ImageStack(ABC):
         raise NotImplementedError()
 
 
-def read_imgs(fname: str, **kwargs) -> ImageStack:
+# fmt:off
+@overload
+def read_imgs(fname: str, *, dtype: ScalarType, **kwargs) -> ImageStack[ScalarType]: ...
+@overload
+def read_imgs(fname: str, *, dtype: None =..., **kwargs) -> ImageStack[np.float32]: ...
+# fmt:on
+
+
+def read_imgs(fname: str, *, dtype=None, **kwargs):  # type: ignore
     """Read image stack."""
+
+    kwargs["dtype"] = dtype or np.float32
+
     ext = os.path.splitext(fname)[-1]
     if ext in [".tif", ".tiff"]:
         return TiffImageStack(fname, **kwargs)
@@ -191,7 +205,7 @@ def save_tiff(
     tifffile.imwrite(fname, data, **kwargs)
 
 
-class NDArrayImageStack(ImageStack):
+class NDArrayImageStack(ImageStack[ScalarType]):
     """NDArray image stack."""
 
     def __init__(
@@ -199,6 +213,8 @@ class NDArrayImageStack(ImageStack):
         imgs: npt.NDArray[Any],
         swap_xy: Optional[bool] = None,
         filp_xy: Optional[bool] = None,
+        *,
+        dtype: ScalarType,
     ) -> None:
         super().__init__()
 
@@ -206,10 +222,6 @@ class NDArrayImageStack(ImageStack):
             imgs = np.expand_dims(imgs, -1)
         assert imgs.ndim == 4, "Should be shape of (X, Y, Z, C)"
 
-        sclar_factor = 1.0
-        if np.issubdtype((dtype := imgs.dtype), np.unsignedinteger):
-            sclar_factor /= UINT_MAX[dtype]
-
         if swap_xy is not None:
             warnings.warn(
                 "flag `swap_xy` now is unnecessary, tifffile will "
@@ -231,12 +243,18 @@ class NDArrayImageStack(ImageStack):
             if filp_xy is True:
                 imgs = np.flip(imgs, (0, 1))  # (X, Y, Z, C)
 
-        self.imgs = imgs.astype(np.float32) * sclar_factor
+        dtype_raw = imgs.dtype
+        self.imgs = imgs.astype(dtype)
+        if np.issubdtype(dtype, np.floating) and np.issubdtype(
+            dtype_raw, np.unsignedinteger
+        ):  # TODO: add a option to disable this
+            sclar_factor = 1.0 / UINT_MAX[imgs.dtype]
+            self.imgs *= sclar_factor
 
     def __getitem__(self, key):
         return self.imgs.__getitem__(key)
 
-    def get_full(self) -> npt.NDArray[np.float32]:
+    def get_full(self) -> npt.NDArray[ScalarType]:
         return self.imgs
 
     @property
@@ -244,7 +262,7 @@ class NDArrayImageStack(ImageStack):
         return cast(Tuple[int, int, int, int], self.imgs.shape)
 
 
-class TiffImageStack(NDArrayImageStack):
+class TiffImageStack(NDArrayImageStack[ScalarType]):
     """Tiff image stack."""
 
     def __init__(
@@ -252,6 +270,8 @@ class TiffImageStack(NDArrayImageStack):
         fname: str,
         swap_xy: Optional[bool] = None,
         filp_xy: Optional[bool] = None,
+        *,
+        dtype: ScalarType,
         **kwargs,
     ) -> None:
         with tifffile.TiffFile(fname, **kwargs) as f:
@@ -265,10 +285,10 @@ class TiffImageStack(NDArrayImageStack):
 
         orders = [AXES_ORDER[c] for c in axes]
         imgs = imgs.transpose(np.argsort(orders))
-        super().__init__(imgs, swap_xy=swap_xy, filp_xy=filp_xy)
+        super().__init__(imgs, swap_xy=swap_xy, filp_xy=filp_xy, dtype=dtype)
 
 
-class NrrdImageStack(NDArrayImageStack):
+class NrrdImageStack(NDArrayImageStack[ScalarType]):
     """Nrrd image stack."""
 
     def __init__(
@@ -276,37 +296,41 @@ class NrrdImageStack(NDArrayImageStack):
         fname: str,
         swap_xy: Optional[bool] = None,
         filp_xy: Optional[bool] = None,
+        *,
+        dtype: ScalarType,
         **kwargs,
     ) -> None:
         imgs, header = nrrd.read(fname, **kwargs)
-        super().__init__(imgs, swap_xy=swap_xy, filp_xy=filp_xy)
+        super().__init__(imgs, swap_xy=swap_xy, filp_xy=filp_xy, dtype=dtype)
         self.header = header
 
 
-class V3dImageStack(NDArrayImageStack):
+class V3dImageStack(NDArrayImageStack[ScalarType]):
     """v3d image stack."""
 
-    def __init__(self, fname: str, loader: Raw | PBD, **kwargs) -> None:
+    def __init__(
+        self, fname: str, loader: Raw | PBD, *, dtype: ScalarType, **kwargs
+    ) -> None:
         r = loader()
         imgs = r.load(fname)
-        super().__init__(imgs, **kwargs)
+        super().__init__(imgs, dtype=dtype, **kwargs)
 
 
-class V3drawImageStack(V3dImageStack):
+class V3drawImageStack(V3dImageStack[ScalarType]):
     """v3draw image stack."""
 
-    def __init__(self, fname: str, **kwargs) -> None:
-        super().__init__(fname, loader=Raw, **kwargs)
+    def __init__(self, fname: str, *, dtype: ScalarType, **kwargs) -> None:
+        super().__init__(fname, loader=Raw, dtype=dtype, **kwargs)
 
 
-class V3dpbdImageStack(V3dImageStack):
+class V3dpbdImageStack(V3dImageStack[ScalarType]):
     """v3dpbd image stack."""
 
-    def __init__(self, fname: str, **kwargs) -> None:
-        super().__init__(fname, loader=PBD, **kwargs)
+    def __init__(self, fname: str, *, dtype: ScalarType, **kwargs) -> None:
+        super().__init__(fname, loader=PBD, dtype=dtype, **kwargs)
 
 
-class TeraflyImageStack(ImageStack):
+class TeraflyImageStack(ImageStack[ScalarType]):
     """TeraFly image stack.
 
     References
@@ -328,21 +352,26 @@ class TeraflyImageStack(ImageStack):
     _listdir: Callable[[str], List[str]]
     _read_patch: Callable[[str], npt.NDArray]
 
-    def __init__(self, root: str, *, lru_maxsize: int | None = 128) -> None:
+    def __init__(
+        self, root: str, *, dtype: ScalarType, lru_maxsize: int | None = 128
+    ) -> None:
         r"""
         Parameters
         ----------
         root : str
             The root of terafly which contains directories named as
             `RES(YxXxZ)`.
+        dtype : np.dtype
         lru_maxsize : int or None, default to 128
             Forwarding to `functools.lru_cache`. A decompressed array
             size of (256, 256, 256, 1), which is the typical size of
             terafly image stack, takes about 256 * 256 * 256 * 1 *
             4B = 64MB. A cache size of 128 requires about 8GB memeory.
         """
+
         super().__init__()
         self.root = root
+        self.dtype = dtype
         self.res, self.res_dirs, self.res_patch_sizes = self.get_resolutions(root)
 
         @cache
@@ -350,8 +379,8 @@ class TeraflyImageStack(ImageStack):
             return os.listdir(path)
 
         @lru_cache(maxsize=lru_maxsize)
-        def read_patch(path: str) -> npt.NDArray[np.float32]:
-            return read_imgs(path).get_full()
+        def read_patch(path: str) -> npt.NDArray[ScalarType]:
+            return read_imgs(path, dtype=dtype).get_full()
 
         self._listdir, self._read_patch = listdir, read_patch
 
@@ -382,7 +411,7 @@ class TeraflyImageStack(ImageStack):
 
     def get_patch(
         self, starts, ends, strides: int | Vec3i = 1, res_level=-1
-    ) -> npt.NDArray[np.float32]:
+    ) -> npt.NDArray[ScalarType]:
         """Get patch of image stack.
 
         Returns
@@ -397,7 +426,7 @@ class TeraflyImageStack(ImageStack):
         assert np.equal(strides, [1, 1, 1]).all()  #  TODO: support stride
 
         shape_out = np.concatenate([ends - starts, [1]])
-        out = np.zeros(shape_out, dtype=np.float32)
+        out = np.zeros(shape_out, dtype=self.dtype)
         self._get_range(starts, ends, res_level, out=out)
 
         # flip y-axis to makes it a left-handed coordinate system

swcgeom/transforms/base.py

@@ -1,7 +1,7 @@
 """Transformation in tree."""
 
 from abc import ABC, abstractmethod
-from typing import Any, Generic, TypeVar, cast, overload
+from typing import Any, Generic, TypeVar, overload
 
 __all__ = ["Transform", "Transforms", "Identity"]
 
@@ -15,26 +15,49 @@ class Transform(ABC, Generic[T, K]):
     r"""An abstract class representing a :class:`Transform`.
 
     All transforms that represent a map from `T` to `K`.
-
-    Methods
-    -------
-    __call__(x: T) -> K
-        All subclasses should overwrite :meth:`__call__`, supporting
-        applying transform in `x`.
-    __repr__() -> str
-        Subclasses could also optionally overwrite :meth:`__repr__`.
-        Avoid using the underscore `_` because it is used by
-        `Transforms`. If not provied, class name will be a default
-        value.
     """
 
     @abstractmethod
     def __call__(self, x: T) -> K:
-        """Apply transform."""
+        """Apply transform.
+
+        Notes
+        -----
+        All subclasses should overwrite :meth:`__call__`, supporting
+        applying transform in `x`.
+        """
         raise NotImplementedError()
 
     def __repr__(self) -> str:
-        return self.__class__.__name__
+        classname = self.__class__.__name__
+        repr_ = self.extra_repr()
+        return f"{classname}({repr_})"
+
+    def extra_repr(self):
+        """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.
+
+        Examples
+        --------
+        class Foo(Transform[T, K]):
+            def __init__(self, my_parameter: int = 1):
+                self.my_parameter = my_parameter
+
+            def extra_repr(self):
+                return f"my_parameter={self.my_parameter}"
+
+        Notes
+        -----
+        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]):
@@ -81,23 +104,20 @@ class Transforms(Transform[T, K]):
         for transform in self.transforms:
             x = transform(x)
 
-        return cast(K, x)
+        return x  # type: ignore
 
     def __getitem__(self, idx: int) -> Transform[Any, Any]:
         return self.transforms[idx]
 
-    def __repr__(self) -> str:
-        return "_".join([str(transform) for transform in self])
-
     def __len__(self) -> int:
         return len(self.transforms)
 
+    def extra_repr(self) -> str:
+        return ", ".join([str(transform) for transform in self])
+
 
 class Identity(Transform[T, T]):
     """Resurn input as-is."""
 
     def __call__(self, x: T) -> T:
         return x
-
-    def __repr__(self) -> str:
-        return ""

swcgeom/transforms/branch.py

@@ -48,9 +48,6 @@ class BranchLinearResampler(_BranchResampler):
         super().__init__()
         self.n_nodes = n_nodes
 
-    def __repr__(self) -> str:
-        return f"BranchLinearResampler-{self.n_nodes}"
-
     def resample(self, xyzr: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
         """Resampling by linear interpolation, DO NOT keep original node.
 
@@ -75,6 +72,9 @@ class BranchLinearResampler(_BranchResampler):
         r = np.interp(xvals, xp, xyzr[:, 3])
         return cast(npt.NDArray[np.float32], np.stack([x, y, z, r], axis=1))
 
+    def extra_repr(self):
+        return f"n_nodes={self.n_nodes}"
+
 
 class BranchConvSmoother(Transform[Branch, Branch[DictSWC]]):
     r"""Smooth the branch by sliding window."""
@@ -100,8 +100,8 @@ class BranchConvSmoother(Transform[Branch, Branch[DictSWC]]):
 
         return x
 
-    def __repr__(self) -> str:
-        return f"BranchConvSmoother-{self.n_nodes}"
+    def extra_repr(self):
+        return f"n_nodes={self.n_nodes}"
 
 
 class BranchStandardizer(Transform[Branch, Branch[DictSWC]]):

swcgeom/transforms/geometry.py

@@ -72,6 +72,9 @@ class RadiusReseter(Generic[T], Transform[T, T]):
         new_tree.ndata[new_tree.names.r] = r
         return new_tree
 
+    def extra_repr(self):
+        return f"r={self.r:.4f}"
+
 
 class AffineTransform(Generic[T], Transform[T, T]):
     """Apply affine matrix."""
@@ -85,10 +88,18 @@ class AffineTransform(Generic[T], Transform[T, T]):
         tm: npt.NDArray[np.float32],
         center: Center = "origin",
         *,
-        fmt: str,
+        fmt: Optional[str] = None,
         names: Optional[SWCNames] = None,
     ) -> None:
-        self.tm, self.center, self.fmt = tm, center, fmt
+        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 "
@@ -111,9 +122,6 @@ class AffineTransform(Generic[T], Transform[T, T]):
 
         return self.apply(x, tm)
 
-    def __repr__(self) -> str:
-        return self.fmt
-
     @staticmethod
     def apply(x: T, tm: npt.NDArray[np.float32]) -> T:
         xyzw = x.xyzw().dot(tm.T).T
@@ -130,8 +138,11 @@ class Translate(Generic[T], AffineTransform[T]):
     """Translate SWC."""
 
     def __init__(self, tx: float, ty: float, tz: float, **kwargs) -> None:
-        fmt = f"Translate-{tx}-{ty}-{tz}"
-        super().__init__(translate3d(tx, ty, tz), fmt=fmt, **kwargs)
+        super().__init__(translate3d(tx, ty, tz), **kwargs)
+        self.tx, self.ty, self.tz = tx, ty, tz
+
+    def extra_repr(self):
+        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:
@@ -158,8 +169,7 @@ class Scale(Generic[T], AffineTransform[T]):
     def __init__(
         self, sx: float, sy: float, sz: float, center: Center = "root", **kwargs
     ) -> None:
-        fmt = f"Scale-{sx}-{sy}-{sz}"
-        super().__init__(scale3d(sx, sy, sz), center=center, fmt=fmt, **kwargs)
+        super().__init__(scale3d(sx, sy, sz), center=center, **kwargs)
 
     @classmethod
     def transform(  # pylint: disable=too-many-arguments
@@ -180,6 +190,12 @@ class Rotate(Generic[T], AffineTransform[T]):
     ) -> 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
+
+    def extra_repr(self):
+        return f"n={self.n}, theta={self.theta:.4f}, center={self.center}"  # TODO: imporve format of n
 
     @classmethod
     def transform(
@@ -197,9 +213,11 @@ 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, fmt=f"RotateX-{theta}", **kwargs
-        )
+        super().__init__(rotate3d_x(theta), center=center, **kwargs)
+        self.theta = theta
+
+    def extra_repr(self):
+        return f"center={self.center}, theta={self.theta:.4f}"
 
     @classmethod
     def transform(cls, x: T, theta: float, center: Center = "root", **kwargs) -> T:
@@ -210,9 +228,12 @@ 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, fmt=f"RotateX-{theta}", **kwargs
-        )
+        super().__init__(rotate3d_y(theta), center=center, **kwargs)
+        self.theta = theta
+        self.center = center
+
+    def extra_repr(self):
+        return f"theta={self.theta:.4f}, center={self.center}"
 
     @classmethod
     def transform(cls, x: T, theta: float, center: Center = "root", **kwargs) -> T:
@@ -223,9 +244,12 @@ 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, fmt=f"RotateX-{theta}", **kwargs
-        )
+        super().__init__(rotate3d_z(theta), center=center, **kwargs)
+        self.theta = theta
+        self.center = center
+
+    def extra_repr(self):
+        return f"theta={self.theta:.4f}, center={self.center}"
 
     @classmethod
     def transform(cls, x: T, theta: float, center: Center = "root", **kwargs) -> T: