Multi-Resolution Plugin Architecture

This section covers the multi-resolution plugin architecture that enables efficient processing where algorithms are computed at low resolution and applied to full resolution data.

Overview

The scaled processing plugin architecture in ZarrNii allows for efficient multi-resolution operations. This is particularly useful for algorithms that can be computed efficiently at lower resolution and then applied to the full-resolution data. Common use cases include:

Bias field correction
Background estimation
Denoising operations
Global intensity normalization

ZarrNii uses the pluggy framework for its plugin system, providing a flexible and extensible architecture for creating custom processing plugins.

Plugin Interface

All scaled processing plugins must inherit from ScaledProcessingPlugin and implement the required hook methods decorated with @hookimpl:

`lowres_func(lowres_array: np.ndarray) -> np.ndarray`

This function processes the downsampled data and returns a result that will be upsampled and applied to the full-resolution data.

`highres_func(fullres_array: dask.array, upsampled_output: dask.array) -> dask.array`

This function receives the full-resolution dask array and the upsampled output (already upsampled to match the full-resolution shape), and applies the operation blockwise. The upsampling is handled internally by the apply_scaled_processing method.

`scaled_processing_plugin_name() -> str`

Returns the name of the plugin.

`scaled_processing_plugin_description() -> str`

Returns a description of what the plugin does.

Basic Usage

from zarrnii import ZarrNii, GaussianBiasFieldCorrection

# Load your data
znimg = ZarrNii.from_nifti("path/to/image.nii")

# Apply bias field correction
corrected = znimg.apply_scaled_processing(
    GaussianBiasFieldCorrection(sigma=5.0),
    downsample_factor=4
)

# Save result
corrected.to_nifti("corrected_image.nii")

Built-in Plugins

GaussianBiasFieldCorrection

A simple bias field correction plugin that estimates smooth bias fields using Gaussian smoothing at low resolution and applies correction by division.

# Basic usage with default parameters
corrected = znimg.apply_scaled_processing(GaussianBiasFieldCorrection())

# Custom parameters
corrected = znimg.apply_scaled_processing(
    GaussianBiasFieldCorrection(sigma=3.0, mode='constant'),
    downsample_factor=8
)

Parameters: - sigma: Standard deviation for Gaussian smoothing (default: 5.0) - mode: Boundary condition for smoothing (default: 'reflect')

N4BiasFieldCorrection

A more sophisticated bias field correction plugin that uses the N4 algorithm from ANTsPy for superior bias field estimation at low resolution, then applies correction by division.

Installation: Requires antspyx to be installed:

# Install with N4 support
pip install 'zarrnii[n4]'
# or install antspyx directly
pip install antspyx

from zarrnii import N4BiasFieldCorrection

# Basic usage with default parameters
corrected = znimg.apply_scaled_processing(N4BiasFieldCorrection())

# Custom parameters for more control
corrected = znimg.apply_scaled_processing(
    N4BiasFieldCorrection(
        spline_spacing=150.0,
        convergence={'iters': [25, 25], 'tol': 0.001},
        shrink_factor=2
    ),
    downsample_factor=4
)

Parameters: - spline_spacing: Spacing between knots for spline fitting (default: 200.0) - convergence: Dictionary with 'iters' (list) and 'tol' (float) for convergence criteria (default: {'iters': [50], 'tol': 0.001}) - shrink_factor: Shrink factor for processing (default: 1)

Creating Custom Plugins

You can create custom plugins by inheriting from ScaledProcessingPlugin and implementing the required hook methods:

from zarrnii.plugins import ScaledProcessingPlugin
from zarrnii.plugins.scaled_processing.base import hookimpl
import numpy as np
import dask.array as da
from scipy import ndimage

class CustomPlugin(ScaledProcessingPlugin):
    """Custom scaled processing plugin example."""

    def __init__(self, param1=1.0, **kwargs):
        super().__init__(param1=param1, **kwargs)
        self.param1 = param1

    @hookimpl
    def lowres_func(self, lowres_array: np.ndarray) -> np.ndarray:
        """Process low-resolution data."""
        # Your low-resolution algorithm here
        # Example: compute some correction map
        correction_map = np.ones_like(lowres_array) * self.param1
        return correction_map

    @hookimpl
    def highres_func(self, fullres_array: da.Array, upsampled_output: da.Array) -> da.Array:
        """Apply correction to full-resolution data."""
        # The upsampling is handled internally by apply_scaled_processing
        # This example shows a simple multiplication operation
        # Apply correction directly (both arrays are same size)
        result = fullres_array * upsampled_output
        return result

    @hookimpl
    def scaled_processing_plugin_name(self) -> str:
        """Return the name of the plugin."""
        return "Custom Plugin"

    @hookimpl
    def scaled_processing_plugin_description(self) -> str:
        """Return a description of the plugin."""
        return "A custom multi-resolution processing plugin"

# Method 1: Direct usage with ZarrNii (recommended for simple cases)
result = znimg.apply_scaled_processing(CustomPlugin(param1=2.0))

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

pm = get_plugin_manager()
plugin = CustomPlugin(param1=2.0)
pm.register(plugin)

# Now the plugin is available through the plugin manager
# and can be discovered by other tools

External Plugin Development

External plugins allow you to package and distribute your custom plugins as separate Python packages that can be discovered and used by ZarrNii. Here's a complete example of how to create an external plugin:

Step 1: Create Your Plugin Package

Create a new Python package with the following structure:

my_zarrnii_plugin/
├── setup.py or pyproject.toml
└── my_zarrnii_plugin/
    ├── __init__.py
    └── my_plugin.py

Step 2: Implement Your Plugin

In my_plugin.py:

"""Custom external plugin for ZarrNii."""
from zarrnii.plugins import ScaledProcessingPlugin
from zarrnii.plugins.scaled_processing.base import hookimpl
import numpy as np
import dask.array as da
from scipy import ndimage

class MyExternalPlugin(ScaledProcessingPlugin):
    """An example external plugin for demonstration."""

    def __init__(self, smoothing_sigma=2.0, scale_factor=1.5, **kwargs):
        """Initialize the plugin with custom parameters.

        Args:
            smoothing_sigma: Sigma for Gaussian smoothing
            scale_factor: Scaling factor for the correction
        """
        super().__init__(
            smoothing_sigma=smoothing_sigma,
            scale_factor=scale_factor,
            **kwargs
        )
        self.smoothing_sigma = smoothing_sigma
        self.scale_factor = scale_factor

    @hookimpl
    def lowres_func(self, lowres_array: np.ndarray) -> np.ndarray:
        """Compute correction map at low resolution."""
        if lowres_array.size == 0:
            raise ValueError("Input array is empty")

        # Apply smoothing to estimate low-frequency components
        smoothed = ndimage.gaussian_filter(
            lowres_array.astype(np.float32),
            sigma=self.smoothing_sigma
        )

        # Scale the correction map
        correction_map = smoothed * self.scale_factor

        return correction_map

    @hookimpl
    def highres_func(self, fullres_array: da.Array, upsampled_output: da.Array) -> da.Array:
        """Apply the correction to full-resolution data."""
        # Apply the correction by division
        epsilon = np.finfo(np.float32).eps
        corrected = fullres_array / da.maximum(upsampled_output, epsilon)
        return corrected

    @hookimpl
    def scaled_processing_plugin_name(self) -> str:
        """Return plugin name."""
        return "My External Plugin"

    @hookimpl
    def scaled_processing_plugin_description(self) -> str:
        """Return plugin description."""
        return (
            f"External plugin with smoothing (sigma={self.smoothing_sigma}) "
            f"and scaling (factor={self.scale_factor})"
        )

Step 3: Make Your Plugin Discoverable

In __init__.py:

"""My ZarrNii Plugin Package."""
from .my_plugin import MyExternalPlugin

__all__ = ["MyExternalPlugin"]

Step 4: Configure Your Package

In pyproject.toml:

[project]
name = "my-zarrnii-plugin"
version = "0.1.0"
description = "My custom ZarrNii processing plugin"
dependencies = [
    "zarrnii>=0.1.0",
    "scipy>=1.11.0",
]

[project.entry-points."zarrnii.plugins"]
my_external_plugin = "my_zarrnii_plugin:MyExternalPlugin"

Step 5: Use Your External Plugin

After installing your plugin package (pip install my-zarrnii-plugin):

from zarrnii import ZarrNii
from my_zarrnii_plugin import MyExternalPlugin

# Load your data
znimg = ZarrNii.from_nifti("input.nii")

# Use the external plugin directly
result = znimg.apply_scaled_processing(
    MyExternalPlugin(smoothing_sigma=3.0, scale_factor=2.0),
    downsample_factor=4
)

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

pm = get_plugin_manager()
plugin = MyExternalPlugin(smoothing_sigma=3.0, scale_factor=2.0)
pm.register(plugin)

# Query registered plugins
for p in pm.get_plugins():
    names = pm.hook.scaled_processing_plugin_name()
    descriptions = pm.hook.scaled_processing_plugin_description()
    print(f"Plugin: {names[0]} - {descriptions[0]}")

Advanced Usage

Custom Downsampling Factors

# Use different downsampling factors
result = znimg.apply_scaled_processing(
    GaussianBiasFieldCorrection(),
    downsample_factor=8  # 8x downsampling
)

Custom Chunk Sizes

# Specify custom chunk sizes for low-resolution processing
# The chunk_size parameter controls the chunking of low-resolution intermediate results
result = znimg.apply_scaled_processing(
    GaussianBiasFieldCorrection(),
    chunk_size=(1, 32, 32, 32)  # Used for low-res processing chunks
)

Temporary File Options

The framework uses temporary OME-Zarr files to break up the dask computation graph for better performance. You can control this behavior:

# Disable temporary file usage (may impact performance on large datasets)
result = znimg.apply_scaled_processing(
    GaussianBiasFieldCorrection(),
    use_temp_zarr=False
)

# Use custom temporary file location
result = znimg.apply_scaled_processing(
    GaussianBiasFieldCorrection(),
    temp_zarr_path="/custom/path/temp_processing.ome.zarr"
)

Plugin Class vs Instance

# Using plugin class (parameters passed as kwargs)
result1 = znimg.apply_scaled_processing(GaussianBiasFieldCorrection, sigma=3.0)

# Using plugin instance (parameters set during initialization)
plugin = GaussianBiasFieldCorrection(sigma=3.0)
result2 = znimg.apply_scaled_processing(plugin)

Performance Considerations

Downsampling Factor: Higher factors reduce computation time but may reduce accuracy
Chunk Sizes: Optimize for your memory constraints and processing requirements
Algorithm Complexity: The lowres_func runs on small numpy arrays, while highres_func uses dask for scalability
Temporary Files: The default temporary OME-Zarr approach breaks up dask computation graphs for better performance on large datasets. Disable only if you have specific memory/disk constraints
Dask-based Upsampling: Uses ZarrNii's .upsample() method which leverages dask for efficient parallel upsampling

Integration with Other Operations

The scaled processing plugins integrate seamlessly with other ZarrNii operations:

# Chain operations
result = (znimg
    .apply_scaled_processing(GaussianBiasFieldCorrection())
    .downsample(level=1)
    .segment_otsu())

# Save to different formats
result.to_nifti("processed.nii")
result.to_ome_zarr("processed.ome.zarr")