Segmentation Plugin Example

This example demonstrates how to use the segmentation plugin system in ZarrNii. The plugin system is built on the pluggy framework, providing a flexible and extensible architecture for implementing custom segmentation algorithms.

Plugin Architecture Overview

ZarrNii's segmentation plugins use pluggy hook specifications and implementations: - Hook specifications are defined in zarrnii.plugins (the core zarrnii package) - Hook implementations (using @hookimpl decorator) provide the actual functionality - Plugins are plain Python classes — no inheritance from any base class is required - Plugins can be used directly or registered with the plugin manager for discovery

Basic Otsu Segmentation

import numpy as np
import dask.array as da
from zarrnii import ZarrNii
from zarrnii.plugins import LocalOtsuSegmentation

# Load or create your image data
# For this example, we'll create synthetic bimodal data
np.random.seed(42)
image_data = np.random.normal(0.2, 0.05, (1, 50, 100, 100))  # Background
image_data[0, 20:30, 40:60, 40:60] = np.random.normal(0.8, 0.05, (10, 20, 20))  # Foreground

# Create ZarrNii instance
darr = da.from_array(image_data, chunks=(1, 25, 50, 50))
znimg = ZarrNii.from_darr(darr, axes_order="ZYX", orientation="RAS")

# Method 1: Using the convenience method
segmented = znimg.segment_otsu(nbins=256)

# Method 2: Using the plugin directly
plugin = LocalOtsuSegmentation(nbins=256)
segmented = znimg.segment(plugin)

# Method 3: Using plugin class with parameters
segmented = znimg.segment(LocalOtsuSegmentation, nbins=128)

# The result is a new ZarrNii instance with binary segmentation
print(f"Original shape: {znimg.shape}")
print(f"Segmented shape: {segmented.shape}")
print(f"Segmented dtype: {segmented.data.dtype}")
print(f"Unique values: {np.unique(segmented.data.compute())}")

# Save segmented result as OME-Zarr
segmented.to_ome_zarr("segmented_image.ome.zarr")

Custom Chunk Processing

For large datasets, you can control the chunk size for blockwise processing:

# Segment with custom chunk size for memory efficiency
custom_chunks = (1, 10, 25, 25)
segmented = znimg.segment_otsu(chunk_size=custom_chunks)

# The segmentation will be applied block-wise using dask
result_data = segmented.data.compute()

Creating Custom Segmentation Plugins

You can create your own segmentation plugins as plain Python classes using the @hookimpl decorator from zarrnii.plugins. No inheritance is required.

from zarrnii.plugins import hookimpl
import numpy as np

class ThresholdSegmentation:
    """Simple threshold-based segmentation plugin."""

    def __init__(self, threshold: float = 0.5):
        """Initialize the plugin.

        Args:
            threshold: Threshold value for segmentation
        """
        self.threshold = threshold

    @hookimpl
    def segment(self, image: np.ndarray, metadata=None) -> np.ndarray:
        """Apply threshold segmentation.

        Args:
            image: Input image as numpy array
            metadata: Optional metadata dictionary

        Returns:
            Binary segmentation mask
        """
        binary_mask = image > self.threshold
        return binary_mask.astype(np.uint8)

    @hookimpl
    def segmentation_plugin_name(self) -> str:
        """Return the name of the plugin."""
        return "Threshold Segmentation"

    @hookimpl
    def segmentation_plugin_description(self) -> str:
        """Return a description of the plugin."""
        return f"Simple thresholding at value {self.threshold}"

# Method 1: Direct usage with ZarrNii (recommended for simple cases)
custom_plugin = ThresholdSegmentation(threshold=0.3)
segmented = znimg.segment(custom_plugin)

# Method 2: Using the plugin manager (recommended for external plugins)
from zarrnii.plugins import get_plugin_manager

pm = get_plugin_manager()
plugin = ThresholdSegmentation(threshold=0.3)
pm.register(plugin)

# Now the plugin is available through the plugin manager
# and can be discovered by other tools
registered_plugins = pm.get_plugins()
print(f"Registered {len(registered_plugins)} plugins")

External Segmentation Plugin Development

External plugins allow you to package and distribute your custom segmentation algorithms as separate Python packages. They only need zarrnii as a dependency.

Step 1: Create Your Plugin Package Structure

my_segmentation_plugin/
├── pyproject.toml
└── my_segmentation_plugin/
    ├── __init__.py
    └── adaptive_threshold.py

Step 2: Implement Your Segmentation Plugin

In adaptive_threshold.py:

"""Adaptive threshold segmentation plugin."""
from zarrnii.plugins import hookimpl
import numpy as np
from skimage.filters import threshold_local

class AdaptiveThresholdSegmentation:
    """Adaptive thresholding segmentation for images with varying illumination."""

    def __init__(self, block_size=35, offset=10, method='gaussian'):
        """Initialize adaptive threshold segmentation.

        Args:
            block_size: Size of pixel neighborhood for threshold calculation
            offset: Constant subtracted from weighted mean
            method: Method for computing threshold ('gaussian' or 'mean')
        """
        self.block_size = block_size
        self.offset = offset
        self.method = method

    @hookimpl
    def segment(self, image: np.ndarray, metadata=None) -> np.ndarray:
        """Segment image using adaptive thresholding.

        Args:
            image: Input image as numpy array
            metadata: Optional metadata (unused)

        Returns:
            Binary segmentation mask
        """
        if image.size == 0:
            raise ValueError("Input image is empty")

        if image.ndim < 2:
            raise ValueError("Input image must be at least 2D")

        # Work with 2D slice for threshold computation
        work_image = image
        if work_image.ndim > 2:
            # Use first channel/slice if multi-dimensional
            if work_image.ndim == 3 and work_image.shape[0] <= 4:
                work_image = work_image[0]
            elif work_image.ndim > 3:
                work_image = work_image.reshape(-1, *work_image.shape[-2:])[0]

        # Compute adaptive threshold
        threshold = threshold_local(
            work_image,
            block_size=self.block_size,
            offset=self.offset,
            method=self.method
        )

        # Apply threshold to original image
        binary_mask = image > threshold

        return binary_mask.astype(np.uint8)

    @hookimpl
    def segmentation_plugin_name(self) -> str:
        """Return the name of the plugin."""
        return "Adaptive Threshold Segmentation"

    @hookimpl
    def segmentation_plugin_description(self) -> str:
        """Return a description of the plugin."""
        return (
            f"Adaptive thresholding using {self.method} method "
            f"(block_size={self.block_size}, offset={self.offset})"
        )

Step 3: Configure Package Discovery

In __init__.py:

"""My Segmentation Plugin Package."""
from .adaptive_threshold import AdaptiveThresholdSegmentation

__all__ = ["AdaptiveThresholdSegmentation"]

In pyproject.toml:

[project]
name = "my-segmentation-plugin"
version = "0.1.0"
description = "Adaptive threshold segmentation plugin for ZarrNii"
dependencies = [
    "zarrnii>=0.1.0",
    "scikit-image>=0.21.0",
]

[project.entry-points."zarrnii"]
adaptive_threshold = "my_segmentation_plugin:AdaptiveThresholdSegmentation"

Step 4: Use Your External Plugin

After installing (pip install my-segmentation-plugin):

from zarrnii import ZarrNii
from my_segmentation_plugin import AdaptiveThresholdSegmentation

# Load your data
znimg = ZarrNii.from_ome_zarr("input.ome.zarr")

# Use the plugin directly
segmented = znimg.segment(
    AdaptiveThresholdSegmentation(block_size=51, offset=5)
)

# Or register with the plugin manager for discovery
from zarrnii.plugins import get_plugin_manager

pm = get_plugin_manager()
plugin = AdaptiveThresholdSegmentation(block_size=51, offset=5)
pm.register(plugin)

# Call hooks through the plugin manager
test_image = znimg.data.compute()
results = pm.hook.segment(image=test_image)

# Get plugin information
names = pm.hook.segmentation_plugin_name()
descriptions = pm.hook.segmentation_plugin_description()
print(f"Using: {names[0]} - {descriptions[0]}")

Working with Multi-channel Images

The segmentation plugins automatically handle multi-channel images:

# Create multi-channel test data
multichannel_data = np.random.rand(3, 50, 100, 100)  # 3 channels
multichannel_data[0, 20:30, 40:60, 40:60] += 0.5  # Add signal to first channel

darr = da.from_array(multichannel_data, chunks=(1, 25, 50, 50))
znimg = ZarrNii.from_darr(darr, axes_order="ZYX", orientation="RAS")

# Segment - will use first channel for threshold calculation
# but preserve all channel dimensions in output
segmented = znimg.segment_otsu()
print(f"Input shape: {znimg.shape}")       # (3, 50, 100, 100)
print(f"Output shape: {segmented.shape}")   # (3, 50, 100, 100)

Integration with Existing Workflows

The segmentation plugins integrate seamlessly with other ZarrNii operations:

# Complete workflow: load, downsample, segment, save
znimg = ZarrNii.from_ome_zarr("input_image.ome.zarr", level=1)

# Downsample for faster processing
downsampled = znimg.downsample(factors=2, spatial_dims=["z", "y", "x"])

# Apply segmentation
segmented = downsampled.segment_otsu(nbins=128)

# Crop to region of interest
bbox_min = (10, 20, 20)
bbox_max = (40, 80, 80)
cropped = segmented.crop_with_bounding_box(bbox_min, bbox_max)

# Save final result
cropped.to_ome_zarr("processed_segmentation.ome.zarr")