diff --git a/docs/refactor/consolidation-20261009/L5-arraybridge-operations.md b/docs/refactor/consolidation-20261009/L5-arraybridge-operations.md new file mode 100644 index 0000000000..0e159fab89 --- /dev/null +++ b/docs/refactor/consolidation-20261009/L5-arraybridge-operations.md @@ -0,0 +1,105 @@ +# L5: Per-backend processors written once against arraybridge + +**Head audited:** `openhcs` `main` at `b987b7b89` (#1190); `arraybridge` `main` at `31632d8` (0.3.8). **Rules:** [00-RULES.md](00-RULES.md). **Step 2** (replaces C4 in [01-INDEX.md](01-INDEX.md)). +**Shared abstractions:** none from [02-SHARED-ABSTRACTIONS.md](02-SHARED-ABSTRACTIONS.md). *Builds:* the `ArrayOperations` family in arraybridge, keyed by memory type, with `ArrayOperations.for_memory(...)`. *Uses:* `EnumKeyedStrategyMixin` for the projection and spatial-bin families. + +## What is wrong + +**Six processor modules restate the same image operations once per array library, and the arraybridge family that already models "what differs per array library" is reached through a string roster.** + +`openhcs/processing/backends/processors/` holds `numpy_processor.py` (1,086 lines), `cupy_processor.py` (1,235), `torch_processor.py` (868), `jax_processor.py` (768), `tensorflow_processor.py` (835) and `pyclesperanto_processor.py` (819): 5,611 lines, plus `percentile_utils.py` (187), which already writes percentile normalization once against injected `percentile_func`/`clip_func`/`ones_like_func` callables. + +Census at `b987b7b89`. Thirteen names are defined in all six files (78 definitions); `_gaussian_blur` in three; `spatial_bin_2d`, `spatial_bin_3d` and `crop` in two. Each cell says how the backend's definition relates to the others: **=** identical modulo the array library (the same steps, each step one library call), **≠** a genuinely different algorithm or numeric rule that per-backend parity must keep. + +| Operation | numpy | cupy | torch | jax | tensorflow | pyclesperanto | +|---|---|---|---|---|---|---| +| `_validate_3d_array` | = | = | = | = | = | = (shape only) | +| `create_weight_mask` | = | = | = | = | = | = | +| `create_linear_weight_mask` | = | = | = (ramp includes its endpoint) | = | = (ramp includes its endpoint) | ≠ min(1, distance/margin) ramps | +| `apply_mask` | = | = | = | = | = | ≠ native `multiply_images`, no float32 detour | +| `max_projection` | = | = | = | = | = | = (z-projection kernel) | +| `mean_projection` | = | = | = | = | = | ≠ keeps the float32 z-mean, no cast back | +| `create_projection` + per-file strategy family | = (+MIN) | = | = | = | = | = except MEAN (above) | +| `spatial_bin_2d`/`_3d` + per-file strategy family | = | = | – | – | – | – | +| `crop` | = | = | – | – | – | – | +| `stack_equalize_histogram` | = | = | ≠ `histc` + floor lookup | ≠ `take` lookup, own bin width | ≠ `histogram_fixed_width` | ≠ 3-D CLAHE | +| `percentile_normalize` | = | = | ≠ `quantile`, sampling above 10M, uint16 output | ≠ `isclose`, uint16 output | ≠ floor-index percentile | ≠ min/max | +| `stack_percentile_normalize` | = | = | ≠ | ≠ | ≠ | ≠ min/max | +| `create_composite` (weight validation) | = | = | ≠ skips weights ≤ 0, no normalization check | = | ≠ takes a list of stacks | = | +| `create_composite` (arithmetic) | ≠ numba kernel, NumPy pairwise order | = | ≠ | = | ≠ | ≠ native kernels | +| `sharpen` | ≠ scikit-image | ≠ vectorized ndimage | ≠ zero-padded blur | ≠ reflect-padded blur, int→uint16 | ≠ `tf.image` blur | ≠ native, no normalization | +| `tophat` | ≠ scikit-image | ≠ ndimage + own resize | ≠ unfold morphology | ≠ per-pixel loops | ≠ depthwise conv | ≠ `top_hat_sphere` | + +Totals: the six files hold 145 top-level definitions (numpy 34, cupy 36, torch 17, jax 18, tensorflow 17, pyclesperanto 23). Identical modulo the array library: 55 definitions of 13 operations (counting composite weight validation separately). Genuinely different: 31 definitions (`sharpen` ×6, `tophat` ×6, histogram equalization ×4, percentile normalization ×8, composite arithmetic ×4, and pyclesperanto's linear mask, `apply_mask` and `mean_projection`). Backend-only operations (numpy 7, cupy 7, pyclesperanto 4) are not restatements. + +Restated mechanisms beyond the operations: +- Five copies of a `StackProjectionStrategy` family (numpy, cupy, torch, jax, tensorflow, pyclesperanto: six roots, 13 leaves) and two copies of a `SpatialBinStrategy` family (numpy, cupy: two roots, 8 leaves), all with the same dispatch. +- Six copies of the framework-import boilerplate (`optional_import_placeholder`, the `OPENHCS_SUBPROCESS_NO_GPU` branch in three files, `_check_pyclesperanto_available`, the TensorFlow version probe). +- `PyclesperantoMaskDimensionality` is an enum whose `apply` branches on `self is MASK_2D` (rule 1a). +- In arraybridge, `MemoryType` reaches its `ArrayOperations` leaf through a string field naming a module-level singleton (`"NUMPY_OPERATIONS"` … six constants); the leaf does not know its memory type. + +Defects found while measuring (kept, because parity is the acceptance check; listed for the owner): +- TensorFlow: every decorated call fails in arraybridge with `Cannot interpret 'tf.uint16' as a data type` (the TensorFlow leaf inherits the NumPy `dtype_name`; the `_tensorflow_dtype_name` helper is unused; fixed in this surface's arraybridge PR); undecorated, `sharpen` calls `tf.image.gaussian_blur` (not a TensorFlow API), percentile normalization and histogram equalization fail on dtype mismatches, and `tophat`'s "erosion" is a sum convolution. +- torch: `tophat` raises in `F.pad`; decorated `max_projection`, `mean_projection` and `sharpen` on uint16 are nondeterministic run to run on the desktop GPU (the undecorated bodies are deterministic). +- pyclesperanto 0.19: `apply_mask` (2-D mask), `create_composite` and `equalize_histogram_per_slice` fail on plane assignment (`index_x`). + +## Target + +arraybridge owns what differs per array library: + +```python +class ArrayOperations(ABC, metaclass=AutoRegisterMeta): # keyed by memory_type + memory_type: ClassVar[str | None] = None + @classmethod + def for_memory(cls, memory_type) -> ArrayOperations # one instance per leaf + module # the framework, imported on first use + accepts, astype, max, min, mean, sum, prepend_axis, # primitives that differ per library + linspace, ones, assign, outer, floor +class NumpyArrayOperations(MutableArrayOperations): memory_type = "numpy" +... # one leaf per memory type +``` + +`MemoryType` finds its leaf with `ArrayOperations.for_memory(self)`; the operations-name field and the six `*_OPERATIONS` constants are deleted. + +OpenHCS writes each shared operation once, in `processors/image_stack_operations.py`, against an `ArrayOperations` leaf: stack validation, weight masks, mask application, max/mean/min projection and `create_projection` (one `StackProjection` family), spatial binning (one `SpatialBin` family), crop, histogram equalization and percentile normalization for the NumPy-compatible libraries, composite weight validation and the weighted composite. `percentile_utils.py` is folded in and deleted. Each backend module keeps only its registrations: a decorated declaration per public name (registered identity, signature and contract unchanged) that calls the shared operation with its leaf, plus the operations that genuinely differ (the ≠ cells). Framework names for annotations stay `optional_import_placeholder` handles (the decorators evaluate annotations at registration, so a lazy proxy would import or fail on every absent framework); the three `OPENHCS_SUBPROCESS_NO_GPU` branches, `_check_pyclesperanto_available`, the TensorFlow version probe and `PyclesperantoMaskDimensionality` are deleted. + +Decision with a default: `NumpyStackProjectionMethod` stays as numpy's `create_projection` annotation, because saved pipelines import it (rule 2 forbids changing durable pipelines without an owner decision). `StackProjectionMethod` gains `MIN`; both enums key the one `StackProjection` family by value. + +## Guards + +`tests/unit/test_processor_operations_written_once.py`: +- AST: no backend module defines a class keyed by `StackProjectionMethod`/`SpatialBinMethod`, a `_validate_3d_array`, or reads `OPENHCS_SUBPROCESS_NO_GPU`; `percentile_utils.py` does not exist. +- AST: every function defined in two or more backend modules delegates to a shared operation, except the census's ≠ cells (listed once in the test). +- Each backend's `OPERATIONS` is `ArrayOperations.for_memory` of its own memory type. +- In arraybridge: `types.py` has no `_OPERATIONS` string and `array_operations.py` no module-level operations instance; every `MemoryType` resolves a leaf through `ArrayOperations.for_memory`. + +## Tests + +- arraybridge: one family test (every `MemoryType` has exactly one leaf; `for_memory` returns it; primitives on NumPy). +- OpenHCS: one family test over the six backends: each `create_projection` method enum is served by the `StackProjection` family; the shared operations agree with the NumPy reference on the installed backends. `test_processor_method_strategies.py` is rewritten to the family level. +- Numeric parity per backend (not committed): every public declaration on fixed inputs, before and after, decorated and undecorated; NumPy and JAX here, CuPy, torch, TensorFlow and pyclesperanto on the desktop GPU. + +## New-case experiments + +Today a new array library needs a sixth copy of every shared operation. After: one `ArrayOperations` leaf in arraybridge with the primitives, and one registration per operation. + +## Done when + +Every = cell above is a delegation to one shared operation; the per-file strategy families, validators and import boilerplate are gone; parity holds per backend; the guards pass; the arraybridge PR is open with a minor bump and green CI. + +## Result + +Measured on the branch (`image_stack_operations.py` 435 lines; the six backend modules 4,141 lines, from 5,611; `percentile_utils.py` deleted). + +Numeric parity, every public declaration on fixed uint16/uint8/float32/constant stacks (`.claude/l5/parity_harness.py`, not committed), called decorated and undecorated, before (`b987b7b89`) and after: + +| Backend | Where | Cases | Exact | Changed | +|---|---|---|---|---| +| numpy | here and desktop | 236 | 236 | none (the unregistered `min_projection` helper is deleted; `create_projection(MIN)` replaces it) | +| jax (CPU) | here and desktop | 140 | 140 | +6 new `create_projection(MIN)` cases | +| cupy | desktop RTX 2070 SUPER | 282 | 282 | +6 MIN | +| pyclesperanto 0.19 | desktop GPU (OpenCL) | 188 | 188 | +6 MIN | +| torch 2.6 (CUDA) | desktop GPU | 70 undecorated | 70 | +MIN; decorated calls are nondeterministic on main itself (two main runs differed in 2 of 70 decorated cases); after, 1 to 2 decorated cases differ per run, in operations this surface did not change | +| tensorflow 2.19 (GPU) | desktop GPU | 70 undecorated | 62 same, 8 error→value | `create_projection` (6) and `create_weight_mask` (2) used to fail through the nested decorated call; their values now equal TensorFlow's own projection and mask bodies. Decorated calls, all failing on main, now run and equal the undecorated results (33 values, 40 same errors) | + +Exact means bitwise-equal arrays and equal dtypes; errors compare by exception type. No case needed a tolerance. JAX on the GPU could not be checked: the desktop's JAX CUDA plugin finds no visible GPU, so JAX ran on the CPU on both machines. diff --git a/external/arraybridge b/external/arraybridge index 31632d8ce3..f614ab9bbe 160000 --- a/external/arraybridge +++ b/external/arraybridge @@ -1 +1 @@ -Subproject commit 31632d8ce31016f0fe6ef28df61f49da397bf3e6 +Subproject commit f614ab9bbe8a7fb8372e548e57f0d42896bdeb2c diff --git a/openhcs/core/pipeline/funcstep_contract_validator.py b/openhcs/core/pipeline/funcstep_contract_validator.py index b39bc2995e..5d1fedf952 100644 --- a/openhcs/core/pipeline/funcstep_contract_validator.py +++ b/openhcs/core/pipeline/funcstep_contract_validator.py @@ -274,8 +274,8 @@ def _resolve_relative_import( """Resolve an ImportFrom module and level to an absolute module name. Args: - module: The ImportFrom module (e.g. 'percentile_utils' for - `from .percentile_utils import ...`). + module: The ImportFrom module (e.g. 'image_stack_operations' for + `from .image_stack_operations import ...`). level: The ImportFrom level (1='.', 2='..', ...). Returns: diff --git a/openhcs/processing/backends/processors/__init__.py b/openhcs/processing/backends/processors/__init__.py index 4bcd5110fb..b9524b9dc6 100644 --- a/openhcs/processing/backends/processors/__init__.py +++ b/openhcs/processing/backends/processors/__init__.py @@ -1,9 +1,9 @@ """ -Image processors for different backends. +Image processors for the array backends (numpy, cupy, torch, tensorflow, jax, +pyclesperanto). -This module contains image processing functions implemented for different -computational backends (numpy, cupy, torch, tensorflow, jax). - -Each processor module provides the same set of functions but optimized -for its specific backend. +Operations shared by several backends are written once in +``image_stack_operations`` against arraybridge ``ArrayOperations``; each +``_processor`` module registers them for its memory type and keeps +only the operations whose algorithm differs on that backend. """ diff --git a/openhcs/processing/backends/processors/cupy_processor.py b/openhcs/processing/backends/processors/cupy_processor.py index a1d9fd11ff..7126155a57 100644 --- a/openhcs/processing/backends/processors/cupy_processor.py +++ b/openhcs/processing/backends/processors/cupy_processor.py @@ -1,28 +1,25 @@ -""" -CuPy Image Processor Implementation - -This module implements the ImageProcessorInterface using CuPy as the backend. -It leverages GPU acceleration for image processing operations. +"""CuPy processor registrations. -Doctrinal Clauses: -- Clause 3 — Declarative Primacy: All functions are pure and stateless -- Clause 88 — No Inferred Capabilities: Explicit CuPy dependency -- Clause 106-A — Declared Memory Types: All methods specify CuPy arrays +Shared operations are written once in ``image_stack_operations``; this module +registers them for CuPy and holds the CuPy-specific algorithms (sharpening, +top-hat, Sobel kernels). """ from __future__ import annotations -import logging -import os from abc import abstractmethod -from typing import Any, List, Optional, Tuple +from typing import List, Optional, Tuple +from arraybridge import MemoryType +from arraybridge.array_operations import ArrayOperations from metaclass_registry import AutoRegisterMeta from openhcs.core.artifacts import ArtifactSpec, ImageArtifactType +from openhcs.utils.import_utils import optional_import_placeholder from openhcs.core.memory import cupy as cupy_func from openhcs.core.pipeline.function_contracts import artifact_inputs from metaclass_registry.strategies import EnumKeyedStrategyMixin +from openhcs.processing.backends.processors import image_stack_operations as operations from openhcs.processing.backends.processors.method_axes import ( EdgeMagnitudeMethod, ScipyBoundaryMode, @@ -30,28 +27,10 @@ SpatialBinMethod, StackProjectionMethod, ) -from openhcs.utils.import_utils import optional_import_placeholder -logger = logging.getLogger(__name__) +cp = optional_import_placeholder("cupy") -# Check if we're in subprocess runner mode and should skip GPU imports -if os.getenv("OPENHCS_SUBPROCESS_NO_GPU") == "1": - # Subprocess runner mode - skip GPU imports - cp = None - ndimage = None - cucim_filters = None - logger.info("Subprocess runner mode - skipping cupy import") -else: - # Normal mode - import CuPy as an optional dependency - cp = optional_import_placeholder("cupy") - ndimage = None - if cp: - ndimage = optional_import_placeholder("cupyx.scipy.ndimage") - - # Import CuCIM for edge detection - cucim_filters = optional_import_placeholder("cucim.skimage.filters") - -logger = logging.getLogger(__name__) +OPERATIONS = ArrayOperations.for_memory(MemoryType.CUPY) @cupy_func @@ -69,56 +48,9 @@ def create_linear_weight_mask( Returns: 2D CuPy weight mask of shape (height, width) """ - # The compiler will ensure this function is only called when CuPy is available - # No need to check for CuPy availability here - - margin_y = int(cp.floor(height * margin_ratio)) - margin_x = int(cp.floor(width * margin_ratio)) - - weight_y = cp.ones(height, dtype=cp.float32) - if margin_y > 0: - ramp_top = cp.linspace(0, 1, margin_y, endpoint=False) - ramp_bottom = cp.linspace(1, 0, margin_y, endpoint=False) - weight_y[:margin_y] = ramp_top - weight_y[-margin_y:] = ramp_bottom - - weight_x = cp.ones(width, dtype=cp.float32) - if margin_x > 0: - ramp_left = cp.linspace(0, 1, margin_x, endpoint=False) - ramp_right = cp.linspace(1, 0, margin_x, endpoint=False) - weight_x[:margin_x] = ramp_left - weight_x[-margin_x:] = ramp_right - - # Create 2D weight mask - weight_mask = cp.outer(weight_y, weight_x) - - return weight_mask - - -def _validate_3d_array(array: Any, name: str = "input") -> None: - """ - Validate that the input is a 3D CuPy array. - - Args: - array: Array to validate - name: Name of the array for error messages - - Raises: - TypeError: If the array is not a CuPy array - ValueError: If the array is not 3D - ImportError: If CuPy is not available - """ - # The compiler will ensure this function is only called when CuPy is available - # No need to check for CuPy availability here - - if not isinstance(array, cp.ndarray): - raise TypeError( - f"{name} must be a CuPy array, got {type(array)}. " - f"No automatic conversion is performed to maintain explicit contracts." - ) - - if array.ndim != 3: - raise ValueError(f"{name} must be a 3D array, got {array.ndim}D") + return operations.linear_weight_mask( + OPERATIONS, height, width, margin_ratio, ramp_endpoint=False + ) @cupy_func @@ -139,7 +71,9 @@ def sharpen( Returns: Sharpened 3D CuPy array of shape (Z, Y, X) """ - _validate_3d_array(image) + from cupyx.scipy import ndimage + + operations.require_stack(OPERATIONS, image) # Store original dtype dtype = image.dtype @@ -202,33 +136,14 @@ def percentile_normalize( high_percentile: Upper percentile (0-100) target_min: Target minimum value (auto-detected from dtype if None) target_max: Target maximum value (auto-detected from dtype if None) - preserve_dtype: If True, output same dtype as input. If False, use target range. + preserve_dtype: No effect; the output always has the input dtype. Returns: - Normalized 3D CuPy array of shape (Z, Y, X) with same or specified dtype + Normalized 3D CuPy array of shape (Z, Y, X) with the input dtype """ - _validate_3d_array(image) - - # Import shared utilities - from .percentile_utils import resolve_target_range, slice_percentile_normalize_core - - # Auto-detect target range based on input dtype if not specified - target_min, target_max = resolve_target_range(image.dtype, target_min, target_max) - - # Use shared core logic with CuPy-specific functions - return slice_percentile_normalize_core( - image=image, - low_percentile=low_percentile, - high_percentile=high_percentile, - target_min=target_min, - target_max=target_max, - percentile_func=cp.percentile, - clip_func=cp.clip, - ones_like_func=cp.ones_like, - zeros_like_func=lambda arr, dtype=None: cp.zeros_like( - arr, dtype=dtype or cp.float32 - ), - preserve_dtype=preserve_dtype, + del preserve_dtype # the output always has the input dtype + return operations.percentile_normalize_planes( + OPERATIONS, image, low_percentile, high_percentile, target_min, target_max ) @@ -253,30 +168,14 @@ def stack_percentile_normalize( high_percentile: Upper percentile (0-100) target_min: Target minimum value (auto-detected from dtype if None) target_max: Target maximum value (auto-detected from dtype if None) - preserve_dtype: If True, output same dtype as input. If False, use target range. + preserve_dtype: No effect; the output always has the input dtype. Returns: - Normalized 3D CuPy array of shape (Z, Y, X) with same or specified dtype + Normalized 3D CuPy array of shape (Z, Y, X) with the input dtype """ - _validate_3d_array(stack) - - # Import shared utilities - from .percentile_utils import percentile_normalize_core, resolve_target_range - - # Auto-detect target range based on input dtype if not specified - target_min, target_max = resolve_target_range(stack.dtype, target_min, target_max) - - # Use shared core logic with CuPy-specific functions - return percentile_normalize_core( - stack=stack, - low_percentile=low_percentile, - high_percentile=high_percentile, - target_min=target_min, - target_max=target_max, - percentile_func=lambda arr, pct: cp.percentile(arr, pct), - clip_func=cp.clip, - ones_like_func=cp.ones_like, - preserve_dtype=preserve_dtype, + del preserve_dtype # the output always has the input dtype + return operations.percentile_normalize_stack( + OPERATIONS, stack, low_percentile, high_percentile, target_min, target_max ) @@ -294,70 +193,7 @@ def create_composite( Returns: Composite 3D CuPy array of shape (1, Y, X) """ - # 🔄 MEMORY CONVERSION LOGGING: Log what we actually received - import logging - - logger = logging.getLogger(__name__) - logger.info( - f"🔄 CREATE_COMPOSITE: Called with stack type: {type(stack)}, shape: {getattr(stack, 'shape', 'no shape')}" - ) - logger.info( - f"🔄 CREATE_COMPOSITE: weights parameter - type: {type(weights)}, value: {weights}" - ) - - # Validate input is 3D array - if not hasattr(stack, "shape") or len(stack.shape) != 3: - raise TypeError( - f"stack must be a 3D CuPy array, got shape: {getattr(stack, 'shape', 'no shape')}" - ) - - n_slices, height, width = stack.shape - - # Default weights if none provided - if weights is None: - # Equal weights for all slices - weights = [1.0 / n_slices] * n_slices - logger.info(f"🔄 CREATE_COMPOSITE: Using default equal weights: {weights}") - elif isinstance(weights, (list, tuple)): - # Convert tuple to list if needed - weights = list(weights) - logger.info(f"🔄 CREATE_COMPOSITE: Using provided weights: {weights}") - if len(weights) != n_slices: - raise ValueError( - f"Number of weights ({len(weights)}) must match number of slices ({n_slices})" - ) - else: - # Log the problematic type and value for debugging - logger.error( - f"🔄 CREATE_COMPOSITE: Invalid weights type - expected list/tuple/None, got {type(weights)}: {weights}" - ) - raise TypeError( - f"weights must be a list of values or None, got {type(weights)}: {weights}" - ) - - # Normalize weights to sum to 1 - weight_sum = sum(weights) - if weight_sum == 0: - raise ValueError("Sum of weights cannot be zero") - normalized_weights = [w / weight_sum for w in weights] - - # Convert weights to CuPy array for efficient computation - # CRITICAL: Use float32 for weights to preserve fractional values, not stack.dtype - weights_array = cp.array(normalized_weights, dtype=cp.float32) - - # Reshape weights for broadcasting: (N, 1, 1) to multiply with (N, Y, X) - weights_array = weights_array.reshape(n_slices, 1, 1) - - # Create composite by weighted sum along the first axis - # Convert stack to float32 for computation to avoid precision loss - stack_float = stack.astype(cp.float32) - weighted_stack = stack_float * weights_array - composite_slice = cp.sum(weighted_stack, axis=0, keepdims=True) # Keep as (1, Y, X) - - # Convert back to original dtype - composite_slice = composite_slice.astype(stack.dtype) - - return composite_slice + return operations.weighted_composite(OPERATIONS, stack, weights) @artifact_inputs(ArtifactSpec.input("mask", ImageArtifactType, parameter_name="mask")) @@ -376,35 +212,7 @@ def apply_mask(image: "cp.ndarray", mask: "cp.ndarray") -> "cp.ndarray": Returns: Masked 3D CuPy array of shape (Z, Y, X) """ - _validate_3d_array(image) - - # Handle 2D mask (apply to each Z-slice) - if isinstance(mask, cp.ndarray) and mask.ndim == 2: - if mask.shape != image.shape[1:]: - raise ValueError( - f"2D mask shape {mask.shape} doesn't match image slice shape {image.shape[1:]}" - ) - - # Apply 2D mask to all Z-slices simultaneously using broadcasting - GPU PARALLELIZED - # Broadcasting mask from (Y, X) to (1, Y, X) allows vectorized operation across all slices - mask_3d = mask[None, :, :] # Shape: (1, Y, X) - result = image.astype(cp.float32) * mask_3d.astype(cp.float32) - - return result.astype(image.dtype) - - # Handle 3D mask - if isinstance(mask, cp.ndarray) and mask.ndim == 3: - if mask.shape != image.shape: - raise ValueError( - f"3D mask shape {mask.shape} doesn't match image shape {image.shape}" - ) - - # Apply 3D mask directly - MATCH NUMPY EXACTLY - masked = image.astype(cp.float32) * mask.astype(cp.float32) - return masked.astype(image.dtype) - - # If we get here, the mask is neither 2D nor 3D CuPy array - raise TypeError(f"mask must be a 2D or 3D CuPy array, got {type(mask)}") + return operations.apply_mask(OPERATIONS, image, mask) @cupy_func @@ -421,11 +229,7 @@ def create_weight_mask( Returns: 2D CuPy weight mask of shape (Y, X) """ - if not isinstance(shape, tuple) or len(shape) != 2: - raise TypeError("shape must be a tuple of (height, width)") - - height, width = shape - return create_linear_weight_mask(height, width, margin_ratio) + return create_linear_weight_mask(*operations.weight_mask_extent(shape), margin_ratio) @cupy_func @@ -439,11 +243,7 @@ def max_projection(stack: "cp.ndarray") -> "cp.ndarray": Returns: 3D CuPy array of shape (1, Y, X) """ - _validate_3d_array(stack) - - # Create max projection - projection_2d = cp.max(stack, axis=0) - return projection_2d.reshape(1, projection_2d.shape[0], projection_2d.shape[1]) + return operations.max_projection(OPERATIONS, stack) @cupy_func @@ -457,73 +257,7 @@ def mean_projection(stack: "cp.ndarray") -> "cp.ndarray": Returns: 3D CuPy array of shape (1, Y, X) """ - _validate_3d_array(stack) - - # Create mean projection - projection_2d = cp.mean(stack, axis=0).astype(stack.dtype) - return projection_2d.reshape(1, projection_2d.shape[0], projection_2d.shape[1]) - - -class CupySpatialBinStrategy( - EnumKeyedStrategyMixin[SpatialBinMethod], metaclass=AutoRegisterMeta -): - __enum_member_attr__ = "method" - - @abstractmethod - def apply(self, array: "cp.ndarray", axis: tuple[int, ...]) -> "cp.ndarray": - raise NotImplementedError - - -class CupyMeanSpatialBinStrategy(CupySpatialBinStrategy): - method = SpatialBinMethod.MEAN - - def apply(self, array: "cp.ndarray", axis: tuple[int, ...]) -> "cp.ndarray": - return cp.mean(array, axis=axis) - - -class CupySumSpatialBinStrategy(CupySpatialBinStrategy): - method = SpatialBinMethod.SUM - - def apply(self, array: "cp.ndarray", axis: tuple[int, ...]) -> "cp.ndarray": - return cp.sum(array, axis=axis) - - -class CupyMaxSpatialBinStrategy(CupySpatialBinStrategy): - method = SpatialBinMethod.MAX - - def apply(self, array: "cp.ndarray", axis: tuple[int, ...]) -> "cp.ndarray": - return cp.max(array, axis=axis) - - -class CupyMinSpatialBinStrategy(CupySpatialBinStrategy): - method = SpatialBinMethod.MIN - - def apply(self, array: "cp.ndarray", axis: tuple[int, ...]) -> "cp.ndarray": - return cp.min(array, axis=axis) - - -class CupyStackProjectionStrategy( - EnumKeyedStrategyMixin[StackProjectionMethod], metaclass=AutoRegisterMeta -): - __enum_member_attr__ = "method" - - @abstractmethod - def apply(self, stack: "cp.ndarray") -> "cp.ndarray": - raise NotImplementedError - - -class CupyMaxStackProjectionStrategy(CupyStackProjectionStrategy): - method = StackProjectionMethod.MAX - - def apply(self, stack: "cp.ndarray") -> "cp.ndarray": - return max_projection(stack) - - -class CupyMeanStackProjectionStrategy(CupyStackProjectionStrategy): - method = StackProjectionMethod.MEAN - - def apply(self, stack: "cp.ndarray") -> "cp.ndarray": - return mean_projection(stack) + return operations.mean_projection(OPERATIONS, stack) @cupy_func @@ -546,35 +280,7 @@ def spatial_bin_2d( Returns: Binned 3D CuPy array of shape (Z, Y//bin_size, X//bin_size) """ - _validate_3d_array(stack) - - if bin_size <= 0: - raise ValueError("bin_size must be positive") - - z_slices, height, width = stack.shape - - # Calculate output dimensions - new_height = height // bin_size - new_width = width // bin_size - - if new_height == 0 or new_width == 0: - raise ValueError( - f"bin_size {bin_size} is too large for image dimensions {height}x{width}" - ) - - # Crop to make dimensions divisible by bin_size - crop_height = new_height * bin_size - crop_width = new_width * bin_size - cropped_stack = stack[:, :crop_height, :crop_width] - - # Reshape for binning: (Z, new_height, bin_size, new_width, bin_size) - reshaped = cropped_stack.reshape( - z_slices, new_height, bin_size, new_width, bin_size - ) - - result = CupySpatialBinStrategy.for_enum_member(method).apply(reshaped, axis=(2, 4)) - - return result.astype(stack.dtype) + return operations.spatial_bin_2d(OPERATIONS, stack, bin_size, method) @cupy_func @@ -597,39 +303,7 @@ def spatial_bin_3d( Returns: Binned 3D CuPy array of shape (Z//bin_size, Y//bin_size, X//bin_size) """ - _validate_3d_array(stack) - - if bin_size <= 0: - raise ValueError("bin_size must be positive") - - depth, height, width = stack.shape - - # Calculate output dimensions - new_depth = depth // bin_size - new_height = height // bin_size - new_width = width // bin_size - - if new_depth == 0 or new_height == 0 or new_width == 0: - raise ValueError( - f"bin_size {bin_size} is too large for stack dimensions {depth}x{height}x{width}" - ) - - # Crop to make dimensions divisible by bin_size - crop_depth = new_depth * bin_size - crop_height = new_height * bin_size - crop_width = new_width * bin_size - cropped_stack = stack[:crop_depth, :crop_height, :crop_width] - - # Reshape for 3D binning: (new_depth, bin_size, new_height, bin_size, new_width, bin_size) - reshaped = cropped_stack.reshape( - new_depth, bin_size, new_height, bin_size, new_width, bin_size - ) - - result = CupySpatialBinStrategy.for_enum_member(method).apply( - reshaped, axis=(1, 3, 5) - ) - - return result.astype(stack.dtype) + return operations.spatial_bin_3d(OPERATIONS, stack, bin_size, method) @cupy_func @@ -654,42 +328,10 @@ def stack_equalize_histogram( Returns: Equalized 3D CuPy array of shape (Z, Y, X) """ - _validate_3d_array(stack) - - # Remember input dtype to preserve it - input_dtype = stack.dtype - - # MATCH NUMPY EXACTLY - Flatten the entire stack to compute the global histogram - flat_stack = stack.flatten() - - # Calculate the histogram and cumulative distribution function (CDF) - MATCH NUMPY EXACTLY - hist, bin_edges = cp.histogram(flat_stack, bins=bins, range=(range_min, range_max)) - cdf = hist.cumsum() - - # Normalize the CDF to the input dtype range - # Avoid division by zero - if cdf[-1] > 0: - if cp.issubdtype(input_dtype, cp.integer): - dtype_info = cp.iinfo(input_dtype) - cdf = dtype_info.max * cdf / cdf[-1] - else: - # For float dtypes, normalize to [0, 1] - cdf = cdf / cdf[-1] - - # Use linear interpolation to map input values to equalized values - MATCH NUMPY EXACTLY - equalized_stack = cp.interp(stack.flatten(), bin_edges[:-1], cdf).reshape( - stack.shape + return operations.stack_equalize_histogram( + OPERATIONS, stack, bins, range_min, range_max ) - # Convert back to input dtype - if cp.issubdtype(input_dtype, cp.integer): - dtype_info = cp.iinfo(input_dtype) - return cp.clip(equalized_stack, dtype_info.min, dtype_info.max).astype( - input_dtype - ) - else: - return equalized_stack.astype(input_dtype) - @cupy_func def create_projection( @@ -701,14 +343,12 @@ def create_projection( Args: stack: 3D CuPy array of shape (Z, Y, X) - method: Projection method (max_projection, mean_projection) + method: Projection method served by ``image_stack_operations.StackProjection`` Returns: 3D CuPy array of shape (1, Y, X) """ - _validate_3d_array(stack) - - return CupyStackProjectionStrategy.for_enum_member(method).apply(stack) + return operations.create_projection(OPERATIONS, stack, method) @cupy_func @@ -748,39 +388,9 @@ def crop( cp.ndarray Cropped 3D array of shape (depth, height, width) """ - _validate_3d_array(input_image) - - # Validate crop parameters - if width <= 0 or height <= 0 or depth <= 0: - raise ValueError( - f"Crop dimensions must be positive: width={width}, height={height}, depth={depth}" - ) - - if start_x < 0 or start_y < 0 or start_z < 0: - raise ValueError( - f"Start coordinates must be non-negative: start_x={start_x}, start_y={start_y}, start_z={start_z}" - ) - - # Get input dimensions - input_depth, input_height, input_width = input_image.shape - - # Calculate end coordinates - end_x = start_x + width - end_y = start_y + height - end_z = start_z + depth - - # Validate bounds - if end_x > input_width or end_y > input_height or end_z > input_depth: - raise ValueError( - f"Crop region extends beyond image bounds. " - f"Image shape: {input_image.shape}, " - f"Crop region: ({start_z}:{end_z}, {start_y}:{end_y}, {start_x}:{end_x})" - ) - - # Perform the crop using CuPy slicing - cropped = input_image[start_z:end_z, start_y:end_y, start_x:end_x] - - return cropped + return operations.crop( + OPERATIONS, input_image, start_x, start_y, start_z, width, height, depth + ) def _create_disk_cupy(radius: int) -> "cp.ndarray": @@ -868,7 +478,9 @@ def tophat( Returns: Filtered 3D CuPy array of shape (Z, Y, X) """ - _validate_3d_array(image) + from cupyx.scipy import ndimage + + operations.require_stack(OPERATIONS, image) # Process each Z-slice independently - IMPROVED MATCH TO NUMPY # NOTE: For loop is ALGORITHMICALLY NECESSARY here due to complex multi-step @@ -924,9 +536,6 @@ def _get_sobel_2d_kernel(): """Get or create the Sobel 2D parallel kernel.""" global _sobel_2d_parallel if _sobel_2d_parallel is None: - if not cp: - raise ImportError("CuPy is required for GPU Sobel operations") - _sobel_2d_parallel = cp.ElementwiseKernel( "raw T input, int32 Y, int32 X, int32 mode, T cval", "T output", @@ -994,7 +603,7 @@ def sobel_2d_vectorized( Returns: Edge magnitude as 3D CuPy array of shape (Z, Y, X) """ - _validate_3d_array(image) + operations.require_stack(OPERATIONS, image) Z, Y, X = image.shape input_float = image.astype(cp.float32) @@ -1038,8 +647,9 @@ def _sobel_3d_voxel( cval: float, ) -> "cp.ndarray": """Apply the 3-D Sobel implementation at its library string boundary.""" + from cupyx.scipy import ndimage - _validate_3d_array(image) + operations.require_stack(OPERATIONS, image) # Convert to float32 for processing image_float = image.astype(cp.float32) @@ -1085,7 +695,9 @@ def sobel_components( - If include_z=False: (sobel_x, sobel_y) - 2D gradients per slice - If include_z=True: (sobel_x, sobel_y, sobel_z) - 3D gradients """ - _validate_3d_array(image) + from cupyx.scipy import ndimage + + operations.require_stack(OPERATIONS, image) # Convert to float32 for processing image_float = image.astype(cp.float32) @@ -1161,7 +773,7 @@ def edge_magnitude( Returns: Edge magnitude as 3D CuPy array of shape (Z, Y, X) """ - _validate_3d_array(image) + operations.require_stack(OPERATIONS, image) return CupyEdgeMagnitudeStrategy.for_enum_member(method).apply( image, mode=mode, cval=cval @@ -1212,13 +824,7 @@ def sobel( CuCIM's sobel normalizes integer inputs to [0, 1] range and returns float64. This wrapper always preserves the input dtype by forcing dtype conversion. """ - if not cucim_filters: - raise ImportError( - "CuCIM is required for sobel edge detection but is not available" - ) - - # Import here to avoid circular dependency - from openhcs.constants.constants import MemoryType + from cucim.skimage import filters as cucim_filters # Store original dtype original_dtype = image.dtype diff --git a/openhcs/processing/backends/processors/image_stack_operations.py b/openhcs/processing/backends/processors/image_stack_operations.py new file mode 100644 index 0000000000..f0eacf92fa --- /dev/null +++ b/openhcs/processing/backends/processors/image_stack_operations.py @@ -0,0 +1,435 @@ +"""Image-stack operations written once against arraybridge ``ArrayOperations``. + +Each backend processor module registers these operations under its own +memory type and passes its ``ArrayOperations`` leaf. A leaf supplies only the +primitives its array library does differently, so a result here equals what +the backend's own library calls produce. Operations whose algorithm genuinely +differs per backend stay in that backend's module. + +Functions taking ``xp``-style calls through ``ops.array_module`` are used only +by the NumPy-compatible libraries (NumPy, CuPy). +""" + +from __future__ import annotations + +from abc import abstractmethod +from collections.abc import Sequence +from typing import Any + +import numpy as np +from arraybridge import MemoryType +from arraybridge.array_operations import ArrayOperations +from metaclass_registry import AutoRegisterMeta +from metaclass_registry.strategies import EnumKeyedStrategyMixin + +from openhcs.processing.backends.processors.method_axes import ( + SpatialBinMethod, + StackProjectionMethod, +) + + +def require_stack(ops: ArrayOperations, array: Any, name: str = "input") -> None: + """Require one native array of the leaf's framework with three axes.""" + + if not ops.accepts(array): + display_name = MemoryType(ops.memory_type).display_name + raise TypeError( + f"{name} must be a {display_name} array, got {type(array)}. " + "No automatic conversion is performed." + ) + if len(array.shape) != 3: + raise ValueError(f"{name} must be a 3D array, got {len(array.shape)}D") + + +# Integer dtypes whose composites and equalized stacks torch and TensorFlow +# clamp to the dtype's limits; other dtypes are cast directly. +CLAMPED_INTEGER_DTYPES = frozenset( + {"uint8", "uint16", "uint32", "int8", "int16", "int32", "int64"} +) + + +# ----- weight masks ----- + + +def weight_mask_extent(shape: Any) -> tuple[int, int]: + """Validate a ``(height, width)`` weight-mask shape.""" + + if not isinstance(shape, tuple) or len(shape) != 2: + raise TypeError("shape must be a tuple of (height, width)") + return shape + + +def linear_weight_mask( + ops: ArrayOperations, + height: int, + width: int, + margin_ratio: float, + *, + ramp_endpoint: bool, +) -> Any: + """Weights ramping linearly from 0 at the edges to 1 inside the margin. + + ``ramp_endpoint`` states whether each edge ramp reaches 1 on its last + margin pixel (torch and TensorFlow) or stops one step short (NumPy, CuPy, + JAX). + """ + + def axis_weights(length: int) -> Any: + margin = int(ops.floor(length * margin_ratio)) + weights = ops.ones(length, np.float32) + if margin > 0: + rising = ops.linspace(0, 1, margin, endpoint=ramp_endpoint) + falling = ops.linspace(1, 0, margin, endpoint=ramp_endpoint) + weights = ops.assign(weights, slice(None, margin), rising) + weights = ops.assign(weights, slice(-margin, None), falling) + return weights + + return ops.outer(axis_weights(height), axis_weights(width)) + + +# ----- masks and projections ----- + + +def apply_mask(ops: ArrayOperations, image: Any, mask: Any) -> Any: + """Multiply a stack by a 2-D (per plane) or 3-D mask in float32.""" + + require_stack(ops, image) + if not ops.accepts(mask) or len(mask.shape) not in (2, 3): + raise TypeError(f"mask must be a 2D or 3D native array, got {type(mask)}") + expected = tuple(image.shape[1:]) if len(mask.shape) == 2 else tuple(image.shape) + if tuple(mask.shape) != expected: + raise ValueError( + f"{len(mask.shape)}D mask shape {tuple(mask.shape)} doesn't match " + f"image shape {expected}" + ) + masked = ops.astype(image, np.float32) * ops.astype(mask, np.float32) + return ops.astype(masked, image.dtype) + + +def max_projection(ops: ArrayOperations, stack: Any) -> Any: + """Maximum over the leading axis, shaped ``(1, Y, X)``.""" + + require_stack(ops, stack) + return ops.prepend_axis(ops.max(stack, 0)) + + +def mean_projection(ops: ArrayOperations, stack: Any) -> Any: + """Mean over the leading axis in the input dtype, shaped ``(1, Y, X)``.""" + + require_stack(ops, stack) + return ops.prepend_axis(ops.astype(ops.mean(stack, 0), stack.dtype)) + + +def min_projection(ops: ArrayOperations, stack: Any) -> Any: + """Minimum over the leading axis, shaped ``(1, Y, X)``.""" + + require_stack(ops, stack) + return ops.prepend_axis(ops.min(stack, 0)) + + +class StackProjection( + EnumKeyedStrategyMixin[StackProjectionMethod], metaclass=AutoRegisterMeta +): + """One leading-axis projection, keyed by its method value.""" + + __enum_member_attr__ = "method" + + @abstractmethod + def project(self, ops: ArrayOperations, stack: Any) -> Any: + """Return the projection shaped ``(1, Y, X)``.""" + + +class MaxStackProjection(StackProjection): + method = StackProjectionMethod.MAX + + def project(self, ops: ArrayOperations, stack: Any) -> Any: + return max_projection(ops, stack) + + +class MeanStackProjection(StackProjection): + method = StackProjectionMethod.MEAN + + def project(self, ops: ArrayOperations, stack: Any) -> Any: + return mean_projection(ops, stack) + + +class MinStackProjection(StackProjection): + method = StackProjectionMethod.MIN + + def project(self, ops: ArrayOperations, stack: Any) -> Any: + return min_projection(ops, stack) + + +def create_projection(ops: ArrayOperations, stack: Any, method: Any) -> Any: + """Project a stack with the method's registered projection. + + ``method`` is any enum whose values name ``StackProjection`` members. + """ + + require_stack(ops, stack) + return StackProjection.for_enum_member(method).project(ops, stack) + + +# ----- spatial binning (NumPy-compatible libraries) ----- + + +class SpatialBin(EnumKeyedStrategyMixin[SpatialBinMethod], metaclass=AutoRegisterMeta): + """Reduction applied within each spatial bin.""" + + __enum_member_attr__ = "method" + + @abstractmethod + def reduce(self, ops: ArrayOperations, array: Any, axis: tuple[int, ...]) -> Any: + """Reduce the bin axes.""" + + +class MeanSpatialBin(SpatialBin): + method = SpatialBinMethod.MEAN + + def reduce(self, ops: ArrayOperations, array: Any, axis: tuple[int, ...]) -> Any: + return ops.mean(array, axis) + + +class SumSpatialBin(SpatialBin): + method = SpatialBinMethod.SUM + + def reduce(self, ops: ArrayOperations, array: Any, axis: tuple[int, ...]) -> Any: + return ops.sum(array, axis) + + +class MaxSpatialBin(SpatialBin): + method = SpatialBinMethod.MAX + + def reduce(self, ops: ArrayOperations, array: Any, axis: tuple[int, ...]) -> Any: + return ops.max(array, axis) + + +class MinSpatialBin(SpatialBin): + method = SpatialBinMethod.MIN + + def reduce(self, ops: ArrayOperations, array: Any, axis: tuple[int, ...]) -> Any: + return ops.min(array, axis) + + +def _require_bin_size(bin_size: int) -> None: + if bin_size <= 0: + raise ValueError("bin_size must be positive") + + +def spatial_bin_2d( + ops: ArrayOperations, stack: Any, bin_size: int, method: SpatialBinMethod +) -> Any: + """Bin each plane in ``bin_size`` x ``bin_size`` blocks.""" + + require_stack(ops, stack) + _require_bin_size(bin_size) + planes, height, width = stack.shape + new_height, new_width = height // bin_size, width // bin_size + if new_height == 0 or new_width == 0: + raise ValueError( + f"bin_size {bin_size} is too large for image dimensions {height}x{width}" + ) + cropped = stack[:, : new_height * bin_size, : new_width * bin_size] + blocks = cropped.reshape(planes, new_height, bin_size, new_width, bin_size) + binned = SpatialBin.for_enum_member(method).reduce(ops, blocks, (2, 4)) + return ops.astype(binned, stack.dtype) + + +def spatial_bin_3d( + ops: ArrayOperations, stack: Any, bin_size: int, method: SpatialBinMethod +) -> Any: + """Bin the stack in ``bin_size``-sided cubes.""" + + require_stack(ops, stack) + _require_bin_size(bin_size) + depth, height, width = stack.shape + new_depth, new_height, new_width = ( + depth // bin_size, + height // bin_size, + width // bin_size, + ) + if new_depth == 0 or new_height == 0 or new_width == 0: + raise ValueError( + f"bin_size {bin_size} is too large for stack dimensions " + f"{depth}x{height}x{width}" + ) + cropped = stack[ + : new_depth * bin_size, : new_height * bin_size, : new_width * bin_size + ] + blocks = cropped.reshape( + new_depth, bin_size, new_height, bin_size, new_width, bin_size + ) + binned = SpatialBin.for_enum_member(method).reduce(ops, blocks, (1, 3, 5)) + return ops.astype(binned, stack.dtype) + + +# ----- cropping, equalization, normalization (NumPy-compatible libraries) ----- + + +def crop( + ops: ArrayOperations, + image: Any, + start_x: int, + start_y: int, + start_z: int, + width: int, + height: int, + depth: int, +) -> Any: + """Return the ``(depth, height, width)`` substack at the start coordinates.""" + + require_stack(ops, image) + if width <= 0 or height <= 0 or depth <= 0: + raise ValueError( + "Crop dimensions must be positive: " + f"width={width}, height={height}, depth={depth}" + ) + if start_x < 0 or start_y < 0 or start_z < 0: + raise ValueError( + "Start coordinates must be non-negative: " + f"start_x={start_x}, start_y={start_y}, start_z={start_z}" + ) + input_depth, input_height, input_width = image.shape + end_x, end_y, end_z = start_x + width, start_y + height, start_z + depth + if end_x > input_width or end_y > input_height or end_z > input_depth: + raise ValueError( + "Crop region extends beyond image bounds. " + f"Image shape: {image.shape}, " + f"Crop region: ({start_z}:{end_z}, {start_y}:{end_y}, {start_x}:{end_x})" + ) + return image[start_z:end_z, start_y:end_y, start_x:end_x] + + +def stack_equalize_histogram( + ops: ArrayOperations, stack: Any, bins: int, range_min: float, range_max: float +) -> Any: + """Equalize with one histogram over the whole stack.""" + + require_stack(ops, stack) + xp = ops.array_module + input_dtype = stack.dtype + is_integer = xp.issubdtype(input_dtype, xp.integer) + hist, bin_edges = xp.histogram( + stack.flatten(), bins=bins, range=(range_min, range_max) + ) + cdf = hist.cumsum() + if cdf[-1] > 0: + if is_integer: + cdf = xp.iinfo(input_dtype).max * cdf / cdf[-1] + else: + cdf = cdf / cdf[-1] + equalized = xp.interp(stack.flatten(), bin_edges[:-1], cdf).reshape(stack.shape) + if is_integer: + limits = xp.iinfo(input_dtype) + return xp.clip(equalized, limits.min, limits.max).astype(input_dtype) + return equalized.astype(input_dtype) + + +def dtype_target_range(dtype: Any) -> tuple[float, float]: + """The natural intensity range of ``dtype``: integer limits, or [0, 1].""" + + np_dtype = np.dtype(dtype) + if np.issubdtype(np_dtype, np.integer): + limits = np.iinfo(np_dtype) + return limits.min, limits.max + if np_dtype in (np.float16, np.float32, np.float64): + return 0.0, 1.0 + return 0, 65535 + + +def _target_range( + dtype: Any, target_min: float | None, target_max: float | None +) -> tuple[float, float]: + natural_min, natural_max = dtype_target_range(dtype) + return ( + natural_min if target_min is None else target_min, + natural_max if target_max is None else target_max, + ) + + +def percentile_normalize_planes( + ops: ArrayOperations, + image: Any, + low_percentile: float, + high_percentile: float, + target_min: float | None, + target_max: float | None, +) -> Any: + """Stretch each plane's percentile range to the target range.""" + + require_stack(ops, image) + xp = ops.array_module + target_min, target_max = _target_range(image.dtype, target_min, target_max) + result = xp.zeros_like(image, dtype=np.float32) + for z in range(image.shape[0]): + p_low, p_high = xp.percentile(image[z], (low_percentile, high_percentile)) + if p_high == p_low: + result[z] = xp.ones_like(image[z]) * target_min + continue + clipped = xp.clip(image[z], p_low, p_high) + result[z] = (clipped - p_low) * (target_max - target_min) / ( + p_high - p_low + ) + target_min + return result.astype(image.dtype) + + +def percentile_normalize_stack( + ops: ArrayOperations, + stack: Any, + low_percentile: float, + high_percentile: float, + target_min: float | None, + target_max: float | None, +) -> Any: + """Stretch the whole stack's percentile range to the target range.""" + + require_stack(ops, stack) + xp = ops.array_module + target_min, target_max = _target_range(stack.dtype, target_min, target_max) + p_low = xp.percentile(stack, low_percentile) + p_high = xp.percentile(stack, high_percentile) + if p_high == p_low: + return (xp.ones_like(stack) * target_min).astype(stack.dtype) + clipped = xp.clip(stack, p_low, p_high) + normalized = (clipped - p_low) * (target_max - target_min) / ( + p_high - p_low + ) + target_min + return normalized.astype(stack.dtype) + + +# ----- composites ----- + + +def composite_weights(weights: Sequence[float] | None, count: int) -> list[float]: + """Validate per-plane weights and normalize them to sum to one.""" + + if weights is None: + weights = [1.0 / count] * count + elif isinstance(weights, (list, tuple)): + weights = list(weights) + if len(weights) != count: + raise ValueError( + f"Number of weights ({len(weights)}) must match number of slices ({count})" + ) + else: + raise TypeError( + f"weights must be a list of values or None, got {type(weights)}: {weights}" + ) + total = sum(weights) + if total == 0: + raise ValueError("Sum of weights cannot be zero") + return [weight / total for weight in weights] + + +def weighted_composite( + ops: ArrayOperations, stack: Any, weights: Sequence[float] | None +) -> Any: + """Float32 weighted sum of the planes, shaped ``(1, Y, X)``, in the input dtype.""" + + require_stack(ops, stack) + count = stack.shape[0] + plane_weights = ops.astype( + ops.array_module.asarray(composite_weights(weights, count)), np.float32 + ).reshape(count, 1, 1) + weighted = ops.astype(stack, np.float32) * plane_weights + return ops.astype(ops.prepend_axis(ops.sum(weighted, 0)), stack.dtype) diff --git a/openhcs/processing/backends/processors/jax_processor.py b/openhcs/processing/backends/processors/jax_processor.py index d0b3993ee5..a5cbf5f481 100644 --- a/openhcs/processing/backends/processors/jax_processor.py +++ b/openhcs/processing/backends/processors/jax_processor.py @@ -1,38 +1,30 @@ -""" -JAX Image Processor Implementation - -This module implements the ImageProcessorInterface using JAX as the backend. -It leverages GPU acceleration for image processing operations. +"""JAX processor registrations. -Doctrinal Clauses: -- Clause 3 — Declarative Primacy: All functions are pure and stateless -- Clause 88 — No Inferred Capabilities: Explicit JAX dependency -- Clause 106-A — Declared Memory Types: All methods specify JAX arrays +Shared operations are written once in ``image_stack_operations``; this module +registers them for JAX and holds the JAX-specific algorithms. """ from __future__ import annotations -import logging -from abc import abstractmethod -from typing import Any, List, Optional, Tuple +from typing import List, Optional, Tuple -from metaclass_registry import AutoRegisterMeta +from arraybridge import MemoryType +from arraybridge.array_operations import ArrayOperations -# Import JAX as an optional dependency from openhcs.core.artifacts import ArtifactSpec, ImageArtifactType -from openhcs.core.lazy_gpu_imports import jax +from openhcs.utils.import_utils import optional_import_placeholder from openhcs.core.memory import jax as jax_func from openhcs.core.pipeline.function_contracts import artifact_inputs -from metaclass_registry.strategies import EnumKeyedStrategyMixin +from openhcs.processing.backends.processors import image_stack_operations as operations from openhcs.processing.backends.processors.method_axes import ( StackProjectionMethod, ) -from openhcs.utils.import_utils import optional_import_placeholder -jnp = optional_import_placeholder("jax.numpy") if jax else None -lax = jax.lax if jax else None +jax = optional_import_placeholder("jax") +jnp = optional_import_placeholder("jax.numpy") +lax = optional_import_placeholder("jax.lax") -logger = logging.getLogger(__name__) +OPERATIONS = ArrayOperations.for_memory(MemoryType.JAX) @jax_func @@ -50,56 +42,9 @@ def create_linear_weight_mask( Returns: 2D JAX weight mask of shape (height, width) """ - # The compiler will ensure this function is only called when JAX is available - # No need to check for JAX availability here - - margin_y = int(jnp.floor(height * margin_ratio)) - margin_x = int(jnp.floor(width * margin_ratio)) - - weight_h = jnp.ones(height, dtype=jnp.float32) - if margin_y > 0: - ramp_top = jnp.linspace(0, 1, margin_y, endpoint=False) - ramp_bottom = jnp.linspace(1, 0, margin_y, endpoint=False) - weight_h = weight_h.at[:margin_y].set(ramp_top) - weight_h = weight_h.at[-margin_y:].set(ramp_bottom) - - weight_x = jnp.ones(width, dtype=jnp.float32) - if margin_x > 0: - ramp_left = jnp.linspace(0, 1, margin_x, endpoint=False) - ramp_right = jnp.linspace(1, 0, margin_x, endpoint=False) - weight_x = weight_x.at[:margin_x].set(ramp_left) - weight_x = weight_x.at[-margin_x:].set(ramp_right) - - # Create 2D weight mask using outer product - weight_mask = jnp.outer(weight_h, weight_x) - - return weight_mask - - -def _validate_3d_array(array: Any, name: str = "input") -> None: - """ - Validate that the input is a 3D JAX array. - - Args: - array: Array to validate - name: Name of the array for error messages - - Raises: - TypeError: If the array is not a JAX array - ValueError: If the array is not 3D - ImportError: If JAX is not available - """ - # The compiler will ensure this function is only called when JAX is available - # No need to check for JAX availability here - - if not isinstance(array, jnp.ndarray): - raise TypeError( - f"{name} must be a JAX array, got {type(array)}. " - f"No automatic conversion is performed to maintain explicit contracts." - ) - - if array.ndim != 3: - raise ValueError(f"{name} must be a 3D array, got {array.ndim}D") + return operations.linear_weight_mask( + OPERATIONS, height, width, margin_ratio, ramp_endpoint=False + ) @jax_func @@ -191,7 +136,7 @@ def sharpen( Returns: Sharpened 3D JAX array of shape (Z, Y, X) """ - _validate_3d_array(image) + operations.require_stack(OPERATIONS, image) # Store original dtype dtype = image.dtype @@ -255,7 +200,7 @@ def percentile_normalize( Returns: Normalized 3D JAX array of shape (Z, Y, X) """ - _validate_3d_array(image) + operations.require_stack(OPERATIONS, image) # Process each Z-slice independently result_list = [] @@ -332,7 +277,7 @@ def stack_percentile_normalize( Returns: Normalized 3D JAX array of shape (Z, Y, X) """ - _validate_3d_array(stack) + operations.require_stack(OPERATIONS, stack) # Calculate global percentiles across the entire stack p_low = jnp.percentile(stack, low_percentile) @@ -366,52 +311,7 @@ def create_composite( Returns: Composite 3D JAX array of shape (1, Y, X) """ - # Validate input is 3D array - _validate_3d_array(stack) - - n_slices, height, width = stack.shape - - # Default weights if none provided - if weights is None: - # Equal weights for all slices - weights = [1.0 / n_slices] * n_slices - elif isinstance(weights, (list, tuple)): - # Convert tuple to list if needed - weights = list(weights) - if len(weights) != n_slices: - raise ValueError( - f"Number of weights ({len(weights)}) must match number of slices ({n_slices})" - ) - else: - raise TypeError( - f"weights must be a list of values or None, got {type(weights)}: {weights}" - ) - - # Normalize weights to sum to 1 - weight_sum = sum(weights) - if weight_sum == 0: - raise ValueError("Sum of weights cannot be zero") - normalized_weights = [w / weight_sum for w in weights] - - # Convert weights to JAX array for efficient computation - # CRITICAL: Use float32 for weights to preserve fractional values, not stack.dtype - weights_array = jnp.array(normalized_weights, dtype=jnp.float32) - - # Reshape weights for broadcasting: (N, 1, 1) to multiply with (N, Y, X) - weights_array = weights_array.reshape(n_slices, 1, 1) - - # Create composite by weighted sum along the first axis - # Convert stack to float32 for computation to avoid precision loss - stack_float = stack.astype(jnp.float32) - weighted_stack = stack_float * weights_array - composite_slice = jnp.sum( - weighted_stack, axis=0, keepdims=True - ) # Keep as (1, Y, X) - - # Convert back to original dtype - composite_slice = composite_slice.astype(stack.dtype) - - return composite_slice + return operations.weighted_composite(OPERATIONS, stack, weights) @artifact_inputs(ArtifactSpec.input("mask", ImageArtifactType, parameter_name="mask")) @@ -430,36 +330,7 @@ def apply_mask(image: "jnp.ndarray", mask: "jnp.ndarray") -> "jnp.ndarray": Returns: Masked 3D JAX array of shape (Z, Y, X) """ - _validate_3d_array(image) - - # Handle 2D mask (apply to each Z-slice) - if isinstance(mask, jnp.ndarray) and mask.ndim == 2: - if mask.shape != image.shape[1:]: - raise ValueError( - f"2D mask shape {mask.shape} doesn't match image slice shape {image.shape[1:]}" - ) - - # Apply 2D mask to each Z-slice - result_list = [] - for z in range(image.shape[0]): - result_list.append(image[z].astype(jnp.float32) * mask.astype(jnp.float32)) - - result = jnp.stack(result_list, axis=0) - return result.astype(image.dtype) - - # Handle 3D mask - if isinstance(mask, jnp.ndarray) and mask.ndim == 3: - if mask.shape != image.shape: - raise ValueError( - f"3D mask shape {mask.shape} doesn't match image shape {image.shape}" - ) - - # Apply 3D mask directly - masked = image.astype(jnp.float32) * mask.astype(jnp.float32) - return masked.astype(image.dtype) - - # If we get here, the mask is neither 2D nor 3D JAX array - raise TypeError(f"mask must be a 2D or 3D JAX array, got {type(mask)}") + return operations.apply_mask(OPERATIONS, image, mask) @jax_func @@ -476,11 +347,7 @@ def create_weight_mask( Returns: 2D JAX weight mask of shape (Y, X) """ - if not isinstance(shape, tuple) or len(shape) != 2: - raise TypeError("shape must be a tuple of (height, width)") - - height, width = shape - return create_linear_weight_mask(height, width, margin_ratio) + return create_linear_weight_mask(*operations.weight_mask_extent(shape), margin_ratio) @jax_func @@ -494,11 +361,7 @@ def max_projection(stack: "jnp.ndarray") -> "jnp.ndarray": Returns: 3D JAX array of shape (1, Y, X) """ - _validate_3d_array(stack) - - # Create max projection - projection_2d = jnp.max(stack, axis=0) - return jnp.expand_dims(projection_2d, axis=0) + return operations.max_projection(OPERATIONS, stack) @jax_func @@ -512,11 +375,7 @@ def mean_projection(stack: "jnp.ndarray") -> "jnp.ndarray": Returns: 3D JAX array of shape (1, Y, X) """ - _validate_3d_array(stack) - - # Create mean projection - projection_2d = jnp.mean(stack.astype(jnp.float32), axis=0).astype(stack.dtype) - return jnp.expand_dims(projection_2d, axis=0) + return operations.mean_projection(OPERATIONS, stack) @jax_func @@ -541,7 +400,7 @@ def stack_equalize_histogram( Returns: Equalized 3D JAX array of shape (Z, Y, X) """ - _validate_3d_array(stack) + operations.require_stack(OPERATIONS, stack) # Remember input dtype to preserve it input_dtype = stack.dtype @@ -591,33 +450,6 @@ def stack_equalize_histogram( return equalized_stack.astype(input_dtype) -class JaxStackProjectionStrategy( - EnumKeyedStrategyMixin[StackProjectionMethod], - metaclass=AutoRegisterMeta, -): - __enum_member_attr__ = "method" - - """JAX stack-projection dispatch owner.""" - - @abstractmethod - def apply(self, stack: "jnp.ndarray") -> "jnp.ndarray": - """Apply this projection method.""" - - -class JaxMaxStackProjectionStrategy(JaxStackProjectionStrategy): - method = StackProjectionMethod.MAX - - def apply(self, stack: "jnp.ndarray") -> "jnp.ndarray": - return max_projection(stack) - - -class JaxMeanStackProjectionStrategy(JaxStackProjectionStrategy): - method = StackProjectionMethod.MEAN - - def apply(self, stack: "jnp.ndarray") -> "jnp.ndarray": - return mean_projection(stack) - - @jax_func def create_projection( stack: "jnp.ndarray", @@ -628,14 +460,12 @@ def create_projection( Args: stack: 3D JAX array of shape (Z, Y, X) - method: Projection method (max_projection, mean_projection) + method: Projection method served by ``image_stack_operations.StackProjection`` Returns: 3D JAX array of shape (1, Y, X) """ - _validate_3d_array(stack) - - return JaxStackProjectionStrategy.for_enum_member(method).apply(stack) + return operations.create_projection(OPERATIONS, stack, method) @jax_func @@ -656,7 +486,7 @@ def tophat( Returns: Filtered 3D JAX array of shape (Z, Y, X) """ - _validate_3d_array(image) + operations.require_stack(OPERATIONS, image) # Process each Z-slice independently result_list = [] diff --git a/openhcs/processing/backends/processors/method_axes.py b/openhcs/processing/backends/processors/method_axes.py index 58d7847411..73a5eb51a1 100644 --- a/openhcs/processing/backends/processors/method_axes.py +++ b/openhcs/processing/backends/processors/method_axes.py @@ -15,10 +15,11 @@ class SpatialBinMethod(Enum): class StackProjectionMethod(Enum): - """Max/mean reductions shared by processor backends.""" + """Leading-axis projections served by ``image_stack_operations.StackProjection``.""" MAX = "max_projection" MEAN = "mean_projection" + MIN = "min_projection" class EdgeMagnitudeMethod(Enum): diff --git a/openhcs/processing/backends/processors/numpy_processor.py b/openhcs/processing/backends/processors/numpy_processor.py index 02c288cadd..5b41c859de 100644 --- a/openhcs/processing/backends/processors/numpy_processor.py +++ b/openhcs/processing/backends/processors/numpy_processor.py @@ -1,22 +1,17 @@ -""" -NumPy Image Processor Implementation - -This module implements the ImageProcessorInterface using NumPy as the backend. -It serves as the reference implementation and works on CPU. +"""NumPy processor registrations. -Doctrinal Clauses: -- Clause 3 — Declarative Primacy: All functions are pure and stateless -- Clause 88 — No Inferred Capabilities: Explicit NumPy dependency -- Clause 106-A — Declared Memory Types: All methods specify NumPy arrays +Shared operations are written once in ``image_stack_operations``; this module +registers them for NumPy and holds the operations whose NumPy algorithm +(scikit-image, the numba composite kernel) differs from other backends. """ from __future__ import annotations -import logging -from abc import abstractmethod from enum import Enum -from typing import Annotated, Any, List, Optional, Tuple +from typing import Annotated, List, Optional, Tuple +from arraybridge import MemoryType +from arraybridge.array_operations import ArrayOperations from metaclass_registry import AutoRegisterMeta import numpy as np from numba import njit @@ -25,13 +20,10 @@ from skimage import transform as trans from openhcs.core.artifacts import ArtifactSpec, ImageArtifactType - -# Use direct import from core memory decorators to avoid circular imports from openhcs.core.memory import numpy as numpy_func from openhcs.core.pipeline.function_contracts import artifact_inputs -from metaclass_registry.strategies import EnumKeyedStrategyMixin -from openhcs.core.runtime_array_values import RuntimeArrayPayload from openhcs.core.processing_preparation import RegisteredNumbaKernelPreparation +from openhcs.processing.backends.processors import image_stack_operations as operations from openhcs.processing.backends.processors.method_axes import ( OrthogonalProjectionPlane, SpatialBinMethod, @@ -41,7 +33,7 @@ VolumetricToSliceContract, ) -logger = logging.getLogger(__name__) +OPERATIONS = ArrayOperations.for_memory(MemoryType.NUMPY) @numpy_func(contract=Pure2DContract) @@ -120,44 +112,6 @@ def non_local_means_denoise_planes( ] -class NumpySpatialBinStrategy( - EnumKeyedStrategyMixin[SpatialBinMethod], metaclass=AutoRegisterMeta -): - __enum_member_attr__ = "method" - - @abstractmethod - def apply(self, array: np.ndarray, axis: tuple[int, ...]) -> np.ndarray: - raise NotImplementedError - - -class NumpyMeanSpatialBinStrategy(NumpySpatialBinStrategy): - method = SpatialBinMethod.MEAN - - def apply(self, array: np.ndarray, axis: tuple[int, ...]) -> np.ndarray: - return np.mean(array, axis=axis) - - -class NumpySumSpatialBinStrategy(NumpySpatialBinStrategy): - method = SpatialBinMethod.SUM - - def apply(self, array: np.ndarray, axis: tuple[int, ...]) -> np.ndarray: - return np.sum(array, axis=axis) - - -class NumpyMaxSpatialBinStrategy(NumpySpatialBinStrategy): - method = SpatialBinMethod.MAX - - def apply(self, array: np.ndarray, axis: tuple[int, ...]) -> np.ndarray: - return np.max(array, axis=axis) - - -class NumpyMinSpatialBinStrategy(NumpySpatialBinStrategy): - method = SpatialBinMethod.MIN - - def apply(self, array: np.ndarray, axis: tuple[int, ...]) -> np.ndarray: - return np.min(array, axis=axis) - - @numpy_func def create_linear_weight_mask( height: int, width: int, margin_ratio: float = 0.1 @@ -173,46 +127,9 @@ def create_linear_weight_mask( Returns: 2D weight mask of shape (height, width) """ - margin_y = int(np.floor(height * margin_ratio)) - margin_x = int(np.floor(width * margin_ratio)) - - weight_y = np.ones(height, dtype=np.float32) - if margin_y > 0: - ramp_top = np.linspace(0, 1, margin_y, endpoint=False) - ramp_bottom = np.linspace(1, 0, margin_y, endpoint=False) - weight_y[:margin_y] = ramp_top - weight_y[-margin_y:] = ramp_bottom - - weight_x = np.ones(width, dtype=np.float32) - if margin_x > 0: - ramp_left = np.linspace(0, 1, margin_x, endpoint=False) - ramp_right = np.linspace(1, 0, margin_x, endpoint=False) - weight_x[:margin_x] = ramp_left - weight_x[-margin_x:] = ramp_right - - # Create 2D weight mask - weight_mask = np.outer(weight_y, weight_x) - - return weight_mask - - -def _validate_3d_array(array: Any, name: str = "input") -> None: - """ - Validate that the input is a 3D NumPy array. - - Args: - array: Array to validate - name: Name of the array for error messages - - Raises: - TypeError: If the array is not a NumPy array - ValueError: If the array is not 3D - """ - if not isinstance(array, (np.ndarray, RuntimeArrayPayload)): - raise TypeError(f"{name} must be a NumPy array, got {type(array)}") - - if array.ndim != 3: - raise ValueError(f"{name} must be a 3D array, got {array.ndim}D") + return operations.linear_weight_mask( + OPERATIONS, height, width, margin_ratio, ramp_endpoint=False + ) @numpy_func @@ -229,7 +146,7 @@ def sharpen(image: np.ndarray, radius: float = 1.0, amount: float = 1.0) -> np.n Returns: Sharpened 3D NumPy array of shape (Z, Y, X) """ - _validate_3d_array(image) + operations.require_stack(OPERATIONS, image) # Store original dtype dtype = image.dtype @@ -288,27 +205,8 @@ def percentile_normalize( Returns: Normalized 3D NumPy array of shape (N, Y, X). """ - _validate_3d_array(image) - - # Import shared utilities - from .percentile_utils import resolve_target_range, slice_percentile_normalize_core - - # Auto-detect target range based on input dtype if not specified - target_min, target_max = resolve_target_range(image.dtype, target_min, target_max) - - # Use shared core logic with NumPy-specific functions - return slice_percentile_normalize_core( - image=image, - low_percentile=low_percentile, - high_percentile=high_percentile, - target_min=target_min, - target_max=target_max, - percentile_func=np.percentile, - clip_func=np.clip, - ones_like_func=np.ones_like, - zeros_like_func=lambda arr, dtype=None: np.zeros_like( - arr, dtype=dtype or np.float32 - ), + return operations.percentile_normalize_planes( + OPERATIONS, image, low_percentile, high_percentile, target_min, target_max ) @@ -336,7 +234,7 @@ def percentile_normalize_plane( plane_index: Leading-axis plane to normalize; zero selects the first plane. """ - _validate_3d_array(image) + operations.require_stack(OPERATIONS, image) if not 0 <= plane_index < image.shape[0]: raise ValueError("plane_index is outside the input stack") @@ -381,24 +279,8 @@ def stack_percentile_normalize( Returns: Normalized 3D NumPy array of shape (N, Y, X). """ - _validate_3d_array(stack) - - # Import shared utilities - from .percentile_utils import resolve_target_range, percentile_normalize_core - - # Auto-detect target range based on input dtype if not specified - target_min, target_max = resolve_target_range(stack.dtype, target_min, target_max) - - # Use shared core logic with NumPy-specific functions - return percentile_normalize_core( - stack=stack, - low_percentile=low_percentile, - high_percentile=high_percentile, - target_min=target_min, - target_max=target_max, - percentile_func=lambda arr, pct: np.percentile(arr, pct, axis=None), - clip_func=np.clip, - ones_like_func=np.ones_like, + return operations.percentile_normalize_stack( + OPERATIONS, stack, low_percentile, high_percentile, target_min, target_max ) @@ -508,32 +390,9 @@ def create_composite( Composite 2D NumPy array of shape (Y, X). OpenHCS contract execution restores the singleton runtime stack axis. """ - # Validate input is 3D array - _validate_3d_array(stack) - - n_slices, height, width = stack.shape - - # Default weights if none provided - if weights is None: - # Equal weights forangeit so that it can only all slices - weights = [1.0 / n_slices] * n_slices - elif isinstance(weights, (list, tuple)): - # Convert tuple to list if needed - weights = list(weights) - if len(weights) != n_slices: - raise ValueError( - f"Number of weights ({len(weights)}) must match number of slices ({n_slices})" - ) - else: - raise TypeError( - f"weights must be a list of values or None, got {type(weights)}: {weights}" - ) - - # Normalize weights to sum to 1 - weight_sum = sum(weights) - if weight_sum == 0: - raise ValueError("Sum of weights cannot be zero") - normalized_weights = [w / weight_sum for w in weights] + operations.require_stack(OPERATIONS, stack) + n_slices = stack.shape[0] + normalized_weights = operations.composite_weights(weights, n_slices) # Convert weights to NumPy array for efficient computation # CRITICAL: Use float32 for weights to preserve fractional values, not stack.dtype @@ -565,35 +424,7 @@ def apply_mask(image: np.ndarray, mask: np.ndarray) -> np.ndarray: Returns: Masked 3D NumPy array of shape (Z, Y, X) """ - _validate_3d_array(image) - - # Handle 2D mask (apply to each Z-slice) - if isinstance(mask, np.ndarray) and mask.ndim == 2: - if mask.shape != image.shape[1:]: - raise ValueError( - f"2D mask shape {mask.shape} doesn't match image slice shape {image.shape[1:]}" - ) - - # Apply 2D mask to each Z-slice - result = np.zeros_like(image) - for z in range(image.shape[0]): - result[z] = image[z].astype(np.float32) * mask.astype(np.float32) - - return result.astype(image.dtype) - - # Handle 3D mask - if isinstance(mask, np.ndarray) and mask.ndim == 3: - if mask.shape != image.shape: - raise ValueError( - f"3D mask shape {mask.shape} doesn't match image shape {image.shape}" - ) - - # Apply 3D mask directly - masked = image.astype(np.float32) * mask.astype(np.float32) - return masked.astype(image.dtype) - - # If we get here, the mask is neither 2D nor 3D NumPy array - raise TypeError(f"mask must be a 2D or 3D NumPy array, got {type(mask)}") + return operations.apply_mask(OPERATIONS, image, mask) @numpy_func @@ -608,11 +439,7 @@ def create_weight_mask(shape: Tuple[int, int], margin_ratio: float = 0.1) -> np. Returns: 2D NumPy weight mask of shape (Y, X) """ - if not isinstance(shape, tuple) or len(shape) != 2: - raise TypeError("shape must be a tuple of (height, width)") - - height, width = shape - return create_linear_weight_mask(height, width, margin_ratio) + return create_linear_weight_mask(*operations.weight_mask_extent(shape), margin_ratio) @numpy_func @@ -623,11 +450,7 @@ def max_projection(stack: np.ndarray) -> np.ndarray: Returns: 3D NumPy array of shape (1, Y, X) """ - _validate_3d_array(stack) - - # Create max projection - projection_2d = np.max(stack, axis=0) - return projection_2d.reshape(1, projection_2d.shape[0], projection_2d.shape[1]) + return operations.max_projection(OPERATIONS, stack) @numpy_func @@ -638,61 +461,20 @@ def mean_projection(stack: np.ndarray) -> np.ndarray: Returns: 3D NumPy array of shape (1, Y, X) """ - _validate_3d_array(stack) - - # Create mean projection - projection_2d = np.mean(stack, axis=0).astype(stack.dtype) - return projection_2d.reshape(1, projection_2d.shape[0], projection_2d.shape[1]) - - -def min_projection(stack: np.ndarray) -> np.ndarray: - """Create a minimum intensity projection from a Z-stack.""" - - _validate_3d_array(stack) - projection_2d = np.min(stack, axis=0) - return projection_2d.reshape(1, *projection_2d.shape) + return operations.mean_projection(OPERATIONS, stack) class NumpyStackProjectionMethod(Enum): - """Stack reductions implemented by the NumPy projection dispatcher.""" + """NumPy ``create_projection`` choices; values name ``StackProjection`` members. + + Kept as its own class because saved pipelines import it from this module. + """ MAX = "max_projection" MEAN = "mean_projection" MIN = "min_projection" -class NumpyStackProjectionStrategy( - EnumKeyedStrategyMixin[NumpyStackProjectionMethod], - metaclass=AutoRegisterMeta, -): - __enum_member_attr__ = "method" - - @abstractmethod - def apply(self, stack: np.ndarray) -> np.ndarray: - raise NotImplementedError - - -class NumpyMaxStackProjectionStrategy(NumpyStackProjectionStrategy): - method = NumpyStackProjectionMethod.MAX - - def apply(self, stack: np.ndarray) -> np.ndarray: - return max_projection(stack) - - -class NumpyMeanStackProjectionStrategy(NumpyStackProjectionStrategy): - method = NumpyStackProjectionMethod.MEAN - - def apply(self, stack: np.ndarray) -> np.ndarray: - return mean_projection(stack) - - -class NumpyMinStackProjectionStrategy(NumpyStackProjectionStrategy): - method = NumpyStackProjectionMethod.MIN - - def apply(self, stack: np.ndarray) -> np.ndarray: - return min_projection(stack) - - @numpy_func def create_orthogonal_projections( stack: np.ndarray, @@ -720,7 +502,7 @@ def create_orthogonal_projections( - Returns data, never performs I/O - Testable in isolation """ - _validate_3d_array(stack) + operations.require_stack(OPERATIONS, stack) result = {} if OrthogonalProjectionPlane.XY in projections: @@ -743,7 +525,7 @@ def gaussian_blur(stack: np.ndarray, sigma: float = 1.0) -> np.ndarray: Returns: Blurred 3D NumPy array of shape (Z, Y, X) """ - _validate_3d_array(stack) + operations.require_stack(OPERATIONS, stack) # Apply Gaussian blur slice-by-slice blurred = np.zeros_like(stack, dtype=np.float64) @@ -771,37 +553,7 @@ def spatial_bin_2d( Returns: Binned 3D NumPy array of shape (Z, Y//bin_size, X//bin_size) """ - _validate_3d_array(stack) - - if bin_size <= 0: - raise ValueError("bin_size must be positive") - - z_slices, height, width = stack.shape - - # Calculate output dimensions - new_height = height // bin_size - new_width = width // bin_size - - if new_height == 0 or new_width == 0: - raise ValueError( - f"bin_size {bin_size} is too large for image dimensions {height}x{width}" - ) - - # Crop to make dimensions divisible by bin_size - crop_height = new_height * bin_size - crop_width = new_width * bin_size - cropped_stack = stack[:, :crop_height, :crop_width] - - # Reshape for binning: (Z, new_height, bin_size, new_width, bin_size) - reshaped = cropped_stack.reshape( - z_slices, new_height, bin_size, new_width, bin_size - ) - - result = NumpySpatialBinStrategy.for_enum_member(method).apply( - reshaped, axis=(2, 4) - ) - - return result.astype(stack.dtype) + return operations.spatial_bin_2d(OPERATIONS, stack, bin_size, method) @numpy_func @@ -821,39 +573,7 @@ def spatial_bin_3d( Returns: Binned 3D NumPy array of shape (Z//bin_size, Y//bin_size, X//bin_size) """ - _validate_3d_array(stack) - - if bin_size <= 0: - raise ValueError("bin_size must be positive") - - depth, height, width = stack.shape - - # Calculate output dimensions - new_depth = depth // bin_size - new_height = height // bin_size - new_width = width // bin_size - - if new_depth == 0 or new_height == 0 or new_width == 0: - raise ValueError( - f"bin_size {bin_size} is too large for stack dimensions {depth}x{height}x{width}" - ) - - # Crop to make dimensions divisible by bin_size - crop_depth = new_depth * bin_size - crop_height = new_height * bin_size - crop_width = new_width * bin_size - cropped_stack = stack[:crop_depth, :crop_height, :crop_width] - - # Reshape for 3D binning: (new_depth, bin_size, new_height, bin_size, new_width, bin_size) - reshaped = cropped_stack.reshape( - new_depth, bin_size, new_height, bin_size, new_width, bin_size - ) - - result = NumpySpatialBinStrategy.for_enum_member(method).apply( - reshaped, axis=(1, 3, 5) - ) - - return result.astype(stack.dtype) + return operations.spatial_bin_3d(OPERATIONS, stack, bin_size, method) @numpy_func @@ -877,42 +597,10 @@ def stack_equalize_histogram( Returns: Equalized 3D NumPy array of shape (Z, Y, X) """ - _validate_3d_array(stack) - - # Remember input dtype to preserve it - input_dtype = stack.dtype - - # Flatten the entire stack to compute the global histogram - flat_stack = stack.flatten() - - # Calculate the histogram and cumulative distribution function (CDF) - hist, bin_edges = np.histogram(flat_stack, bins=bins, range=(range_min, range_max)) - cdf = hist.cumsum() - - # Normalize the CDF to the input dtype range - # Avoid division by zero - if cdf[-1] > 0: - if np.issubdtype(input_dtype, np.integer): - dtype_info = np.iinfo(input_dtype) - cdf = dtype_info.max * cdf / cdf[-1] - else: - # For float dtypes, normalize to [0, 1] - cdf = cdf / cdf[-1] - - # Use linear interpolation to map input values to equalized values - equalized_stack = np.interp(stack.flatten(), bin_edges[:-1], cdf).reshape( - stack.shape + return operations.stack_equalize_histogram( + OPERATIONS, stack, bins, range_min, range_max ) - # Convert back to input dtype - if np.issubdtype(input_dtype, np.integer): - dtype_info = np.iinfo(input_dtype) - return np.clip(equalized_stack, dtype_info.min, dtype_info.max).astype( - input_dtype - ) - else: - return equalized_stack.astype(input_dtype) - @numpy_func(contract=VolumetricToSliceContract) def create_projection( @@ -923,15 +611,12 @@ def create_projection( Create a projection from a stack using the specified method. Args: - method: Projection method owned by ``NumpyStackProjectionStrategy``. + method: Projection method served by ``image_stack_operations.StackProjection``. Returns: 2D NumPy array of shape (Y, X) """ - stack_array = np.asarray(stack) - _validate_3d_array(stack_array) - - return NumpyStackProjectionStrategy.for_enum_member(method).apply(stack_array)[0] + return operations.create_projection(OPERATIONS, np.asarray(stack), method)[0] @numpy_func @@ -971,39 +656,9 @@ def crop( np.ndarray Cropped 3D array of shape (depth, height, width) """ - _validate_3d_array(input_image) - - # Validate crop parameters - if width <= 0 or height <= 0 or depth <= 0: - raise ValueError( - f"Crop dimensions must be positive: width={width}, height={height}, depth={depth}" - ) - - if start_x < 0 or start_y < 0 or start_z < 0: - raise ValueError( - f"Start coordinates must be non-negative: start_x={start_x}, start_y={start_y}, start_z={start_z}" - ) - - # Get input dimensions - input_depth, input_height, input_width = input_image.shape - - # Calculate end coordinates - end_x = start_x + width - end_y = start_y + height - end_z = start_z + depth - - # Validate bounds - if end_x > input_width or end_y > input_height or end_z > input_depth: - raise ValueError( - f"Crop region extends beyond image bounds. " - f"Image shape: {input_image.shape}, " - f"Crop region: ({start_z}:{end_z}, {start_y}:{end_y}, {start_x}:{end_x})" - ) - - # Perform the crop using NumPy slicing - cropped = input_image[start_z:end_z, start_y:end_y, start_x:end_x] - - return cropped + return operations.crop( + OPERATIONS, input_image, start_x, start_y, start_z, width, height, depth + ) @numpy_func @@ -1045,7 +700,7 @@ def tophat( Returns: Background-subtracted 3D NumPy array of shape (N, Y, X). """ - _validate_3d_array(image) + operations.require_stack(OPERATIONS, image) # Process each Z-slice independently result = np.zeros_like(image) diff --git a/openhcs/processing/backends/processors/percentile_utils.py b/openhcs/processing/backends/processors/percentile_utils.py deleted file mode 100644 index fb28b5992a..0000000000 --- a/openhcs/processing/backends/processors/percentile_utils.py +++ /dev/null @@ -1,187 +0,0 @@ -""" -Shared utilities for percentile normalization across all backends. - -This module provides common functionality to ensure consistent behavior -between NumPy, CuPy, PyTorch, JAX, TensorFlow, and other implementations. -""" - -import numpy as np -from typing import Tuple, Union - - -def get_dtype_range(dtype) -> Tuple[Union[int, float], Union[int, float]]: - """ - Get the natural min/max range for a numpy-compatible dtype. - - Args: - dtype: NumPy dtype or equivalent (works with CuPy, PyTorch, etc.) - - Returns: - Tuple of (min_value, max_value) for the dtype - """ - # Convert to numpy dtype for consistent comparison - if hasattr(dtype, "type"): - # Handle CuPy/PyTorch dtypes that have .type attribute - np_dtype = dtype.type - else: - # Handle direct numpy dtypes - np_dtype = dtype - - # Map dtypes to their natural ranges - if np_dtype == np.uint8: - return 0, 255 - elif np_dtype == np.uint16: - return 0, 65535 - elif np_dtype == np.uint32: - return 0, 4294967295 - elif np_dtype == np.uint64: - return 0, 18446744073709551615 - elif np_dtype == np.int8: - return -128, 127 - elif np_dtype == np.int16: - return -32768, 32767 - elif np_dtype == np.int32: - return -2147483648, 2147483647 - elif np_dtype == np.int64: - return -9223372036854775808, 9223372036854775807 - elif np_dtype in (np.float16, np.float32, np.float64): - return 0.0, 1.0 - else: - # Fallback for unknown dtypes - assume 16-bit range - return 0, 65535 - - -def resolve_target_range( - stack_dtype, target_min=None, target_max=None -) -> Tuple[Union[int, float], Union[int, float]]: - """ - Resolve target min/max values, auto-detecting from dtype if not specified. - - Args: - stack_dtype: The dtype of the input stack - target_min: Explicit target minimum (None for auto-detection) - target_max: Explicit target maximum (None for auto-detection) - - Returns: - Tuple of (resolved_min, resolved_max) - """ - if target_min is None or target_max is None: - auto_min, auto_max = get_dtype_range(stack_dtype) - if target_min is None: - target_min = auto_min - if target_max is None: - target_max = auto_max - - return target_min, target_max - - -def percentile_normalize_core( - stack, - low_percentile: float, - high_percentile: float, - target_min: Union[int, float], - target_max: Union[int, float], - percentile_func, - clip_func, - ones_like_func, - preserve_dtype: bool = True, -): - """ - Core percentile normalization logic that works with any array backend. - - This function contains the shared algorithm while allowing different backends - to provide their own array operations (percentile, clip, ones_like). - - Args: - stack: Input array (NumPy, CuPy, PyTorch, etc.) - low_percentile: Lower percentile (0-100) - high_percentile: Upper percentile (0-100) - target_min: Target minimum value - target_max: Target maximum value - percentile_func: Backend-specific percentile function - clip_func: Backend-specific clip function - ones_like_func: Backend-specific ones_like function - preserve_dtype: Whether to preserve input dtype - - Returns: - Normalized array with same backend as input - """ - # Calculate global percentiles across the entire stack - p_low = percentile_func(stack, low_percentile) - p_high = percentile_func(stack, high_percentile) - - # Avoid division by zero - if p_high == p_low: - result = ones_like_func(stack) * target_min - if preserve_dtype: - return result.astype(stack.dtype) - else: - # Legacy behavior: convert to uint16-equivalent - return result.astype( - stack.dtype if hasattr(stack, "dtype") else type(stack) - ) - - # Clip and normalize to target range - clipped = clip_func(stack, p_low, p_high) - normalized = (clipped - p_low) * (target_max - target_min) / ( - p_high - p_low - ) + target_min - - # Always preserve input dtype - return normalized.astype(stack.dtype) - - -def slice_percentile_normalize_core( - image, - low_percentile: float, - high_percentile: float, - target_min: Union[int, float], - target_max: Union[int, float], - percentile_func, - clip_func, - ones_like_func, - zeros_like_func, - preserve_dtype: bool = True, -): - """ - Core slice-by-slice percentile normalization logic. - - Args: - image: Input 3D array (Z, Y, X) - low_percentile: Lower percentile (0-100) - high_percentile: Upper percentile (0-100) - target_min: Target minimum value - target_max: Target maximum value - percentile_func: Backend-specific percentile function - clip_func: Backend-specific clip function - ones_like_func: Backend-specific ones_like function - zeros_like_func: Backend-specific zeros_like function - preserve_dtype: Whether to preserve input dtype - - Returns: - Normalized array with same backend as input - """ - # Process each Z-slice independently - # Use float32 for intermediate calculations to avoid precision loss - result = zeros_like_func( - image, dtype=np.float32 if hasattr(image, "dtype") else None - ) - - for z in range(image.shape[0]): - # Get percentile values for this slice - p_low, p_high = percentile_func(image[z], (low_percentile, high_percentile)) - - # Avoid division by zero - if p_high == p_low: - result[z] = ones_like_func(image[z]) * target_min - continue - - # Clip and normalize to target range - clipped = clip_func(image[z], p_low, p_high) - normalized = (clipped - p_low) * (target_max - target_min) / ( - p_high - p_low - ) + target_min - result[z] = normalized - - # Always preserve input dtype - return result.astype(image.dtype) diff --git a/openhcs/processing/backends/processors/pyclesperanto_processor.py b/openhcs/processing/backends/processors/pyclesperanto_processor.py index ce356294c8..deab9b2b81 100644 --- a/openhcs/processing/backends/processors/pyclesperanto_processor.py +++ b/openhcs/processing/backends/processors/pyclesperanto_processor.py @@ -1,114 +1,35 @@ -""" -pyclesperanto GPU processor for OpenHCS. +"""pyclesperanto processor registrations. -This processor uses pyclesperanto for GPU-accelerated image processing, -providing excellent compatibility with OpenCL devices and seamless integration -with OpenHCS patterns. +Shared operations are written once in ``image_stack_operations``; this module +registers them for pyclesperanto and holds the operations built from +pyclesperanto's own kernels (min/max normalization, CLAHE, top-hat sphere). """ -import logging -import os -from abc import abstractmethod -from enum import Enum from typing import List, Optional -from metaclass_registry import AutoRegisterMeta +from arraybridge import MemoryType +from arraybridge.array_operations import ArrayOperations -# Import OpenHCS decorator from openhcs.core.artifacts import ArtifactSpec, ImageArtifactType +from openhcs.utils.import_utils import optional_import_placeholder from openhcs.core.memory import pyclesperanto as pyclesperanto_func from openhcs.core.pipeline.function_contracts import artifact_inputs -from metaclass_registry.strategies import EnumKeyedStrategyMixin +from openhcs.processing.backends.processors import image_stack_operations as operations from openhcs.processing.backends.processors.method_axes import ( StackProjectionMethod, ) -# Set up logging -logger = logging.getLogger(__name__) - - -class PyclesperantoMaskDimensionality(Enum): - """Closed mask-dimensionality family for GPU mask application.""" - - MASK_2D = 2 - MASK_3D = 3 - - @classmethod - def from_mask(cls, mask: "cle.Array") -> "PyclesperantoMaskDimensionality": - try: - return cls(len(mask.shape)) - except ValueError as exc: - raise TypeError( - f"mask must be a 2D or 3D pyclesperanto Array, got shape {mask.shape}" - ) from exc - - def apply(self, image: "cle.Array", mask: "cle.Array") -> "cle.Array": - if self is PyclesperantoMaskDimensionality.MASK_2D: - if mask.shape != image.shape[1:]: - raise ValueError( - f"2D mask shape {mask.shape} doesn't match image slice shape " - f"{image.shape[1:]}" - ) - - result = cle.create_like(image) - for z in range(image.shape[0]): - result[z] = cle.multiply_images(image[z], mask) - return result - - if mask.shape != image.shape: - raise ValueError( - f"3D mask shape {mask.shape} doesn't match image shape {image.shape}" - ) - return cle.multiply_images(image, mask) - - -# Check if we're in subprocess runner mode and should skip GPU imports -if os.getenv("OPENHCS_SUBPROCESS_NO_GPU") == "1": - # Subprocess runner mode - skip GPU imports - PYCLESPERANTO_AVAILABLE = False - cle = None - logger.info("Subprocess runner mode - skipping pyclesperanto import") -else: - # Normal mode - try to import pyclesperanto - try: - import pyclesperanto as cle - - PYCLESPERANTO_AVAILABLE = True - logger.info("pyclesperanto available - GPU acceleration enabled") - except ImportError: - PYCLESPERANTO_AVAILABLE = False - cle = None - logger.warning( - "pyclesperanto not available - install with: pip install pyclesperanto" - ) - - -def _check_pyclesperanto_available(): - """Check if pyclesperanto is available and raise error if not.""" - if not PYCLESPERANTO_AVAILABLE: - raise ImportError( - "pyclesperanto is required but not available. Install with: pip install pyclesperanto" - ) +cle = optional_import_placeholder("pyclesperanto") - -def _validate_3d_array(array) -> None: - """Validate that the input is a 3D array.""" - if array.ndim != 3: - raise ValueError(f"Expected 3D array, got {array.ndim}D array") +OPERATIONS = ArrayOperations.for_memory(MemoryType.PYCLESPERANTO) def _gpu_minmax_normalize_range(image: "cle.Array") -> tuple: - """Calculate normalization range using min/max instead of percentiles - pure GPU.""" - import pyclesperanto as cle - - # Use min/max instead of percentiles to stay purely on GPU - # This is similar to how CLAHE works internally - min_val = cle.minimum_of_all_pixels(image) - max_val = cle.maximum_of_all_pixels(image) - - # For compatibility, we could apply a small margin based on percentile values - # but for pure GPU operation, we use full min/max range - return float(min_val), float(max_val) + """Return the image's (min, max), computed on the device.""" + return ( + float(cle.minimum_of_all_pixels(image)), + float(cle.maximum_of_all_pixels(image)), + ) @pyclesperanto_func @@ -137,9 +58,7 @@ def per_slice_minmax_normalize( Returns: Normalized 3D pyclesperanto Array of shape (Z, Y, X) with dtype uint16 """ - # Validate 3D array - if len(image.shape) != 3: - raise ValueError(f"Expected 3D array, got {len(image.shape)}D array") + operations.require_stack(OPERATIONS, image) # Build result by concatenating pairs of slices result_slices = [] @@ -200,11 +119,7 @@ def stack_minmax_normalize( Returns: Normalized 3D pyclesperanto Array of shape (Z, Y, X) with dtype uint16 """ - _check_pyclesperanto_available() - - # Validate 3D array - if len(image.shape) != 3: - raise ValueError(f"Expected 3D array, got {len(image.shape)}D array") + operations.require_stack(OPERATIONS, image) # Calculate global min/max range from entire stack (pure GPU) p_low, p_high = _gpu_minmax_normalize_range(image) @@ -247,11 +162,7 @@ def sharpen( Returns: Sharpened 3D pyclesperanto Array of shape (Z, Y, X) """ - _check_pyclesperanto_available() - - # Validate 3D array - if len(image.shape) != 3: - raise ValueError(f"Expected 3D array, got {len(image.shape)}D array") + operations.require_stack(OPERATIONS, image) # Apply 3D Gaussian blur (pyclesperanto handles Z-dimension efficiently) gpu_blurred = cle.gaussian_blur(image, sigma_x=radius, sigma_y=radius, sigma_z=0) @@ -284,21 +195,7 @@ def max_projection(stack: "cle.Array") -> "cle.Array": Returns: 3D pyclesperanto Array of shape (1, Y, X) """ - _check_pyclesperanto_available() - - # Validate 3D array - if len(stack.shape) != 3: - raise ValueError(f"Expected 3D array, got {len(stack.shape)}D array") - - # Create max projection (stays on GPU) - gpu_projection_2d = cle.maximum_z_projection(stack) - - # Reshape to (1, Y, X) by creating a new 3D array - result_shape = (1, gpu_projection_2d.shape[0], gpu_projection_2d.shape[1]) - result = cle.create(result_shape, dtype=gpu_projection_2d.dtype) - result[0] = gpu_projection_2d # Direct assignment - - return result + return operations.max_projection(OPERATIONS, stack) @pyclesperanto_func @@ -314,50 +211,11 @@ def mean_projection(stack: "cle.Array") -> "cle.Array": stack: 3D pyclesperanto Array of shape (Z, Y, X) Returns: - 3D pyclesperanto Array of shape (1, Y, X) + 3D pyclesperanto Array of shape (1, Y, X) holding pyclesperanto's + float32 z-mean (not cast back to the input dtype). """ - _check_pyclesperanto_available() - - # Validate 3D array - if len(stack.shape) != 3: - raise ValueError(f"Expected 3D array, got {len(stack.shape)}D array") - - # Create mean projection (stays on GPU) - gpu_projection_2d = cle.mean_z_projection(stack) - - # Reshape to (1, Y, X) by creating a new 3D array - result_shape = (1, gpu_projection_2d.shape[0], gpu_projection_2d.shape[1]) - result = cle.create(result_shape, dtype=gpu_projection_2d.dtype) - result[0] = gpu_projection_2d # Direct assignment - - return result - - -class PyclesperantoStackProjectionStrategy( - EnumKeyedStrategyMixin[StackProjectionMethod], - metaclass=AutoRegisterMeta, -): - __enum_member_attr__ = "method" - - """pyclesperanto stack-projection dispatch owner.""" - - @abstractmethod - def apply(self, stack: "cle.Array") -> "cle.Array": - """Apply this projection method.""" - - -class PyclesperantoMaxStackProjectionStrategy(PyclesperantoStackProjectionStrategy): - method = StackProjectionMethod.MAX - - def apply(self, stack: "cle.Array") -> "cle.Array": - return max_projection(stack) - - -class PyclesperantoMeanStackProjectionStrategy(PyclesperantoStackProjectionStrategy): - method = StackProjectionMethod.MEAN - - def apply(self, stack: "cle.Array") -> "cle.Array": - return mean_projection(stack) + operations.require_stack(OPERATIONS, stack) + return OPERATIONS.prepend_axis(OPERATIONS.mean(stack, 0)) @pyclesperanto_func @@ -373,18 +231,15 @@ def create_projection( Args: stack: 3D pyclesperanto Array of shape (Z, Y, X) - method: Projection method (max_projection, mean_projection) + method: Projection method served by ``image_stack_operations.StackProjection``; + MEAN keeps pyclesperanto's float32 z-mean (``mean_projection``). Returns: 3D pyclesperanto Array of shape (1, Y, X) """ - _check_pyclesperanto_available() - - # Validate 3D array - if len(stack.shape) != 3: - raise ValueError(f"Expected 3D array, got {len(stack.shape)}D array") - - return PyclesperantoStackProjectionStrategy.for_enum_member(method).apply(stack) + if method is StackProjectionMethod.MEAN: + return mean_projection(stack) + return operations.create_projection(OPERATIONS, stack, method) @pyclesperanto_func @@ -415,11 +270,7 @@ def tophat( Returns: Filtered 3D pyclesperanto Array of shape (Z, Y, X) """ - _check_pyclesperanto_available() - - # Validate 3D array - if len(image.shape) != 3: - raise ValueError(f"Expected 3D array, got {len(image.shape)}D array") + operations.require_stack(OPERATIONS, image) # 1) Downsample entire stack at once (more efficient) scale_factor = 1.0 / downsample_factor @@ -481,13 +332,25 @@ def apply_mask(image: "cle.Array", mask: "cle.Array") -> "cle.Array": Returns: Masked 3D pyclesperanto Array of shape (Z, Y, X) """ - _check_pyclesperanto_available() - - # Validate 3D image - if len(image.shape) != 3: - raise ValueError(f"Expected 3D image array, got {len(image.shape)}D array") - - return PyclesperantoMaskDimensionality.from_mask(mask).apply(image, mask) + operations.require_stack(OPERATIONS, image) + if len(mask.shape) == 3: + if mask.shape != image.shape: + raise ValueError( + f"3D mask shape {mask.shape} doesn't match image shape {image.shape}" + ) + return cle.multiply_images(image, mask) + if len(mask.shape) != 2: + raise TypeError( + f"mask must be a 2D or 3D pyclesperanto Array, got shape {mask.shape}" + ) + if mask.shape != image.shape[1:]: + raise ValueError( + f"2D mask shape {mask.shape} doesn't match image slice shape {image.shape[1:]}" + ) + result = cle.create_like(image) + for z in range(image.shape[0]): + result[z] = cle.multiply_images(image[z], mask) + return result @pyclesperanto_func @@ -508,34 +371,9 @@ def create_composite( Returns: Composite 3D pyclesperanto Array of shape (1, Y, X) """ - _check_pyclesperanto_available() - - # Validate input is 3D array - if len(stack.shape) != 3: - raise ValueError(f"Expected 3D array, got {len(stack.shape)}D array") - - n_slices, height, width = stack.shape - - # Default weights if none provided - if weights is None: - weights = [1.0 / n_slices] * n_slices - elif isinstance(weights, (list, tuple)): - # Convert tuple to list if needed - weights = list(weights) - if len(weights) != n_slices: - raise ValueError( - f"Number of weights ({len(weights)}) must match number of slices ({n_slices})" - ) - else: - raise TypeError( - f"weights must be a list of values or None, got {type(weights)}: {weights}" - ) - - # Normalize weights to sum to 1 - weight_sum = sum(weights) - if weight_sum == 0: - raise ValueError("Sum of weights cannot be zero") - normalized_weights = [w / weight_sum for w in weights] + operations.require_stack(OPERATIONS, stack) + _, height, width = stack.shape + normalized_weights = operations.composite_weights(weights, stack.shape[0]) # Create result array with shape (1, Y, X) result = cle.create((1, height, width), dtype=stack.dtype) @@ -588,11 +426,7 @@ def equalize_histogram_3d( Returns: Equalized 3D pyclesperanto Array of shape (Z, Y, X) """ - _check_pyclesperanto_available() - - # Validate 3D array - if len(stack.shape) != 3: - raise ValueError(f"Expected 3D array, got {len(stack.shape)}D array") + operations.require_stack(OPERATIONS, stack) # Use 3D CLAHE with 3D tiles (cubes) gpu_equalized = cle.clahe(stack, tile_size=tile_size, clip_limit=clip_limit) @@ -630,11 +464,7 @@ def equalize_histogram_per_slice( Returns: Equalized 3D pyclesperanto Array of shape (Z, Y, X) """ - _check_pyclesperanto_available() - - # Validate 3D array - if len(stack.shape) != 3: - raise ValueError(f"Expected 3D array, got {len(stack.shape)}D array") + operations.require_stack(OPERATIONS, stack) # Create result array result = cle.create_like(stack) @@ -812,8 +642,4 @@ def create_weight_mask(shape: tuple, margin_ratio: float = 0.1) -> "cle.Array": Returns: 2D pyclesperanto Array of shape (Y, X) """ - if not isinstance(shape, tuple) or len(shape) != 2: - raise TypeError("shape must be a tuple of (height, width)") - - height, width = shape - return create_linear_weight_mask(height, width, margin_ratio) + return create_linear_weight_mask(*operations.weight_mask_extent(shape), margin_ratio) diff --git a/openhcs/processing/backends/processors/tensorflow_processor.py b/openhcs/processing/backends/processors/tensorflow_processor.py index 0706ba23d0..f0b88191b3 100644 --- a/openhcs/processing/backends/processors/tensorflow_processor.py +++ b/openhcs/processing/backends/processors/tensorflow_processor.py @@ -1,54 +1,28 @@ -""" -TensorFlow Image Processor Implementation - -This module implements the ImageProcessorInterface using TensorFlow as the backend. -It leverages GPU acceleration for image processing operations. +"""TensorFlow processor registrations. -Doctrinal Clauses: -- Clause 3 — Declarative Primacy: All functions are pure and stateless -- Clause 88 — No Inferred Capabilities: Explicit TensorFlow dependency -- Clause 106-A — Declared Memory Types: All methods specify TensorFlow tensors +Shared operations are written once in ``image_stack_operations``; this module +registers them for TensorFlow and holds the TensorFlow-specific algorithms. """ from __future__ import annotations -import logging -from abc import abstractmethod -from typing import Any, List, Optional, Tuple +from typing import List, Optional, Tuple -from metaclass_registry import AutoRegisterMeta -from packaging.version import parse as parse_version +from arraybridge import MemoryType +from arraybridge.array_operations import ArrayOperations from openhcs.core.artifacts import ArtifactSpec, ImageArtifactType +from openhcs.utils.import_utils import optional_import_placeholder from openhcs.core.memory import tensorflow as tensorflow_func -from openhcs.core.lazy_gpu_imports import tf from openhcs.core.pipeline.function_contracts import artifact_inputs -from metaclass_registry.strategies import EnumKeyedStrategyMixin +from openhcs.processing.backends.processors import image_stack_operations as operations from openhcs.processing.backends.processors.method_axes import ( StackProjectionMethod, ) -# Define error variable -TENSORFLOW_ERROR = "" +tf = optional_import_placeholder("tensorflow") -# Check TensorFlow version for DLPack compatibility if available -if tf is not None: - try: - tf_version = parse_version(tf.__version__) - min_version = parse_version("2.12.0") - - if tf_version < min_version: - TENSORFLOW_ERROR = ( - f"TensorFlow version {tf.__version__} is not supported for DLPack operations. " - f"Version 2.12.0 or higher is required for stable DLPack support. " - f"Clause 88 violation: Cannot infer DLPack capability." - ) - tf = None - except Exception as e: - TENSORFLOW_ERROR = str(e) - tf = None - -logger = logging.getLogger(__name__) +OPERATIONS = ArrayOperations.for_memory(MemoryType.TENSORFLOW) @tensorflow_func @@ -66,70 +40,9 @@ def create_linear_weight_mask( Returns: 2D TensorFlow weight mask of shape (height, width) """ - # The compiler will ensure this function is only called when TensorFlow is available - # No need to check for TensorFlow availability here - - margin_y = int(tf.math.floor(height * margin_ratio)) - margin_x = int(tf.math.floor(width * margin_ratio)) - - weight_y = tf.ones(height, dtype=tf.float32) - if margin_y > 0: - ramp_top = tf.linspace(0.0, 1.0, margin_y) - ramp_bottom = tf.linspace(1.0, 0.0, margin_y) - - # Update slices of the weight_y tensor - weight_y = tf.tensor_scatter_nd_update( - weight_y, tf.stack([tf.range(margin_y)], axis=1), ramp_top - ) - weight_y = tf.tensor_scatter_nd_update( - weight_y, - tf.stack([tf.range(height - margin_y, height)], axis=1), - ramp_bottom, - ) - - weight_x = tf.ones(width, dtype=tf.float32) - if margin_x > 0: - ramp_left = tf.linspace(0.0, 1.0, margin_x) - ramp_right = tf.linspace(1.0, 0.0, margin_x) - - # Update slices of the weight_x tensor - weight_x = tf.tensor_scatter_nd_update( - weight_x, tf.stack([tf.range(margin_x)], axis=1), ramp_left - ) - weight_x = tf.tensor_scatter_nd_update( - weight_x, tf.stack([tf.range(width - margin_x, width)], axis=1), ramp_right - ) - - # Create 2D weight mask using outer product - weight_mask = tf.tensordot(weight_y, weight_x, axes=0) - - return weight_mask - - -def _validate_3d_array(array: Any, name: str = "input") -> None: - """ - Validate that the input is a 3D TensorFlow tensor. - - Args: - array: Array to validate - name: Name of the array for error messages - - Raises: - TypeError: If the array is not a TensorFlow tensor - ValueError: If the array is not 3D - ImportError: If TensorFlow is not available - """ - # The compiler will ensure this function is only called when TensorFlow is available - # No need to check for TensorFlow availability here - - if not isinstance(array, tf.Tensor): - raise TypeError( - f"{name} must be a TensorFlow tensor, got {type(array)}. " - f"No automatic conversion is performed to maintain explicit contracts." - ) - - if len(array.shape) != 3: - raise ValueError(f"{name} must be a 3D tensor, got {len(array.shape)}D") + return operations.linear_weight_mask( + OPERATIONS, height, width, margin_ratio, ramp_endpoint=True + ) def _gaussian_blur(image: "tf.Tensor", sigma: float) -> "tf.Tensor": @@ -175,7 +88,7 @@ def sharpen( Returns: Sharpened 3D TensorFlow tensor of shape (Z, Y, X) """ - _validate_3d_array(image) + operations.require_stack(OPERATIONS, image) # Store original dtype dtype = image.dtype @@ -239,7 +152,7 @@ def percentile_normalize( Returns: Normalized 3D TensorFlow tensor of shape (Z, Y, X) """ - _validate_3d_array(image) + operations.require_stack(OPERATIONS, image) # Process each Z-slice independently result_list = [] @@ -305,7 +218,7 @@ def stack_percentile_normalize( Returns: Normalized 3D TensorFlow tensor of shape (Z, Y, X) """ - _validate_3d_array(stack) + operations.require_stack(OPERATIONS, stack) # Calculate global percentiles across the entire stack using TensorFlow Probability # This is memory-efficient and doesn't require sorting the entire array @@ -369,7 +282,7 @@ def create_composite( # Validate all images are 3D TensorFlow tensors with the same shape for i, img in enumerate(images): - _validate_3d_array(img, f"images[{i}]") + operations.require_stack(OPERATIONS, img, f"images[{i}]") if img.shape != images[0].shape: raise ValueError( f"All images must have the same shape. " @@ -413,24 +326,8 @@ def create_composite( composite /= total_weight # Convert back to original dtype (usually uint16) - if dtype in [tf.uint8, tf.uint16, tf.uint32, tf.int8, tf.int16, tf.int32, tf.int64]: - # Get the maximum value for the specific integer dtype - if dtype == tf.uint8: - max_val = 255 - elif dtype == tf.uint16: - max_val = 65535 - elif dtype == tf.uint32: - max_val = 4294967295 - elif dtype == tf.int8: - max_val = 127 - elif dtype == tf.int16: - max_val = 32767 - elif dtype == tf.int32: - max_val = 2147483647 - elif dtype == tf.int64: - max_val = 9223372036854775807 - - composite = tf.cast(tf.clip_by_value(composite, 0, max_val), dtype) + if OPERATIONS.dtype_name(dtype) in operations.CLAMPED_INTEGER_DTYPES: + composite = tf.cast(tf.clip_by_value(composite, 0, dtype.max), dtype) else: composite = tf.cast(composite, dtype) @@ -453,38 +350,7 @@ def apply_mask(image: "tf.Tensor", mask: "tf.Tensor") -> "tf.Tensor": Returns: Masked 3D TensorFlow tensor of shape (Z, Y, X) """ - _validate_3d_array(image) - - # Handle 2D mask (apply to each Z-slice) - if isinstance(mask, tf.Tensor) and len(mask.shape) == 2: - if mask.shape != image.shape[1:]: - raise ValueError( - f"2D mask shape {mask.shape} doesn't match image slice shape {image.shape[1:]}" - ) - - # Apply 2D mask to each Z-slice - result_list = [] - for z in range(image.shape[0]): - result_list.append( - tf.cast(image[z], tf.float32) * tf.cast(mask, tf.float32) - ) - - result = tf.stack(result_list, axis=0) - return tf.cast(result, image.dtype) - - # Handle 3D mask - if isinstance(mask, tf.Tensor) and len(mask.shape) == 3: - if mask.shape != image.shape: - raise ValueError( - f"3D mask shape {mask.shape} doesn't match image shape {image.shape}" - ) - - # Apply 3D mask directly - masked = tf.cast(image, tf.float32) * tf.cast(mask, tf.float32) - return tf.cast(masked, image.dtype) - - # If we get here, the mask is neither 2D nor 3D TensorFlow tensor - raise TypeError(f"mask must be a 2D or 3D TensorFlow tensor, got {type(mask)}") + return operations.apply_mask(OPERATIONS, image, mask) @tensorflow_func @@ -501,11 +367,7 @@ def create_weight_mask( Returns: 2D TensorFlow weight mask of shape (Y, X) """ - if not isinstance(shape, tuple) or len(shape) != 2: - raise TypeError("shape must be a tuple of (height, width)") - - height, width = shape - return create_linear_weight_mask(height, width, margin_ratio) + return create_linear_weight_mask(*operations.weight_mask_extent(shape), margin_ratio) @tensorflow_func @@ -519,11 +381,7 @@ def max_projection(stack: "tf.Tensor") -> "tf.Tensor": Returns: 3D TensorFlow tensor of shape (1, Y, X) """ - _validate_3d_array(stack) - - # Create max projection - projection_2d = tf.reduce_max(stack, axis=0) - return tf.expand_dims(projection_2d, axis=0) + return operations.max_projection(OPERATIONS, stack) @tensorflow_func @@ -537,13 +395,7 @@ def mean_projection(stack: "tf.Tensor") -> "tf.Tensor": Returns: 3D TensorFlow tensor of shape (1, Y, X) """ - _validate_3d_array(stack) - - # Create mean projection - projection_2d = tf.cast( - tf.reduce_mean(tf.cast(stack, tf.float32), axis=0), stack.dtype - ) - return tf.expand_dims(projection_2d, axis=0) + return operations.mean_projection(OPERATIONS, stack) @tensorflow_func @@ -568,7 +420,7 @@ def stack_equalize_histogram( Returns: Equalized 3D TensorFlow tensor of shape (Z, Y, X) """ - _validate_3d_array(stack) + operations.require_stack(OPERATIONS, stack) # Remember input dtype to preserve it input_dtype = stack.dtype @@ -590,31 +442,8 @@ def stack_equalize_histogram( # Normalize the CDF to the input dtype range # Avoid division by zero if tf.reduce_max(cdf) > 0: - if input_dtype in [ - tf.uint8, - tf.uint16, - tf.uint32, - tf.int8, - tf.int16, - tf.int32, - tf.int64, - ]: - # Get dtype max value - if input_dtype == tf.uint8: - dtype_max = 255.0 - elif input_dtype == tf.uint16: - dtype_max = 65535.0 - elif input_dtype == tf.uint32: - dtype_max = 4294967295.0 - elif input_dtype == tf.int8: - dtype_max = 127.0 - elif input_dtype == tf.int16: - dtype_max = 32767.0 - elif input_dtype == tf.int32: - dtype_max = 2147483647.0 - else: # int64 - dtype_max = 9223372036854775807.0 - cdf = dtype_max * cdf / tf.cast(cdf[-1], tf.float32) + if OPERATIONS.dtype_name(input_dtype) in operations.CLAMPED_INTEGER_DTYPES: + cdf = float(input_dtype.max) * cdf / tf.cast(cdf[-1], tf.float32) else: # For float dtypes, normalize to [0, 1] cdf = cdf / tf.cast(cdf[-1], tf.float32) @@ -641,33 +470,6 @@ def stack_equalize_histogram( return tf.cast(equalized_stack, input_dtype) -class TensorFlowStackProjectionStrategy( - EnumKeyedStrategyMixin[StackProjectionMethod], - metaclass=AutoRegisterMeta, -): - __enum_member_attr__ = "method" - - """TensorFlow stack-projection dispatch owner.""" - - @abstractmethod - def apply(self, stack: "tf.Tensor") -> "tf.Tensor": - """Apply this projection method.""" - - -class TensorFlowMaxStackProjectionStrategy(TensorFlowStackProjectionStrategy): - method = StackProjectionMethod.MAX - - def apply(self, stack: "tf.Tensor") -> "tf.Tensor": - return max_projection(stack) - - -class TensorFlowMeanStackProjectionStrategy(TensorFlowStackProjectionStrategy): - method = StackProjectionMethod.MEAN - - def apply(self, stack: "tf.Tensor") -> "tf.Tensor": - return mean_projection(stack) - - @tensorflow_func def create_projection( stack: "tf.Tensor", @@ -678,14 +480,12 @@ def create_projection( Args: stack: 3D TensorFlow tensor of shape (Z, Y, X) - method: Projection method (max_projection, mean_projection) + method: Projection method served by ``image_stack_operations.StackProjection`` Returns: 3D TensorFlow tensor of shape (1, Y, X) """ - _validate_3d_array(stack) - - return TensorFlowStackProjectionStrategy.for_enum_member(method).apply(stack) + return operations.create_projection(OPERATIONS, stack, method) @tensorflow_func @@ -706,7 +506,7 @@ def tophat( Returns: Filtered 3D TensorFlow tensor of shape (Z, Y, X) """ - _validate_3d_array(image) + operations.require_stack(OPERATIONS, image) # Process each Z-slice independently result_list = [] diff --git a/openhcs/processing/backends/processors/torch_processor.py b/openhcs/processing/backends/processors/torch_processor.py index d3dd6286f2..82b9b63273 100644 --- a/openhcs/processing/backends/processors/torch_processor.py +++ b/openhcs/processing/backends/processors/torch_processor.py @@ -1,47 +1,38 @@ -""" -PyTorch Image Processor Implementation - -This module implements the ImageProcessorInterface using PyTorch as the backend. -It leverages GPU acceleration for image processing operations. +"""PyTorch processor registrations. -Doctrinal Clauses: -- Clause 3 — Declarative Primacy: All functions are pure and stateless -- Clause 88 — No Inferred Capabilities: Explicit PyTorch dependency -- Clause 106-A — Declared Memory Types: All methods specify PyTorch tensors +Shared operations are written once in ``image_stack_operations``; this module +registers them for PyTorch and holds the PyTorch-specific algorithms. """ from __future__ import annotations import logging -import os -from abc import abstractmethod -from typing import Any, List, Optional, Tuple +from typing import List, Optional, Tuple -from metaclass_registry import AutoRegisterMeta +from arraybridge import MemoryType +from arraybridge.array_operations import ArrayOperations from openhcs.core.artifacts import ArtifactSpec, ImageArtifactType +from openhcs.utils.import_utils import optional_import_placeholder from openhcs.core.memory import torch as torch_func from openhcs.core.pipeline.function_contracts import artifact_inputs -from metaclass_registry.strategies import EnumKeyedStrategyMixin +from openhcs.processing.backends.processors import image_stack_operations as operations from openhcs.processing.backends.processors.method_axes import ( StackProjectionMethod, ) -from openhcs.utils.import_utils import optional_import_placeholder logger = logging.getLogger(__name__) -# Check if we're in subprocess runner mode and should skip GPU imports -if os.getenv("OPENHCS_SUBPROCESS_NO_GPU") == "1": - # Subprocess runner mode - skip GPU imports - torch = None - F = None - HAS_TORCH = False - logger.info("Subprocess runner mode - skipping torch import") -else: - # Normal mode - import PyTorch as an optional dependency - torch = optional_import_placeholder("torch") - F = optional_import_placeholder("torch.nn.functional") if torch else None - HAS_TORCH = bool(torch) +torch = optional_import_placeholder("torch") + +OPERATIONS = ArrayOperations.for_memory(MemoryType.TORCH) + + +def _integer_limits(dtype: "torch.dtype") -> Optional["torch.iinfo"]: + """``torch.iinfo`` for the clamped integer dtypes, else ``None``.""" + if OPERATIONS.dtype_name(dtype) not in operations.CLAMPED_INTEGER_DTYPES: + return None + return torch.iinfo(dtype) def create_linear_weight_mask( @@ -58,56 +49,9 @@ def create_linear_weight_mask( Returns: 2D PyTorch weight mask of shape (height, width) """ - if not torch: - raise ImportError("PyTorch is required for TorchImageProcessor") - - margin_y = int(torch.floor(torch.tensor(height * margin_ratio))) - margin_x = int(torch.floor(torch.tensor(width * margin_ratio))) - - weight_y = torch.ones(height, dtype=torch.float32) - if margin_y > 0: - ramp_top = torch.linspace(0, 1, margin_y, dtype=torch.float32) - ramp_bottom = torch.linspace(1, 0, margin_y, dtype=torch.float32) - weight_y[:margin_y] = ramp_top - weight_y[-margin_y:] = ramp_bottom - - weight_x = torch.ones(width, dtype=torch.float32) - if margin_x > 0: - ramp_left = torch.linspace(0, 1, margin_x, dtype=torch.float32) - ramp_right = torch.linspace(1, 0, margin_x, dtype=torch.float32) - weight_x[:margin_x] = ramp_left - weight_x[-margin_x:] = ramp_right - - # Create 2D weight mask using outer product - weight_mask = torch.outer(weight_y, weight_x) - - return weight_mask - - -def _validate_3d_array(array: Any, name: str = "input") -> None: - """ - Validate that the input is a 3D PyTorch tensor. - - Args: - array: Array to validate - name: Name of the array for error messages - - Raises: - TypeError: If the array is not a PyTorch tensor - ValueError: If the array is not 3D - ImportError: If PyTorch is not available - """ - if not torch: - raise ImportError("PyTorch is required for TorchImageProcessor") - - if not isinstance(array, torch.Tensor): - raise TypeError( - f"{name} must be a PyTorch tensor, got {type(array)}. " - f"No automatic conversion is performed to maintain explicit contracts." - ) - - if array.ndim != 3: - raise ValueError(f"{name} must be a 3D tensor, got {array.ndim}D") + return operations.linear_weight_mask( + OPERATIONS, height, width, margin_ratio, ramp_endpoint=True + ) def _gaussian_blur(image: "torch.Tensor", sigma: float) -> "torch.Tensor": @@ -140,6 +84,7 @@ def _gaussian_blur(image: "torch.Tensor", sigma: float) -> "torch.Tensor": # Add batch and channel dimensions to image img = image.unsqueeze(0).unsqueeze(0) + F = torch.nn.functional # Apply separable convolution blurred = F.conv2d(img, kernel_x, padding=(kernel_size // 2, 0)) @@ -166,7 +111,7 @@ def sharpen( Returns: Sharpened 3D PyTorch tensor of shape (Z, Y, X) """ - _validate_3d_array(image) + operations.require_stack(OPERATIONS, image) # Store original dtype dtype = image.dtype @@ -228,7 +173,7 @@ def percentile_normalize( Returns: Normalized 3D PyTorch tensor of shape (Z, Y, X) """ - _validate_3d_array(image) + operations.require_stack(OPERATIONS, image) # Process each Z-slice independently result = torch.zeros_like(image, dtype=torch.float32) @@ -310,7 +255,7 @@ def stack_percentile_normalize( Returns: Normalized 3D PyTorch tensor of shape (Z, Y, X) """ - _validate_3d_array(stack) + operations.require_stack(OPERATIONS, stack) # Calculate global percentiles across the entire stack # Handle large tensors that exceed PyTorch's quantile() size limits @@ -378,7 +323,7 @@ def create_composite( Composite 3D PyTorch tensor of shape (1, Y, X) """ # Validate input is 3D tensor - _validate_3d_array(stack) + operations.require_stack(OPERATIONS, stack) n_images, height, width = stack.shape @@ -418,32 +363,9 @@ def create_composite( composite /= total_weight # Convert back to original dtype (usually uint16) - if dtype in [ - torch.uint8, - torch.uint16, - torch.uint32, - torch.int8, - torch.int16, - torch.int32, - torch.int64, - ]: - # Get the maximum value for the specific integer dtype - if dtype == torch.uint8: - max_val = 255 - elif dtype == torch.uint16: - max_val = 65535 - elif dtype == torch.uint32: - max_val = 4294967295 - elif dtype == torch.int8: - max_val = 127 - elif dtype == torch.int16: - max_val = 32767 - elif dtype == torch.int32: - max_val = 2147483647 - elif dtype == torch.int64: - max_val = 9223372036854775807 - - composite = torch.clamp(composite, 0, max_val).to(dtype) + limits = _integer_limits(dtype) + if limits is not None: + composite = torch.clamp(composite, 0, limits.max).to(dtype) else: composite = composite.to(dtype) @@ -467,35 +389,7 @@ def apply_mask(image: "torch.Tensor", mask: "torch.Tensor") -> "torch.Tensor": Returns: Masked 3D PyTorch tensor of shape (Z, Y, X) - dimensionality preserved """ - _validate_3d_array(image) - - # Handle 2D mask (apply to each Z-slice) - if isinstance(mask, torch.Tensor) and mask.ndim == 2: - if mask.shape != image.shape[1:]: - raise ValueError( - f"2D mask shape {mask.shape} doesn't match image slice shape {image.shape[1:]}" - ) - - # Apply 2D mask to each Z-slice - result = torch.zeros_like(image) - for z in range(image.shape[0]): - result[z] = image[z].float() * mask.float() - - return result.to(image.dtype) - - # Handle 3D mask - if isinstance(mask, torch.Tensor) and mask.ndim == 3: - if mask.shape != image.shape: - raise ValueError( - f"3D mask shape {mask.shape} doesn't match image shape {image.shape}" - ) - - # Apply 3D mask directly - masked = image.float() * mask.float() - return masked.to(image.dtype) - - # If we get here, the mask is neither 2D nor 3D PyTorch tensor - raise TypeError(f"mask must be a 2D or 3D PyTorch tensor, got {type(mask)}") + return operations.apply_mask(OPERATIONS, image, mask) @torch_func @@ -512,11 +406,7 @@ def create_weight_mask( Returns: 2D PyTorch weight mask of shape (Y, X) """ - if not isinstance(shape, tuple) or len(shape) != 2: - raise TypeError("shape must be a tuple of (height, width)") - - height, width = shape - return create_linear_weight_mask(height, width, margin_ratio) + return create_linear_weight_mask(*operations.weight_mask_extent(shape), margin_ratio) @torch_func @@ -530,24 +420,7 @@ def max_projection(stack: "torch.Tensor") -> "torch.Tensor": Returns: 3D PyTorch tensor of shape (1, Y, X) """ - _validate_3d_array(stack) - - # Store original dtype for conversion back - original_dtype = stack.dtype - - # Convert to float32 if needed for GPU operations - if stack.dtype == torch.uint16: - stack_float = stack.float() - else: - stack_float = stack - - # Create max projection - projection_2d = torch.max(stack_float, dim=0)[0] - - # Convert back to original dtype - projection_2d = projection_2d.to(original_dtype) - - return projection_2d.reshape(1, projection_2d.shape[0], projection_2d.shape[1]) + return operations.max_projection(OPERATIONS, stack) @torch_func @@ -561,21 +434,7 @@ def mean_projection(stack: "torch.Tensor") -> "torch.Tensor": Returns: 3D PyTorch tensor of shape (1, Y, X) """ - _validate_3d_array(stack) - - # Store original dtype for conversion back - original_dtype = stack.dtype - - # Convert to float32 for mean calculation (always needed for mean) - stack_float = stack.float() - - # Create mean projection - projection_2d = torch.mean(stack_float, dim=0) - - # Convert back to original dtype - projection_2d = projection_2d.to(original_dtype) - - return projection_2d.reshape(1, projection_2d.shape[0], projection_2d.shape[1]) + return operations.mean_projection(OPERATIONS, stack) @torch_func @@ -600,7 +459,7 @@ def stack_equalize_histogram( Returns: Equalized 3D PyTorch tensor of shape (Z, Y, X) """ - _validate_3d_array(stack) + operations.require_stack(OPERATIONS, stack) # Remember input dtype to preserve it input_dtype = stack.dtype @@ -635,18 +494,10 @@ def stack_equalize_histogram( # Normalize the CDF to the input dtype range # Avoid division by zero + limits = _integer_limits(input_dtype) if cdf[-1] > 0: - if input_dtype in [ - torch.uint8, - torch.uint16, - torch.uint32, - torch.int8, - torch.int16, - torch.int32, - torch.int64, - ]: - dtype_info = torch.iinfo(input_dtype) - cdf = dtype_info.max * cdf / cdf[-1] + if limits is not None: + cdf = limits.max * cdf / cdf[-1] else: # For float dtypes, normalize to [0, 1] cdf = cdf / cdf[-1] @@ -668,50 +519,12 @@ def stack_equalize_histogram( equalized_stack = equalized_flat.reshape(stack.shape) # Convert back to input dtype - if input_dtype in [ - torch.uint8, - torch.uint16, - torch.uint32, - torch.int8, - torch.int16, - torch.int32, - torch.int64, - ]: - dtype_info = torch.iinfo(input_dtype) - return torch.clamp(equalized_stack, dtype_info.min, dtype_info.max).to( - input_dtype - ) + if limits is not None: + return torch.clamp(equalized_stack, limits.min, limits.max).to(input_dtype) else: return equalized_stack.to(input_dtype) -class TorchStackProjectionStrategy( - EnumKeyedStrategyMixin[StackProjectionMethod], - metaclass=AutoRegisterMeta, -): - __enum_member_attr__ = "method" - - """PyTorch stack-projection dispatch owner.""" - - @abstractmethod - def apply(self, stack: "torch.Tensor") -> "torch.Tensor": - """Apply this projection method.""" - - -class TorchMaxStackProjectionStrategy(TorchStackProjectionStrategy): - method = StackProjectionMethod.MAX - - def apply(self, stack: "torch.Tensor") -> "torch.Tensor": - return max_projection(stack) - - -class TorchMeanStackProjectionStrategy(TorchStackProjectionStrategy): - method = StackProjectionMethod.MEAN - - def apply(self, stack: "torch.Tensor") -> "torch.Tensor": - return mean_projection(stack) - - @torch_func def create_projection( stack: "torch.Tensor", @@ -722,14 +535,12 @@ def create_projection( Args: stack: 3D PyTorch tensor of shape (Z, Y, X) - method: Projection method (max_projection, mean_projection) + method: Projection method served by ``image_stack_operations.StackProjection`` Returns: 3D PyTorch tensor of shape (1, Y, X) """ - _validate_3d_array(stack) - - return TorchStackProjectionStrategy.for_enum_member(method).apply(stack) + return operations.create_projection(OPERATIONS, stack, method) @torch_func @@ -750,7 +561,9 @@ def tophat( Returns: Filtered 3D PyTorch tensor of shape (Z, Y, X) """ - _validate_3d_array(image) + operations.require_stack(OPERATIONS, image) + + F = torch.nn.functional # Store device for later use device = image.device diff --git a/pyproject.toml b/pyproject.toml index d5fb1c3106..3139640895 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -104,7 +104,7 @@ dependencies = [ "objectstate>=1.3.1,<2", "python-introspect>=0.2.5,<0.3", "metaclass-registry>=0.3.0,<0.4", - "arraybridge>=0.3.8,<0.4", + "arraybridge>=0.4.1,<0.5", "polystore>=0.6,<0.7", "pyqt-reactive>=0.4,<0.5", ] diff --git a/tests/unit/test_processor_method_strategies.py b/tests/unit/test_processor_method_strategies.py deleted file mode 100644 index 847c9f3977..0000000000 --- a/tests/unit/test_processor_method_strategies.py +++ /dev/null @@ -1,199 +0,0 @@ -"""Processor method strategy registration tests.""" - -from __future__ import annotations - -import ast -from enum import Enum -import importlib -from pathlib import Path - -import numpy as np - -from openhcs.core.callable_contract import CallableContract -from openhcs.processing.backends.lib_registry.openhcs_registry import OpenHCSRegistry -from openhcs.core.processing_contracts import ( - VolumetricToSliceContract, -) -from openhcs.processing.backends.processors import method_axes, numpy_processor -from openhcs.processing.backends.processors.numpy_processor import ( - NumpyStackProjectionMethod, -) -from openhcs.processing.backends.processors.method_axes import ( - SpatialBinMethod, -) - - -def _enum_members_for_symbol( - tree: ast.Module, - symbol_name: str, -) -> set[str]: - """Resolve one local or imported enum owner without a name-specific mirror.""" - - local_class = next( - ( - node - for node in tree.body - if isinstance(node, ast.ClassDef) and node.name == symbol_name - ), - None, - ) - if local_class is not None: - return { - target.id - for statement in local_class.body - if isinstance(statement, ast.Assign) - for target in statement.targets - if isinstance(target, ast.Name) - } - - imported_symbol = next( - ( - (statement.module, imported.name) - for statement in tree.body - if isinstance(statement, ast.ImportFrom) - for imported in statement.names - if (imported.asname or imported.name) == symbol_name - ), - None, - ) - assert imported_symbol is not None, symbol_name - module_name, attribute_name = imported_symbol - assert module_name is not None - enum_type = getattr(importlib.import_module(module_name), attribute_name) - assert isinstance(enum_type, type) and issubclass(enum_type, Enum) - return {member.name for member in enum_type} - - -def test_numpy_processor_method_strategies_use_distinct_inherited_registries() -> None: - assert set(numpy_processor.NumpySpatialBinStrategy.__registry__) == { - method.value for method in SpatialBinMethod - } - assert set(numpy_processor.NumpyStackProjectionStrategy.__registry__) == { - method.value for method in NumpyStackProjectionMethod - } - assert ( - numpy_processor.NumpySpatialBinStrategy.__registry__ - is not numpy_processor.NumpyStackProjectionStrategy.__registry__ - ) - - -def test_numpy_processor_method_strategies_dispatch_behavior() -> None: - stack = np.arange(2 * 4 * 4, dtype=np.uint16).reshape(2, 4, 4) - - assert numpy_processor.spatial_bin_2d(stack, 2, SpatialBinMethod.MEAN).shape == ( - 2, - 2, - 2, - ) - assert numpy_processor.spatial_bin_3d(stack, 2, SpatialBinMethod.MAX).shape == ( - 1, - 2, - 2, - ) - assert numpy_processor.create_projection( - stack, NumpyStackProjectionMethod.MEAN - ).shape == (4, 4) - np.testing.assert_array_equal( - numpy_processor.create_projection(stack, NumpyStackProjectionMethod.MIN), - np.min(stack, axis=0), - ) - assert ( - CallableContract.from_callable( - numpy_processor.create_projection - ).require_processing_contract() - is VolumetricToSliceContract - ) - - -def test_projection_method_types_exactly_match_registered_dispatch() -> None: - processor_root = ( - Path(__file__).parents[2] / "openhcs/processing/backends/processors" - ) - discovered: dict[str, set[str]] = {} - - for path in processor_root.glob("*_processor.py"): - tree = ast.parse(path.read_text(), filename=str(path)) - functions = { - node.name: node for node in tree.body if isinstance(node, ast.FunctionDef) - } - create_projection = functions.get("create_projection") - if create_projection is None: - continue - method_parameter = next( - argument - for argument in ( - *create_projection.args.args, - *create_projection.args.kwonlyargs, - ) - if argument.arg == "method" - ) - method_type_name = ast.unparse(method_parameter.annotation) - accepted_members = _enum_members_for_symbol(tree, method_type_name) - - strategy_roots = { - node.name - for node in tree.body - if isinstance(node, ast.ClassDef) - and any( - isinstance(base, ast.Subscript) - and ast.unparse(base.value) == "EnumKeyedStrategyMixin" - and ast.unparse(base.slice) == method_type_name - for base in node.bases - ) - and any( - isinstance(statement, ast.Assign) - and any( - isinstance(target, ast.Name) and target.id == "__enum_member_attr__" - for target in statement.targets - ) - and isinstance(statement.value, ast.Constant) - and statement.value.value == "method" - for statement in node.body - ) - } - registered_members = { - statement.value.attr - for node in tree.body - if isinstance(node, ast.ClassDef) - and any( - isinstance(base, ast.Name) and base.id in strategy_roots - for base in node.bases - ) - for statement in node.body - if isinstance(statement, ast.Assign) - and any( - isinstance(target, ast.Name) and target.id == "method" - for target in statement.targets - ) - and isinstance(statement.value, ast.Attribute) - } - discovered[path.name] = registered_members - assert registered_members == accepted_members, path - - assert discovered - - -def test_processor_method_roots_use_shared_enum_strategy_directly() -> None: - method_axes_source = Path(method_axes.__file__).read_text(encoding="utf-8") - - assert "RegisteredProcessorMethodStrategy" not in method_axes_source - - -def test_openhcs_registry_discovery_does_not_replace_public_processor_callables() -> ( - None -): - declared_projection = numpy_processor.create_projection - registry = OpenHCSRegistry() - registry.MODULES_TO_SCAN = [numpy_processor.__name__] - - functions = registry.discover_functions() - - assert numpy_processor.create_projection is declared_projection - assert ( - functions["processors_numpy_processor_create_projection"].func - is not declared_projection - ) - stack = np.arange(2 * 4 * 4, dtype=np.uint16).reshape(2, 4, 4) - assert numpy_processor.create_projection( - stack, NumpyStackProjectionMethod.MEAN - ).shape == (4, 4) diff --git a/tests/unit/test_processor_operations_written_once.py b/tests/unit/test_processor_operations_written_once.py new file mode 100644 index 0000000000..3a1f5bd2c4 --- /dev/null +++ b/tests/unit/test_processor_operations_written_once.py @@ -0,0 +1,184 @@ +"""Processor operations are written once against arraybridge ``ArrayOperations``. + +The six backend processor modules register shared operations from +``image_stack_operations``; a backend keeps its own body only where its +algorithm genuinely differs (L5 receipt census). +""" + +from __future__ import annotations + +import ast +import importlib +import inspect +from pathlib import Path + +import numpy as np +import pytest +from arraybridge import MemoryType +from arraybridge.array_operations import ArrayOperations + +from openhcs.core.callable_contract import CallableContract +from openhcs.core.processing_contracts import VolumetricToSliceContract +from openhcs.processing.backends.lib_registry.openhcs_registry import OpenHCSRegistry +from openhcs.processing.backends.processors import image_stack_operations +from openhcs.processing.backends.processors.method_axes import ( + SpatialBinMethod, + StackProjectionMethod, +) + +PROCESSOR_ROOT = Path(image_stack_operations.__file__).parent +BACKENDS = tuple(member.value for member in MemoryType) + +# Census of L5: backend declarations whose algorithm genuinely differs from the +# other backends' and therefore keeps its own body. +OWN_BODIES = { + "sharpen": set(BACKENDS), + "tophat": set(BACKENDS), + "stack_equalize_histogram": {"torch", "jax", "tensorflow", "pyclesperanto"}, + "percentile_normalize": {"torch", "jax", "tensorflow", "pyclesperanto"}, + "stack_percentile_normalize": {"torch", "jax", "tensorflow", "pyclesperanto"}, + "create_composite": {"numpy", "torch", "tensorflow", "pyclesperanto"}, + "create_linear_weight_mask": {"pyclesperanto"}, + "apply_mask": {"pyclesperanto"}, + "mean_projection": {"pyclesperanto"}, + "_gaussian_blur": {"torch", "jax", "tensorflow"}, +} + + +def _processor_module(backend: str): + return importlib.import_module( + f"openhcs.processing.backends.processors.{backend}_processor" + ) + + +def _processor_trees() -> dict[str, ast.Module]: + return { + backend: ast.parse((PROCESSOR_ROOT / f"{backend}_processor.py").read_text()) + for backend in BACKENDS + } + + +def _calls_shared_operation(function: ast.FunctionDef) -> bool: + """Whether the body delegates to a shared operation (validation alone is not one).""" + return any( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "operations" + and node.attr != "require_stack" + for node in ast.walk(function) + ) + + +def test_each_backend_module_registers_with_its_own_leaf() -> None: + for backend in BACKENDS: + leaf = _processor_module(backend).OPERATIONS + assert leaf is ArrayOperations.for_memory(MemoryType(backend)) + + +def test_operations_shared_by_two_backends_are_written_once() -> None: + definitions: dict[str, dict[str, ast.FunctionDef]] = {} + for backend, tree in _processor_trees().items(): + for node in tree.body: + if isinstance(node, ast.FunctionDef): + definitions.setdefault(node.name, {})[backend] = node + restated = { + (name, backend) + for name, by_backend in definitions.items() + if len(by_backend) > 1 + for backend, function in by_backend.items() + if backend not in OWN_BODIES.get(name, set()) + and not _calls_shared_operation(function) + } + assert not restated + + +def test_processor_modules_declare_no_dispatch_family_validator_or_gpu_switch() -> None: + for backend, tree in _processor_trees().items(): + source = (PROCESSOR_ROOT / f"{backend}_processor.py").read_text() + assert "OPENHCS_SUBPROCESS_NO_GPU" not in source, backend + for node in tree.body: + assert not ( + isinstance(node, ast.FunctionDef) and node.name == "_validate_3d_array" + ), backend + if isinstance(node, ast.ClassDef): + bases = {ast.unparse(base) for base in node.bases} + assert not any( + "StackProjectionMethod" in base or "SpatialBinMethod" in base + for base in bases + ), (backend, node.name) + assert not (PROCESSOR_ROOT / "percentile_utils.py").exists() + + +def test_projection_choices_are_served_by_the_projection_family() -> None: + served = set(image_stack_operations.StackProjection.__registry__) + assert served == {method.value for method in StackProjectionMethod} + assert set(image_stack_operations.SpatialBin.__registry__) == { + method.value for method in SpatialBinMethod + } + for backend in BACKENDS: + create_projection = _processor_module(backend).create_projection + method_type = type(inspect.signature(create_projection).parameters["method"].default) + assert {member.value for member in method_type} <= served, backend + + +def test_numpy_registrations_compute_the_shared_operations() -> None: + numpy_processor = _processor_module("numpy") + stack = np.arange(2 * 4 * 4, dtype=np.uint16).reshape(2, 4, 4) + + assert numpy_processor.spatial_bin_2d(stack, 2, SpatialBinMethod.MEAN).shape == (2, 2, 2) + np.testing.assert_array_equal( + numpy_processor.spatial_bin_3d(stack, 2, SpatialBinMethod.MAX), + stack.reshape(1, 2, 2, 2, 2, 2).max(axis=(1, 3, 5)), + ) + for method, reduce in ( + (numpy_processor.NumpyStackProjectionMethod.MAX, np.max), + (numpy_processor.NumpyStackProjectionMethod.MIN, np.min), + ): + np.testing.assert_array_equal( + numpy_processor.create_projection(stack, method), reduce(stack, axis=0) + ) + np.testing.assert_array_equal( + numpy_processor.create_projection(stack, numpy_processor.NumpyStackProjectionMethod.MEAN), + stack.mean(axis=0).astype(stack.dtype), + ) + assert ( + CallableContract.from_callable(numpy_processor.create_projection) + .require_processing_contract() + is VolumetricToSliceContract + ) + mask = np.array([[1, 0, 1, 0]] * 4, dtype=np.uint16) + np.testing.assert_array_equal(numpy_processor.apply_mask(stack, mask), stack * mask) + with pytest.raises(ValueError, match="3D array"): + numpy_processor.max_projection(stack[0]) + + +def test_jax_registrations_match_numpy() -> None: + jnp = pytest.importorskip("jax.numpy") + numpy_processor = _processor_module("numpy") + jax_processor = _processor_module("jax") + stack = np.random.default_rng(0).integers(0, 4000, (3, 8, 10)).astype(np.uint16) + for name in ("max_projection", "mean_projection"): + np.testing.assert_array_equal( + np.asarray(getattr(jax_processor, name)(jnp.asarray(stack))), + getattr(numpy_processor, name)(stack), + ) + np.testing.assert_allclose( + np.asarray(jax_processor.create_linear_weight_mask(12, 9)), + numpy_processor.create_linear_weight_mask(12, 9), + rtol=1e-6, + ) + + +def test_registry_discovery_does_not_replace_public_processor_callables() -> None: + numpy_processor = _processor_module("numpy") + declared_projection = numpy_processor.create_projection + registry = OpenHCSRegistry() + registry.MODULES_TO_SCAN = [numpy_processor.__name__] + + functions = registry.discover_functions() + + assert numpy_processor.create_projection is declared_projection + assert ( + functions["processors_numpy_processor_create_projection"].func + is not declared_projection + )