diff --git a/REGISTRY.md b/REGISTRY.md index b02c1af..8a306e1 100644 --- a/REGISTRY.md +++ b/REGISTRY.md @@ -7,13 +7,13 @@ Auto-generated by `generate_worker_docs.py` -- do not edit manually. | Category | Count | |----------|------:| -| Annotation Workers | 27 | +| Annotation Workers | 28 | | Property Workers -- Blobs | 10 | | Property Workers -- Points | 8 | | Property Workers -- Lines | 3 | | Property Workers -- Connections | 2 | | Test / Sample Workers | 7 | -| **Total** | **57** | +| **Total** | **58** | ## Annotation Workers @@ -37,6 +37,7 @@ Create new annotations by segmenting images or connecting existing annotations. | Gaussian Blur | Applies Gaussian blur to images | | Yes | [docs](workers/annotations/gaussian_blur/GAUSSIAN_BLUR.md) | | H&E Deconvolution | Deconvolves H&E stains | | Yes | [docs](workers/annotations/h_and_e_deconvolution/H_AND_E_DECONVOLUTION.md) | | Histogram Matching | Corrects images using histogram matching | | Yes | [docs](workers/annotations/histogram_matching/HISTOGRAM_MATCHING.md) | +| Illumination Correction | Stitched-TIFF fallback for grid-locked illumination correction with held-out channel-model selection | | Yes | [docs](workers/annotations/illumination_correction/ILLUMINATION_CORRECTION.md) | | Laplacian of Gaussian | This tool finds spots in an image using the Laplacian of Gaussian method.It uses a filt... | | | [docs](workers/annotations/laplacian_of_gaussian/LAPLACIAN_OF_GAUSSIAN.md) | | Time lapse registration | Corrects images using time lapse registration | | Yes | [docs](workers/annotations/registration/REGISTRATION.md) | | Rolling Ball | Corrects images using a rolling ball | | Yes | [docs](workers/annotations/rolling_ball/ROLLING_BALL.md) | diff --git a/docker-compose.yml b/docker-compose.yml index ba42c16..a0e9da0 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -197,6 +197,15 @@ services: depends_on: - image-processing-base + illumination_correction: + build: + context: . + dockerfile: ./workers/annotations/illumination_correction/Dockerfile + image: annotations/illumination_correction:latest + profiles: ["worker", "image-processing"] + depends_on: + - image-processing-base + registration: build: context: . @@ -447,6 +456,15 @@ services: - histogram_matching profiles: ["test"] + illumination_correction_test: + build: + context: . + dockerfile: ./workers/annotations/illumination_correction/tests/Dockerfile_Test + image: annotations/illumination_correction:test + depends_on: + - illumination_correction + profiles: ["test"] + rolling_ball_test: build: context: . diff --git a/workers/annotations/illumination_correction/Dockerfile b/workers/annotations/illumination_correction/Dockerfile new file mode 100644 index 0000000..bfc095c --- /dev/null +++ b/workers/annotations/illumination_correction/Dockerfile @@ -0,0 +1,33 @@ +FROM nimbusimage/image-processing-base:latest + +# BaSiCPy 2.0 uses a PyTorch backend. Install the cross-platform CPU build and +# its runtime optimizer explicitly, then install BaSiCPy without its stale +# scipy<1.13 metadata pin. The worker base intentionally carries a newer +# SciPy/NumPy stack; the worker tests exercise real fits with darkfield off/on. +SHELL ["conda", "run", "-n", "worker", "/bin/bash", "-c"] +RUN conda install --yes -c conda-forge \ + "pytorch>=2,<3" \ + "pydantic>=2,<3" \ + pooch \ + tqdm && \ + pip install "torch-dct==0.1.6" "hyperactive>=4.4,<5" && \ + pip install --no-deps "basicpy==2.0.0" && \ + python -c "import dask.array; from basicpy import BaSiC; print(BaSiC.__name__)" && \ + conda clean --all --yes + +COPY ./workers/annotations/illumination_correction/entrypoint.py / +COPY ./workers/annotations/illumination_correction/illumination.py / + +LABEL isUPennContrastWorker="" \ + isGPUWorker="false" \ + isAnnotationWorker="" \ + workerVersion="1.1.0" \ + interfaceName="Illumination Correction" \ + interfaceCategory="Image Processing" \ + description="Stitched-TIFF fallback for grid-locked illumination correction with held-out model selection" \ + hasPreview="False" \ + advancedOptionsPanel="False" \ + annotationConfigurationPanel="False" \ + defaultToolName="Illumination Correction" + +# ENTRYPOINT inherited from image-processing-base (run_worker.sh /entrypoint.py). diff --git a/workers/annotations/illumination_correction/ILLUMINATION_CORRECTION.md b/workers/annotations/illumination_correction/ILLUMINATION_CORRECTION.md new file mode 100644 index 0000000..3f8ad4e --- /dev/null +++ b/workers/annotations/illumination_correction/ILLUMINATION_CORRECTION.md @@ -0,0 +1,107 @@ +# Illumination Correction Worker + +Corrects grid-locked uneven illumination in an already-stitched microscopy TIFF and uploads a new multi-frame TIFF. It can use a requested algorithm or automatically choose independently for each selected channel, including choosing no correction. + +This worker is the stitched-TIFF fallback described by the illumination-correction study. The study's preferred v7 workflow starts from raw overlapping acquisition tiles and uses overlap/DCT or CIDRE-style information that is unavailable after stitching; use that upstream workflow when raw tiles and overlap metadata are available. + +## How It Works + +1. Loads the requested reference XY/Z/time plane. +2. Fits the physical stitched-tile grid. Automatic reference mode evaluates every available channel, identifies the dominant cross-channel pitch cluster, and uses its highest-quality member. Only the grid geometry is shared across channels. +3. Fits each selected channel independently on the reference Z plane. +4. In automatic algorithm mode, evaluates on representative held-out Z planes: + - Identity (leave the channel unchanged) + - BaSiC with darkfield disabled and enabled + - Folded log-gradient + - Split-half affine +5. Rejects candidates that damage object-intensity ranking, fine detail, or numeric range, contain non-finite values, or infer an implausible field. Metrics that cannot be measured on a plane are explicitly recorded as unavailable. +6. Finds the Pareto-optimal candidates across the artifact panel, accounts for score uncertainty across validation planes, and prefers the simpler model inside the resulting tie margin. Spot uniformity participates only for channels explicitly marked punctate. +7. Applies the selected channel model across Z only at the reference XY and reference time, preserves other acquisitions and unselected channels, writes a TIFF, and uploads it to Girder. + +Automatic selection is deliberately conservative: if there is no independent Z plane, the identity candidate is returned and the channel is left unchanged with a warning. A manual algorithm can still be requested for a single-plane dataset. + +The workflow is based on the channel-specific evaluation in `FINDINGS_AND_WORKFLOW.md` from the illumination-correction study. The acquisition geometry is estimated once, while flatfield, darkfield, and per-tile gains are never transferred between channels. + +## Interface Parameters + +| Parameter | Type | Default | Description | +|---|---|---|---| +| **Channels to correct** | channelCheckboxes | — | Channels for which independent correction models are fitted and applied. | +| **Algorithm** | select | Automatic (recommended) | Choose automatic comparison, BaSiC, folded log-gradient, or split-half affine. | +| **Reference channel mode** | select | Automatically choose best channel | Select the best grid reference from all channels or use the specified channel. | +| **Reference channel** | channel | 0 | Manual grid-reference channel; ignored in automatic reference mode. | +| **Reference XY** | text | blank | 1-based XY used for grid and model fitting; blank uses the current XY. Only this XY is corrected. | +| **Reference Z** | text | blank | 1-based Z used for fitting; blank uses the current Z. Use a well-focused plane. | +| **Reference Time** | text | blank | 1-based time point used for fitting; blank uses the current time. Only this time point is corrected. | +| **BaSiC darkfield** | select | Automatic | In BaSiC mode, compare both settings or force darkfield on/off. Automatic algorithm mode always evaluates both. | +| **Per-tile gain correction** | checkbox | false | Experimental whole-tile gain correction for BaSiC and folded log-gradient. It is estimated from the fit plane and can absorb biological field differences. It has no effect on split-half affine. | +| **Punctate channels for spot metric** | channelCheckboxes | none | Channels for which position-dependent spot counts may influence automatic selection. Leave empty for non-punctate signal. | +| **Output type** | select | Float32 (recommended) | Keep float32 values for audit, or preserve source dtype. Preserve-dtype output fails if more than `1e-4` of corrected pixels would be clipped. | +| **Validate every corrected plane** | checkbox | true | Recheck object rank, high-frequency detail, and numeric range on every output frame. | +| **Minimum tile pitch** | number | 150 px | Lower bound for physical stitched-tile pitch detection. | +| **Maximum tile pitch** | number | 1400 px | Upper bound for physical stitched-tile pitch detection. | + +## Algorithms + +### BaSiC + +Complete measured seam-to-seam intervals are resampled to a common 256×256 tile coordinate and passed to BaSiCPy 2.0. The inferred field is expanded across the mosaic using the measured, slightly jittered seam locations. When enabled, correction is `(raw - dark) / flat + mean(dark)`, which restores the mean pedestal. Automatic mode tests darkfield both off and on because the additive term is channel-specific. + +### Folded log-gradient + +The method robustly combines aligned gradients of log intensity across physical tiles, integrates the periodic gradient field with a Fourier-domain Poisson solve, and expands the result over the mosaic. It is useful for dense or saturated channels where a low background quantile is contaminated by biology. + +### Split-half affine + +This conservative comparator estimates separable multiplicative and additive position-locked curves. Broad-to-fine spatial bands are retained in proportion to their split-half reproducibility, reducing the chance that non-repeating biological structures enter the correction. + +## Automatic Evaluation + +Artifact metrics are calculated against each channel's raw plane. Models are fitted on the requested reference Z and automatic ranking uses representative held-out Z planes (first, middle, and last when available, excluding the fit plane): + +- Jitter-aware within-tile amplitude (A1) +- Tile-frequency harmonic modulation (A2) +- Held-out within-tile background dependence (A3) +- Background dynamic range (A5) +- Detrended whole-tile scatter (A6) + +The unchanged image is the baseline candidate. A correction must improve beyond the tie/uncertainty margin to displace it. Candidates are rejected when any applicable hard preservation guardrail fails: + +- Object-intensity Spearman rank below 0.98, when at least 10 measurable objects are available +- Locally normalized high-frequency power below 0.90 of raw, when finite source high-frequency power is available +- Any non-finite source or output pixels +- More than `1e-4` newly nonpositive pixels (pre-existing source zeros are not treated as correction damage) + +BaSiC darkfields are also rejected when their mean reaches or exceeds the reference plane's first-percentile image floor. Spot-count uniformity uses a minimum-count requirement and symmetric pseudocounts, and is a soft selection term only for channels marked as punctate. + +## Output and Metadata + +The worker uploads `/tmp/illumination_corrected.tiff` to the source dataset. Girder metadata records: + +- Requested and selected algorithm for every corrected channel +- Full candidate scores and rejection reasons +- Explicit zero-based and one-based reference channel and XY/Z/time coordinates +- Held-out Z planes, correction scope, pitch bounds, and punctate-metric channels +- Measured pitch, seam positions, residuals, and reference-quality reports +- Model diagnostics, output type, validation setting, worker version, and dependency versions + +Channel names, pixel size, and magnification are copied when present in the source tile metadata. Single-frame datasets without `IndexRange` are supported. + +## Implementation Notes and Limitations + +- This worker is for stitched mosaics with a repeatable physical tile pattern. Each image axis must contain at least four candidate periods and enough detected seams to bound complete tiles. +- Automatic reference selection chooses the best channel on the requested reference XY/Z/time plane; it does not search all Z planes for focus. Navigate to, or specify, a representative focused plane. +- The chosen grid and fitted channel models are shared across Z only at the reference XY and time. Other XY positions and time points are preserved because their acquisition-specific grids and gains may differ. +- Automatic mode needs at least two Z planes. On a single-Z dataset it selects identity; choose a manual algorithm only when fitting and evaluating on the same plane is scientifically acceptable. +- Per-tile gain correction is off by default because it can absorb real whole-tile biology. When enabled, it affects only BaSiC and folded log-gradient. +- Float32 is the recommended audit output. Preserve-source-dtype mode allows only negligible (`<=1e-4`) range clipping and otherwise fails with a request to use Float32. +- Conditional guardrails are recorded as unavailable rather than silently passed. Non-finite input/output and newly nonpositive-pixel checks are always required. +- Missing or stale select values are rejected before image loading. The worker reports which setting must be re-selected and raises so Girder records the job as failed rather than successful. +- The worker is CPU-routed (`isGPUWorker=false`). BaSiCPy uses a CPU PyTorch backend, so its image is larger than other classical image-processing workers. The image intentionally overrides BaSiCPy's stale `scipy<1.13` metadata pin while installing its `hyperactive` runtime dependency explicitly; Docker tests run real BaSiC fits with darkfield both off and on. + +## Build and Test + +```bash +./build_workers.sh illumination_correction +./build_workers.sh --build-and-run-tests illumination_correction +``` diff --git a/workers/annotations/illumination_correction/entrypoint.py b/workers/annotations/illumination_correction/entrypoint.py new file mode 100644 index 0000000..7945d7e --- /dev/null +++ b/workers/annotations/illumination_correction/entrypoint.py @@ -0,0 +1,654 @@ +import argparse +import importlib.metadata +import json +import sys + +import numpy as np + +import annotation_client.tiles as tiles +import annotation_client.workers as workers +from annotation_client.utils import sendError, sendProgress, sendWarning +import annotation_utilities.annotation_tools as annotation_tools + +import illumination as correction + + +OUTPUT_PATH = "/tmp/illumination_corrected.tiff" +WORKER_VERSION = "1.1.0" +ALGORITHM_OPTIONS = correction.ALGORITHM_OPTIONS +REFERENCE_CHANNEL_MODE_OPTIONS = correction.REFERENCE_CHANNEL_MODE_OPTIONS +DARKFIELD_OPTIONS = correction.DARKFIELD_OPTIONS +OUTPUT_TYPE_OPTIONS = correction.OUTPUT_TYPE_OPTIONS + + +def interface(image, apiUrl, token): + client = workers.UPennContrastWorkerPreviewClient(apiUrl=apiUrl, token=token) + values = { + "Illumination correction": { + "type": "notes", + "value": ( + "Stitched-TIFF fallback for grid-locked uneven illumination. " + "Automatic mode fits on the reference Z plane, selects on held-out " + "Z planes, and can leave a channel unchanged." + ), + "displayOrder": 0, + }, + "Channels to correct": { + "type": "channelCheckboxes", + "required": True, + "tooltip": "Channels to correct in the uploaded image.", + "displayOrder": 1, + }, + "Algorithm": { + "type": "select", + "items": list(ALGORITHM_OPTIONS), + "default": "Automatic (recommended)", + "tooltip": ( + "Automatic compares BaSiC with darkfield off/on, folded " + "log-gradient, split-half affine, and no correction. It requires " + "an independent Z plane for model selection." + ), + "displayOrder": 2, + }, + "Reference channel mode": { + "type": "select", + "items": list(REFERENCE_CHANNEL_MODE_OPTIONS), + "default": "Automatically choose best channel", + "tooltip": ( + "Automatic mode chooses the most reliable grid estimate in the " + "dominant cross-channel pitch cluster." + ), + "displayOrder": 3, + }, + "Reference channel": { + "type": "channel", + "default": 0, + "tooltip": "Used only when Reference channel mode is manual.", + "displayOrder": 4, + }, + "Reference XY": { + "type": "text", + "default": "", + "vueAttrs": { + "placeholder": "blank = current XY", + "label": "Reference XY (1-based)", + "persistentPlaceholder": True, + "filled": True, + }, + "tooltip": ( + "Physical grid and channel models are fitted at this XY. Only this " + "XY and the reference time point are corrected." + ), + "displayOrder": 5, + }, + "Reference Z": { + "type": "text", + "default": "", + "vueAttrs": { + "placeholder": "blank = current Z", + "label": "Reference Z (1-based)", + "persistentPlaceholder": True, + "filled": True, + }, + "tooltip": "Use a well-focused plane; the fitted model is shared across Z.", + "displayOrder": 6, + }, + "Reference Time": { + "type": "text", + "default": "", + "vueAttrs": { + "placeholder": "blank = current time", + "label": "Reference time (1-based)", + "persistentPlaceholder": True, + "filled": True, + }, + "tooltip": ( + "Time point used to fit the illumination models. Only this time " + "point and the reference XY are corrected." + ), + "displayOrder": 7, + }, + "BaSiC darkfield": { + "type": "select", + "items": list(DARKFIELD_OPTIONS), + "default": "Automatic", + "tooltip": ( + "Automatic evaluates both settings. Darkfield is channel-specific " + "and is rejected when it is physically implausible." + ), + "displayOrder": 8, + }, + "Per-tile gain correction": { + "type": "checkbox", + "default": False, + "tooltip": ( + "Experimental for BaSiC and folded log-gradient only. It estimates " + "whole-tile gains from the fit plane and can remove real biology." + ), + "displayOrder": 9, + }, + "Punctate channels for spot metric": { + "type": "channelCheckboxes", + "required": False, + "tooltip": ( + "Use position-dependent spot counts as a soft selection metric only " + "for channels known to contain punctate signal." + ), + "displayOrder": 10, + }, + "Output type": { + "type": "select", + "items": list(OUTPUT_TYPE_OPTIONS), + "default": "Float32 (recommended)", + "tooltip": ( + "Float32 preserves the fitted numeric range for audit. Preserve " + "source dtype rejects corrections that would materially clip." + ), + "displayOrder": 11, + }, + "Validate every corrected plane": { + "type": "checkbox", + "default": True, + "tooltip": ( + "Check object ranking, fine detail, and numeric range on every " + "corrected frame before uploading." + ), + "displayOrder": 12, + }, + "Minimum tile pitch": { + "type": "number", + "min": 16, + "max": 5000, + "default": 150, + "unit": "pixels", + "tooltip": "Smallest physical stitched-tile pitch to consider.", + "displayOrder": 13, + }, + "Maximum tile pitch": { + "type": "number", + "min": 32, + "max": 10000, + "default": 1400, + "unit": "pixels", + "tooltip": "Largest physical stitched-tile pitch to consider.", + "displayOrder": 14, + }, + } + client.setWorkerImageInterface(image, values) + + +def _parse_reference_coordinate(value, current, name): + if value is None or str(value).strip() == "": + return int(current) + try: + coordinate = int(str(value).strip()) - 1 + except ValueError as exc: + raise ValueError(f"{name} must be a 1-based integer or left blank") from exc + if coordinate < 0: + raise ValueError(f"{name} must be at least 1") + return coordinate + + +def _source_dtype(tile_client, fallback): + try: + return np.dtype(tile_client.tiles.get("dtype", fallback)) + except TypeError: + return np.dtype(fallback) + + +def _dimension_size(tile_metadata, dimension): + key = dimension if dimension.startswith("Index") else f"Index{dimension}" + index_range = tile_metadata.get("IndexRange") or {} + if key in index_range: + return max(int(index_range[key]), 1) + frames = tile_metadata.get("frames") or [] + if frames: + return max( + int(annotation_tools.get_frame_index(frame, key)) for frame in frames + ) + 1 + if key == "IndexC" and tile_metadata.get("channels"): + return max(len(tile_metadata["channels"]), 1) + return 1 + + +def _validate_reference_coordinates(coordinates, tile_metadata): + dimensions = {"XY": "IndexXY", "Z": "IndexZ", "Time": "IndexT"} + for label, dimension in dimensions.items(): + size = _dimension_size(tile_metadata, dimension) + coordinate = coordinates[label] + if coordinate >= size: + raise ValueError( + f"Reference {label} is {coordinate + 1}, but the dataset only has " + f"{size} position{'s' if size != 1 else ''}" + ) + + +def _representative_validation_z(reference_z, num_z): + candidates = {0, num_z // 2, num_z - 1} + return sorted(z for z in candidates if 0 <= z < num_z and z != reference_z) + + +def _validate_output_values(image, output_type, source_dtype, tolerance=1e-4): + values = np.asarray(image) + if not np.isfinite(values).all(): + raise ValueError("The corrected output contains non-finite values") + if output_type == OUTPUT_TYPE_OPTIONS[0]: + return + + dtype = np.dtype(source_dtype) + if np.issubdtype(dtype, np.integer): + limits = np.iinfo(dtype) + elif np.issubdtype(dtype, np.floating): + limits = np.finfo(dtype) + else: + return + outside = (values < limits.min) | (values > limits.max) + fraction = float(np.count_nonzero(outside) / values.size) + if fraction > tolerance: + raise ValueError( + f"{fraction:.3%} of corrected pixels would be clipped when preserving " + f"source dtype {dtype}; use Float32 output instead" + ) + + +def _cast_corrected(image, output_type, source_dtype): + corrected = np.asarray(image, dtype=np.float32) + if output_type == "Float32 (recommended)": + return corrected + dtype = np.dtype(source_dtype) + if np.issubdtype(dtype, np.integer): + limits = np.iinfo(dtype) + corrected = np.clip(np.rint(corrected), limits.min, limits.max) + elif np.issubdtype(dtype, np.floating): + limits = np.finfo(dtype) + corrected = np.clip(corrected, limits.min, limits.max) + return corrected.astype(dtype) + + +def _software_versions(): + versions = {"worker": WORKER_VERSION} + for package in ("numpy", "scipy", "torch", "basicpy", "large-image"): + try: + versions[package] = importlib.metadata.version(package) + except importlib.metadata.PackageNotFoundError: + continue + return versions + + +def _json_safe(value): + if isinstance(value, dict): + return {str(key): _json_safe(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_safe(item) for item in value] + if isinstance(value, np.ndarray): + return _json_safe(value.tolist()) + if isinstance(value, (np.integer, np.floating)): + value = value.item() + if isinstance(value, float) and not np.isfinite(value): + return None + return value + + +def _copy_sink_metadata(sink, tile_metadata): + if "channels" in tile_metadata: + sink.channelNames = tile_metadata["channels"] + for name in ("mm_x", "mm_y", "magnification"): + if name in tile_metadata: + setattr(sink, name, tile_metadata[name]) + + +def compute(datasetId, apiUrl, token, params): + worker_interface = params.get("workerInterface", {}) + try: + channels = annotation_tools.get_selected_channels( + worker_interface.get("Channels to correct"), "Channels to correct" + ) + punctate_channels = annotation_tools.get_selected_channels( + worker_interface.get("Punctate channels for spot metric"), + "Punctate channels for spot metric", + ) + except ValueError as exc: + sendError("Could not read the channel selection.", info=str(exc)) + raise + if not channels: + exc = ValueError("Select at least one channel and run the worker again.") + sendError( + "No channels selected for illumination correction.", + info=str(exc), + ) + raise exc + + try: + algorithm = annotation_tools.get_required_select( + worker_interface.get("Algorithm"), + "Algorithm", + allowed_values=ALGORITHM_OPTIONS, + ) + reference_channel_mode = annotation_tools.get_required_select( + worker_interface.get("Reference channel mode"), + "Reference channel mode", + allowed_values=REFERENCE_CHANNEL_MODE_OPTIONS, + ) + darkfield_mode = annotation_tools.get_required_select( + worker_interface.get("BaSiC darkfield"), + "BaSiC darkfield", + allowed_values=DARKFIELD_OPTIONS, + ) + output_type = annotation_tools.get_required_select( + worker_interface.get("Output type"), + "Output type", + allowed_values=OUTPUT_TYPE_OPTIONS, + ) + pitch_min = float(worker_interface.get("Minimum tile pitch", 150)) + pitch_max = float(worker_interface.get("Maximum tile pitch", 1400)) + if pitch_min <= 0 or pitch_max <= pitch_min: + raise ValueError( + "Maximum tile pitch must be greater than minimum tile pitch" + ) + per_tile_gain = worker_interface.get("Per-tile gain correction", False) + validate_every_plane = worker_interface.get( + "Validate every corrected plane", True + ) + if not isinstance(per_tile_gain, bool): + raise ValueError("Per-tile gain correction must be true or false") + if not isinstance(validate_every_plane, bool): + raise ValueError("Validate every corrected plane must be true or false") + tile = params.get("tile", {}) + coordinates = { + "XY": _parse_reference_coordinate( + worker_interface.get("Reference XY"), + tile.get("XY", 0), + "Reference XY", + ), + "Z": _parse_reference_coordinate( + worker_interface.get("Reference Z"), + tile.get("Z", 0), + "Reference Z", + ), + "Time": _parse_reference_coordinate( + worker_interface.get("Reference Time"), + tile.get("Time", 0), + "Reference Time", + ), + } + reference_channel_setting = 0 + if reference_channel_mode == REFERENCE_CHANNEL_MODE_OPTIONS[1]: + reference_channel_setting = int( + worker_interface.get("Reference channel") + ) + if reference_channel_setting < 0: + raise ValueError("Reference channel must be at least 0") + except (TypeError, ValueError) as exc: + sendError("Could not read the illumination settings.", info=str(exc)) + raise + + tile_client = tiles.UPennContrastDataset( + apiUrl=apiUrl, token=token, datasetId=datasetId + ) + try: + _validate_reference_coordinates(coordinates, tile_client.tiles) + num_channels = _dimension_size(tile_client.tiles, "IndexC") + channels, missing_channels = annotation_tools.split_channel_selection( + channels, num_channels + ) + punctate_channels, missing_punctate = ( + annotation_tools.split_channel_selection( + punctate_channels, num_channels + ) + ) + if not channels: + raise ValueError( + "The selected channels do not exist in this dataset: " + + ", ".join(str(channel + 1) for channel in missing_channels) + ) + if missing_channels: + sendWarning( + "Some selected correction channels do not exist in this dataset.", + info="Ignoring channels " + + ", ".join(str(channel + 1) for channel in missing_channels), + ) + if missing_punctate: + sendWarning( + "Some punctate-metric channels do not exist in this dataset.", + info="Ignoring channels " + + ", ".join(str(channel + 1) for channel in missing_punctate), + ) + if ( + reference_channel_mode == REFERENCE_CHANNEL_MODE_OPTIONS[1] + and reference_channel_setting >= num_channels + ): + raise ValueError( + f"Reference channel {reference_channel_setting + 1} does not exist; " + f"the dataset has {num_channels} channel" + f"{'s' if num_channels != 1 else ''}" + ) + except (TypeError, ValueError) as exc: + sendError("Could not use the saved settings with this dataset.", info=str(exc)) + raise + + try: + grid, reference_channel, reference_reports = correction.choose_reference_grid( + tile_client, + datasetId, + coordinates, + reference_channel_mode, + reference_channel_setting, + pitch_min, + pitch_max, + progress=sendProgress, + ) + except Exception as exc: + sendError("Could not determine the physical tile grid.", info=str(exc)) + raise + + validation_zs = _representative_validation_z( + coordinates["Z"], _dimension_size(tile_client.tiles, "IndexZ") + ) + selections = {} + model_metadata = {} + + for index, channel in enumerate(channels): + sendProgress( + 0.15 + 0.35 * index / max(len(channels), 1), + "Illumination correction", + f"Selecting a correction for channel {channel + 1}", + ) + try: + frame = tile_client.coordinatesToFrameIndex( + coordinates["XY"], coordinates["Z"], coordinates["Time"], channel + ) + raw = correction.as_plane(tile_client.getRegion(datasetId, frame=frame)) + if raw.shape != grid.shape: + raise ValueError( + f"Reference channel {channel + 1} has shape {raw.shape}, but the " + f"shared grid expects {grid.shape}" + ) + + def validation_source(channel=channel): + for z in validation_zs: + validation_frame = tile_client.coordinatesToFrameIndex( + coordinates["XY"], z, coordinates["Time"], channel + ) + validation_plane = correction.as_plane( + tile_client.getRegion(datasetId, frame=validation_frame) + ) + if validation_plane.shape != grid.shape: + raise ValueError( + f"Held-out Z {z + 1} for channel {channel + 1} has " + f"shape {validation_plane.shape}, but the shared grid " + f"expects {grid.shape}" + ) + yield f"held-out Z {z + 1}", validation_plane + + selection = correction.select_model( + raw, + grid, + algorithm, + darkfield_mode, + per_tile_gain, + progress=sendProgress, + validation_source=validation_source if validation_zs else None, + use_spot_uniformity=channel in punctate_channels, + ) + except Exception as exc: + sendError( + f"Could not fit a safe correction for channel {channel + 1}.", + info=str(exc), + ) + raise + if selection.name == "identity": + sendWarning( + f"Channel {channel + 1} was left unchanged.", + info=getattr(selection.model, "diagnostics", {}).get( + "reason", + "No correction beat the identity baseline by a reliable margin.", + ), + ) + selections[channel] = selection + model_metadata[str(channel)] = { + "channel_index_zero_based": channel, + "selected": selection.name, + "artifact_index": selection.artifact_index, + "metrics": selection.metrics, + "diagnostics": getattr(selection.model, "diagnostics", {}), + "candidates": selection.alternatives, + } + + # Delay the heavy writer import until validation and model fitting have passed. + import large_image as li + + sink = li.new() + frames = tile_client.tiles.get("frames") + frame_records = list(enumerate(frames)) if frames else [(None, None)] + + for output_index, (frame_index, frame_metadata) in enumerate(frame_records): + if frame_metadata is None: + source_frame = tile_client.coordinatesToFrameIndex( + coordinates["XY"], + coordinates["Z"], + coordinates["Time"], + params.get("channel", 0), + ) + channel = int(params.get("channel", 0)) + frame_xy = coordinates["XY"] + frame_time = coordinates["Time"] + else: + source_frame = frame_index + channel = int(annotation_tools.get_frame_index(frame_metadata, "IndexC")) + frame_xy = int( + annotation_tools.get_frame_index(frame_metadata, "IndexXY") + ) + frame_time = int( + annotation_tools.get_frame_index(frame_metadata, "IndexT") + ) + raw = correction.as_plane(tile_client.getRegion(datasetId, frame=source_frame)) + source_dtype = _source_dtype(tile_client, raw.dtype) + should_correct = ( + channel in selections + and frame_xy == coordinates["XY"] + and frame_time == coordinates["Time"] + ) + + if should_correct: + try: + corrected = selections[channel].model.apply(raw) + _validate_output_values(corrected, output_type, source_dtype) + output = _cast_corrected(corrected, output_type, source_dtype) + _validate_output_values(output, output_type, source_dtype) + except Exception as exc: + sendError( + f"Could not create a safe output for frame {output_index + 1}.", + info=str(exc), + ) + raise + if validate_every_plane: + validation = correction.preservation_metrics(raw, output) + violations = validation["guardrail_violations"] + if violations: + exc = ValueError("; ".join(violations)) + sendError( + "Correction failed preservation checks on frame " + f"{output_index + 1}.", + info=str(exc), + ) + raise exc + elif output_type == OUTPUT_TYPE_OPTIONS[0]: + output = raw.astype(np.float32) + else: + output = raw.astype(source_dtype, copy=False) + + if frame_metadata is None: + sink.addTile(output, 0, 0, z=0) + else: + sink.addTile( + output, + 0, + 0, + **annotation_tools.frame_to_large_image_params(frame_metadata), + ) + sendProgress( + 0.5 + 0.48 * (output_index + 1) / len(frame_records), + "Illumination correction", + f"Processing frame {output_index + 1}/{len(frame_records)}", + ) + + _copy_sink_metadata(sink, tile_client.tiles) + sink.write(OUTPUT_PATH) + item = tile_client.client.uploadFileToFolder(datasetId, OUTPUT_PATH) + metadata = _json_safe( + { + "tool": "Illumination correction", + "worker_version": WORKER_VERSION, + "software_versions": _software_versions(), + "input_representation": "stitched TIFF fallback", + "indexing": "channel and coordinate keys marked zero_based are 0-based", + "algorithm_requested": algorithm, + "reference_channel_mode": reference_channel_mode, + "reference_channel_zero_based": reference_channel, + "reference_channel_one_based": reference_channel + 1, + "reference_coordinates_zero_based": coordinates, + "reference_coordinates_one_based": { + key: value + 1 for key, value in coordinates.items() + }, + "held_out_z_zero_based": validation_zs, + "held_out_z_one_based": [z + 1 for z in validation_zs], + "reference_candidates": reference_reports, + "grid": grid.as_dict(), + "per_tile_gain": per_tile_gain, + "corrected_channels_zero_based": channels, + "corrected_channels_one_based": [channel + 1 for channel in channels], + "punctate_channels_zero_based": punctate_channels, + "punctate_channels_one_based": [ + channel + 1 for channel in punctate_channels + ], + "output_type": output_type, + "validated_every_plane": validate_every_plane, + "correction_scope": { + "XY_zero_based": coordinates["XY"], + "Time_zero_based": coordinates["Time"], + "Z": "all planes at the reference XY and time", + }, + "pitch_bounds_pixels": {"minimum": pitch_min, "maximum": pitch_max}, + "darkfield_mode": darkfield_mode, + "channel_models": model_metadata, + } + ) + tile_client.client.addMetadataToItem(item["itemId"], metadata) + sendProgress(1.0, "Illumination correction", "Corrected image uploaded") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Correct uneven image illumination") + parser.add_argument("--datasetId", type=str, required=False, action="store") + parser.add_argument("--apiUrl", type=str, required=True, action="store") + parser.add_argument("--token", type=str, required=True, action="store") + parser.add_argument("--request", type=str, required=True, action="store") + parser.add_argument("--parameters", type=str, required=True, action="store") + args = parser.parse_args(sys.argv[1:]) + + parameters = json.loads(args.parameters) + match args.request: + case "compute": + compute(args.datasetId, args.apiUrl, args.token, parameters) + case "interface": + interface(parameters["image"], args.apiUrl, args.token) diff --git a/workers/annotations/illumination_correction/illumination.py b/workers/annotations/illumination_correction/illumination.py new file mode 100644 index 0000000..49683dc --- /dev/null +++ b/workers/annotations/illumination_correction/illumination.py @@ -0,0 +1,1802 @@ +"""Grid-aware uneven-illumination models and automatic model selection. + +The implementation follows the workflow validated in the adjacent illumination +study: estimate one physical acquisition grid, fit each channel independently, +reject candidates that damage biology, and compare the remaining candidates on a +multi-metric artifact panel. +""" + +from __future__ import annotations + +import math +import time +from dataclasses import dataclass, field +from typing import Callable, Iterable + +import numpy as np +from scipy import ndimage + + +MIN_PERIODS = 4 +N_FOLD_BINS = 256 +SEAM_SEARCH_DIVISOR = 6.0 +TILE_N = 256 +ALGORITHM_OPTIONS = ( + "Automatic (recommended)", + "BaSiC", + "Folded log-gradient", + "Split-half affine", +) +REFERENCE_CHANNEL_MODE_OPTIONS = ( + "Automatically choose best channel", + "Use specified channel", +) +DARKFIELD_OPTIONS = ("Automatic", "Enabled", "Disabled") +OUTPUT_TYPE_OPTIONS = ("Float32 (recommended)", "Preserve source dtype") + + +@dataclass(frozen=True) +class TileGrid: + """Measured physical-tile lattice for one stitched mosaic.""" + + pitch_y: float + pitch_x: float + seam_y: float + seam_x: float + height: int + width: int + seams_y: tuple[float, ...] | None = None + seams_x: tuple[float, ...] | None = None + seam_residual_y: float = float("nan") + seam_residual_x: float = float("nan") + prominence_y: float = float("nan") + prominence_x: float = float("nan") + + @property + def shape(self) -> tuple[int, int]: + return self.height, self.width + + @property + def is_valid(self) -> bool: + values = (self.pitch_y, self.pitch_x, self.seam_y, self.seam_x) + if not all(np.isfinite(v) and v > 0 for v in values[:2]): + return False + if not all(np.isfinite(v) for v in values[2:]): + return False + if not self.seams_y or not self.seams_x: + return False + if len(self.seams_y) < 3 or len(self.seams_x) < 3: + return False + pitch_ratio = max(self.pitch_y, self.pitch_x) / min(self.pitch_y, self.pitch_x) + if pitch_ratio > 1.25: + return False + residuals = ( + self.seam_residual_y / max(self.pitch_y, 1e-9), + self.seam_residual_x / max(self.pitch_x, 1e-9), + ) + return all(np.isfinite(v) and v < 0.25 for v in residuals) + + @property + def quality_score(self) -> float: + """Lower is better; balances square pitch, seam residual, and peak isolation.""" + if not self.is_valid: + return float("inf") + pitch_disagreement = abs(math.log(self.pitch_y / self.pitch_x)) + residual = ( + self.seam_residual_y / self.pitch_y + + self.seam_residual_x / self.pitch_x + ) + prominence = sum( + math.log1p(max(float(v), 0.0)) + for v in (self.prominence_y, self.prominence_x) + if np.isfinite(v) + ) + return float(2.0 * pitch_disagreement + residual - 0.05 * prominence) + + def _seam_array(self, axis: str) -> np.ndarray: + seams = self.seams_y if axis == "y" else self.seams_x + pitch = self.pitch_y if axis == "y" else self.pitch_x + phase = self.seam_y if axis == "y" else self.seam_x + extent = self.height if axis == "y" else self.width + + if seams: + values = np.asarray(sorted(set(seams)), dtype=np.float64) + gap = float(np.median(np.diff(values))) if values.size > 1 else pitch + else: + values = np.arange(phase % pitch, extent + pitch, pitch, dtype=np.float64) + gap = pitch + while values[0] > 0: + values = np.concatenate(([values[0] - gap], values)) + while values[-1] < extent: + values = np.concatenate((values, [values[-1] + gap])) + return values + + def complete_seams(self, axis: str) -> np.ndarray: + """Measured seams bounding complete physical tiles inside the image.""" + seams = self.seams_y if axis == "y" else self.seams_x + extent = self.height if axis == "y" else self.width + if not seams: + raise ValueError(f"No measured {axis}-axis seams are available") + values = np.asarray(sorted(set(seams)), dtype=np.float64) + values = values[(values >= 0) & (values <= extent)] + if values.size < 2: + raise ValueError(f"Fewer than two measured {axis}-axis seams are available") + return values + + def _fraction(self, positions: np.ndarray, axis: str) -> np.ndarray: + seams = self._seam_array(axis) + positions = np.asarray(positions, dtype=np.float64) + index = np.clip( + np.searchsorted(seams, positions, side="right") - 1, + 0, + seams.size - 2, + ) + widths = seams[index + 1] - seams[index] + return ((positions - seams[index]) / np.maximum(widths, 1e-9)).astype( + np.float32 + ) + + def u_of(self, y: np.ndarray) -> np.ndarray: + return self._fraction(y, "y") + + def v_of(self, x: np.ndarray) -> np.ndarray: + return self._fraction(x, "x") + + def tile_of_y(self, y: np.ndarray) -> np.ndarray: + seams = self._seam_array("y") + return np.clip( + np.searchsorted(seams, np.asarray(y), side="right") - 1, + 0, + seams.size - 2, + ) + + def tile_of_x(self, x: np.ndarray) -> np.ndarray: + seams = self._seam_array("x") + return np.clip( + np.searchsorted(seams, np.asarray(x), side="right") - 1, + 0, + seams.size - 2, + ) + + def as_dict(self) -> dict: + return { + "pitch_y": float(self.pitch_y), + "pitch_x": float(self.pitch_x), + "seam_y": float(self.seam_y), + "seam_x": float(self.seam_x), + "seams_y": [float(v) for v in self.seams_y or ()], + "seams_x": [float(v) for v in self.seams_x or ()], + "seam_residual_y": float(self.seam_residual_y), + "seam_residual_x": float(self.seam_residual_x), + "prominence_y": float(self.prominence_y), + "prominence_x": float(self.prominence_x), + "quality_score": float(self.quality_score), + } + + +def median_profiles(image: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + return ( + np.median(image, axis=1).astype(np.float64), + np.median(image, axis=0).astype(np.float64), + ) + + +def fold(profile: np.ndarray, pitch: float, bins: int = N_FOLD_BINS) -> np.ndarray: + labels = ((np.arange(profile.size) % pitch) / pitch * bins).astype(int) % bins + total = np.bincount(labels, weights=profile, minlength=bins) + count = np.bincount(labels, minlength=bins) + with np.errstate(invalid="ignore", divide="ignore"): + return np.where(count > 0, total / np.maximum(count, 1), np.nan) + + +def fold_amplitude(profile: np.ndarray, pitch: float, bins: int = N_FOLD_BINS) -> float: + folded = fold(profile, pitch, bins) + if np.all(np.isnan(folded)): + return float("nan") + return float(np.nanpercentile(folded, 97) - np.nanpercentile(folded, 3)) + + +def _detrend(profile: np.ndarray, sigma: float) -> np.ndarray: + trend = ndimage.gaussian_filter1d(profile, max(float(sigma), 1.0), mode="nearest") + return profile / np.maximum(trend, 1e-9) - 1.0 + + +def _periodogram_pitch(profile: np.ndarray, low: float, high: float) -> float: + relative = _detrend(profile, profile.size / 8.0) + n = relative.size + power = np.abs(np.fft.rfft(relative * np.hanning(n))) ** 2 + frequencies = np.fft.rfftfreq(n) + with np.errstate(divide="ignore"): + pitches = 1.0 / frequencies + keep = np.isfinite(pitches) & (pitches >= low) & (pitches <= high) + if not keep.any(): + return float("nan") + indices = np.nonzero(keep)[0] + return float(pitches[indices[int(np.argmax(power[indices]))]]) + + +def _refine_pitch(profile: np.ndarray, initial: float) -> tuple[float, float]: + relative = _detrend(profile, initial) + step = max(0.05, initial / 3000.0) + candidates = np.arange(initial * 0.96, initial * 1.04 + step, step) + amplitudes = np.asarray([fold_amplitude(relative, p) for p in candidates]) + if np.all(np.isnan(amplitudes)): + return float(initial), float("nan") + best = int(np.nanargmax(amplitudes)) + finite = amplitudes[np.isfinite(amplitudes)] + prominence = float( + amplitudes[best] / max(float(np.median(np.abs(finite))), 1e-9) + ) + return float(candidates[best]), prominence + + +def _find_seams( + profile: np.ndarray, pitch: float, phase: float +) -> tuple[np.ndarray, float]: + relative = _detrend(profile, pitch) + smooth = ndimage.gaussian_filter1d( + relative, max(pitch / 40.0, 1.0), mode="nearest" + ) + window = max(int(pitch / SEAM_SEARCH_DIVISOR), 3) + found = [] + for predicted in np.arange(phase % pitch, profile.size, pitch): + center = int(round(predicted)) + start = max(center - window, 0) + stop = min(center + window + 1, profile.size) + if stop - start < 3: + continue + index = start + int(np.argmin(smooth[start:stop])) + if index <= window // 2 or index >= profile.size - 1 - window // 2: + continue + if not (0 < index < profile.size - 1): + continue + left, middle, right = smooth[index - 1 : index + 2] + denominator = left - 2 * middle + right + offset = ( + 0.5 * (left - right) / denominator if abs(denominator) > 1e-12 else 0.0 + ) + found.append(index + float(np.clip(offset, -1.0, 1.0))) + + if len(found) < 3: + return np.asarray([], dtype=np.float64), float("inf") + positions = np.asarray(sorted(found), dtype=np.float64) + lattice = np.round((positions - positions[0]) / pitch) + design = np.column_stack((lattice, np.ones_like(lattice))) + (slope, intercept), *_ = np.linalg.lstsq(design, positions, rcond=None) + residual = float(np.max(np.abs(positions - (slope * lattice + intercept)))) + return positions, residual + + +def _fit_axis( + profile: np.ndarray, pitch_min: float, pitch_max: float +) -> tuple[float, float, np.ndarray, float, float]: + high = min(float(pitch_max), profile.size / MIN_PERIODS) + low = float(pitch_min) + if high <= low: + raise ValueError( + f"A {profile.size}-pixel axis cannot contain {MIN_PERIODS} periods " + f"between {low:g} and {pitch_max:g} pixels" + ) + initial = _periodogram_pitch(profile, low, high) + if not np.isfinite(initial): + raise ValueError("No periodic illumination peak was found") + refined, prominence = _refine_pitch(profile, initial) + folded = fold(_detrend(profile, refined), refined) + phase = float(np.nanargmin(folded)) / folded.size * refined + seams, residual = _find_seams(profile, refined, phase) + if seams.size >= 3: + pitch = float(np.median(np.diff(seams))) + phase = float(seams[0] % pitch) + else: + pitch = refined + return pitch, phase, seams, residual, prominence + + +def fit_grid( + image: np.ndarray, pitch_min: float = 150.0, pitch_max: float = 1400.0 +) -> TileGrid: + image = as_plane(image) + row, column = median_profiles(image) + py, sy, seams_y, ry, prominence_y = _fit_axis(row, pitch_min, pitch_max) + px, sx, seams_x, rx, prominence_x = _fit_axis(column, pitch_min, pitch_max) + return TileGrid( + pitch_y=py, + pitch_x=px, + seam_y=sy, + seam_x=sx, + height=image.shape[0], + width=image.shape[1], + seams_y=tuple(float(v) for v in seams_y) or None, + seams_x=tuple(float(v) for v in seams_x) or None, + seam_residual_y=ry, + seam_residual_x=rx, + prominence_y=prominence_y, + prominence_x=prominence_x, + ) + + +def as_plane(image: np.ndarray) -> np.ndarray: + plane = np.asarray(image).squeeze() + if plane.ndim != 2: + raise ValueError(f"Expected a 2-D image plane, got shape {plane.shape}") + return plane + + +def _frame_image( + tile_client, dataset_id: str, coordinates: dict, channel: int +) -> np.ndarray: + frame = tile_client.coordinatesToFrameIndex( + coordinates["XY"], coordinates["Z"], coordinates["Time"], channel + ) + return as_plane(tile_client.getRegion(dataset_id, frame=frame)) + + +def _available_channels(tile_client) -> list[int]: + frames = tile_client.tiles.get("frames", []) + channels = sorted({int(frame.get("IndexC", 0)) for frame in frames}) + if channels: + return channels + count = int(tile_client.tiles.get("IndexRange", {}).get("IndexC", 1)) + return list(range(max(count, 1))) + + +def choose_reference_grid( + tile_client, + dataset_id: str, + coordinates: dict, + mode: str, + reference_channel: int, + pitch_min: float, + pitch_max: float, + progress: Callable[[float, str, str], None] | None = None, +) -> tuple[TileGrid, int, list[dict]]: + """Return the best fitted grid from the dominant cross-channel cluster.""" + automatic = mode == "Automatically choose best channel" + channels = ( + _available_channels(tile_client) if automatic else [int(reference_channel)] + ) + reports: list[dict] = [] + fitted: list[tuple[int, TileGrid, dict]] = [] + + for index, channel in enumerate(channels): + if progress: + progress( + index / max(len(channels), 1), + "Illumination correction", + f"Evaluating channel {channel + 1} as the grid reference", + ) + try: + image = _frame_image(tile_client, dataset_id, coordinates, channel) + grid = fit_grid(image, pitch_min=pitch_min, pitch_max=pitch_max) + report = { + "channel": channel, + "valid": bool(grid.is_valid), + **grid.as_dict(), + } + if grid.is_valid: + fitted.append((channel, grid, report)) + except Exception as exc: + report = {"channel": channel, "valid": False, "error": str(exc)} + reports.append(report) + + if not fitted: + detail = "; ".join( + f"channel {r['channel'] + 1}: {r.get('error', 'low-confidence grid')}" + for r in reports + ) + raise ValueError( + "No channel produced a reliable physical tile grid. " + f"Check the reference plane and tile-pitch range. {detail}" + ) + + # A real acquisition grid is shared across channels. Prefer a grid supported by + # the most other channels, then the best individual residual/prominence score. + for channel, grid, report in fitted: + agreement = sum( + abs(other.pitch_y / grid.pitch_y - 1.0) <= 0.075 + and abs(other.pitch_x / grid.pitch_x - 1.0) <= 0.075 + for _, other, _ in fitted + ) + report["cross_channel_agreement"] = int(agreement) + + channel, grid, _ = min( + fitted, + key=lambda item: ( + -item[2]["cross_channel_agreement"], + item[1].quality_score, + item[0], + ), + ) + return grid, channel, reports + + +def normalize_flat(field: np.ndarray) -> np.ndarray: + field = np.asarray(field, dtype=np.float32) + if not np.isfinite(field).all() or np.any(field <= 0): + raise ValueError("A fitted flatfield must be finite and strictly positive") + mean = float(np.mean(field)) + if not np.isfinite(mean) or mean < 1e-12: + raise ValueError("A fitted flatfield must have a positive finite mean") + return (field / mean).astype(np.float32) + + +def tile_stack(image: np.ndarray, grid: TileGrid, n: int = TILE_N) -> np.ndarray: + """Resample complete, seam-to-seam physical tiles into normalized coordinates.""" + sy = grid.complete_seams("y") + sx = grid.complete_seams("x") + fraction = np.linspace(0.0, 1.0, int(n), endpoint=False, dtype=np.float32) + source = np.asarray(image, dtype=np.float32) + output = [] + for y0, y1 in zip(sy[:-1], sy[1:]): + ys = y0 + fraction * (y1 - y0) + for x0, x1 in zip(sx[:-1], sx[1:]): + xs = x0 + fraction * (x1 - x0) + coordinates = np.stack(np.meshgrid(ys, xs, indexing="ij")) + output.append( + ndimage.map_coordinates(source, coordinates, order=1, mode="nearest") + ) + if len(output) < 4: + raise ValueError("At least four complete physical tiles are required") + return np.stack(output).astype(np.float32) + + +def expand_tile_field(field: np.ndarray, grid: TileGrid) -> np.ndarray: + field = np.asarray(field, dtype=np.float32) + if field.ndim != 2: + raise ValueError(f"Expected a 2-D tile field, got {field.shape}") + u = grid.u_of(np.arange(grid.height)) + v = grid.v_of(np.arange(grid.width)) + iy = np.clip((u * field.shape[0]).astype(np.int32), 0, field.shape[0] - 1) + ix = np.clip((v * field.shape[1]).astype(np.int32), 0, field.shape[1] - 1) + return field[np.ix_(iy, ix)].astype(np.float32) + + +def smooth_periodic(field: np.ndarray, sigma: float) -> np.ndarray: + return ndimage.gaussian_filter(field, sigma, mode="wrap") + + +def _tile_gain_map(gains: np.ndarray, grid: TileGrid) -> np.ndarray: + sy = grid.complete_seams("y") + sx = grid.complete_seams("x") + values = np.asarray(gains, dtype=np.float32).reshape(len(sy) - 1, len(sx) - 1) + iy = np.clip( + np.searchsorted(sy, np.arange(grid.height), side="right") - 1, + 0, + values.shape[0] - 1, + ) + ix = np.clip( + np.searchsorted(sx, np.arange(grid.width), side="right") - 1, + 0, + values.shape[1] - 1, + ) + return values[np.ix_(iy, ix)].astype(np.float32) + + +def _estimate_tile_gains(corrected_stack: np.ndarray) -> np.ndarray: + levels = np.percentile(corrected_stack, 25.0, axis=(1, 2)) + median = max(float(np.median(levels)), 1e-9) + gains = np.clip(levels / median, 0.5, 2.0) + return (gains / max(float(np.mean(gains)), 1e-9)).astype(np.float32) + + +class PeriodicFieldModel: + """Compact normalized-tile field expanded only while applying a plane.""" + + def __init__( + self, + name: str, + grid: TileGrid, + flat_tile: np.ndarray, + dark: float | np.ndarray = 0.0, + gains: np.ndarray | None = None, + diagnostics: dict | None = None, + ): + self.name = name + self.grid = grid + self.flat_tile = normalize_flat(flat_tile) + dark_values = np.asarray(dark) + if not np.isfinite(dark_values).all(): + raise ValueError("A fitted darkfield must be finite") + self.dark = dark + if gains is not None: + gains = np.asarray(gains, dtype=np.float32) + if not np.isfinite(gains).all() or np.any(gains <= 0): + raise ValueError("Fitted per-tile gains must be finite and positive") + self.gains = gains + self.diagnostics = diagnostics or {} + + def apply(self, image: np.ndarray) -> np.ndarray: + raw = as_plane(image) + if raw.shape != self.grid.shape: + raise ValueError( + f"Model expects planes of shape {self.grid.shape}, got {raw.shape}" + ) + flat = expand_tile_field(self.flat_tile, self.grid) + if self.gains is not None: + flat *= _tile_gain_map(self.gains, self.grid) + flat = np.maximum(flat, 1e-3) + + output = raw.astype(np.float32, copy=True) + if np.ndim(self.dark) == 0: + dark_mean = float(self.dark) + output -= dark_mean + else: + dark_map = expand_tile_field(np.asarray(self.dark), self.grid) + dark_mean = float(np.mean(self.dark)) + output -= dark_map + output /= flat + output += dark_mean + return output.astype(np.float32) + + +class IdentityModel: + """Explicit no-correction baseline used by automatic selection.""" + + def __init__(self, grid: TileGrid, diagnostics: dict | None = None): + self.name = "identity" + self.grid = grid + self.diagnostics = diagnostics or {} + + def apply(self, image: np.ndarray) -> np.ndarray: + raw = as_plane(image) + if raw.shape != self.grid.shape: + raise ValueError( + f"Model expects planes of shape {self.grid.shape}, got {raw.shape}" + ) + return raw.astype(np.float32, copy=True) + + +class SplitHalfAffineModel: + def __init__( + self, + grid: TileGrid, + gain_y: np.ndarray, + gain_x: np.ndarray, + offset_y: np.ndarray, + offset_x: np.ndarray, + diagnostics: dict, + ): + self.name = "split_half_affine" + self.grid = grid + self.gain_y = gain_y.astype(np.float32) + self.gain_x = gain_x.astype(np.float32) + self.offset_y = offset_y.astype(np.float32) + self.offset_x = offset_x.astype(np.float32) + self.diagnostics = diagnostics + + def apply(self, image: np.ndarray) -> np.ndarray: + raw = as_plane(image) + if raw.shape != self.grid.shape: + raise ValueError( + f"Model expects planes of shape {self.grid.shape}, got {raw.shape}" + ) + flat = expand_tile_field(np.outer(self.gain_y, self.gain_x), self.grid) + offset = expand_tile_field( + self.offset_y[:, None] + self.offset_x[None, :], self.grid + ) + return (raw.astype(np.float32) / np.maximum(flat, 1e-3) - offset).astype( + np.float32 + ) + + +def fit_basic( + raw: np.ndarray, + grid: TileGrid, + *, + darkfield: bool, + per_tile_gain: bool = True, + tile_n: int = TILE_N, +) -> PeriodicFieldModel: + """Fit BaSiCPy to the normalized stack of complete physical tiles.""" + try: + from basicpy import BaSiC + except ImportError as exc: + raise RuntimeError("BaSiCPy is not installed in the worker image") from exc + + stack = tile_stack(raw, grid, n=tile_n) + estimator = BaSiC( + get_darkfield=bool(darkfield), + sparse_cost_darkfield=0.01, + fitting_mode="approximate", + sort_intensity=False, + working_size=128, + max_iterations=500, + ) + estimator.fit(stack.astype(np.float32)) + flat_tile = np.asarray(estimator.flatfield, dtype=np.float32).squeeze() + if flat_tile.shape != (tile_n, tile_n): + raise ValueError( + f"BaSiC returned an unexpected flatfield shape {flat_tile.shape}" + ) + + if darkfield: + dark_tile = np.asarray(estimator.darkfield, dtype=np.float32).squeeze() + if dark_tile.ndim == 0 or dark_tile.size == 1: + dark_tile = np.full_like(flat_tile, float(dark_tile)) + else: + dark_tile = np.zeros_like(flat_tile) + if dark_tile.shape != flat_tile.shape: + raise ValueError( + f"BaSiC returned an unexpected darkfield shape {dark_tile.shape}" + ) + + flat_tile = normalize_flat(flat_tile) + gains = None + if per_tile_gain: + corrected_stack = (stack - dark_tile[None]) / np.maximum(flat_tile[None], 1e-3) + gains = _estimate_tile_gains(corrected_stack) + + name = "basic_darkfield_on" if darkfield else "basic_darkfield_off" + diagnostics = { + "darkfield": bool(darkfield), + "per_tile_gain": bool(per_tile_gain), + "tile_n": int(tile_n), + "n_tiles": int(stack.shape[0]), + "darkfield_mean": float(np.mean(dark_tile)), + "darkfield_range": [float(np.min(dark_tile)), float(np.max(dark_tile))], + "flatfield_range": [float(np.min(flat_tile)), float(np.max(flat_tile))], + "tile_gain_sd_pct": ( + float(100.0 * np.std(gains)) if gains is not None else 0.0 + ), + } + return PeriodicFieldModel( + name, + grid, + flat_tile, + dark=dark_tile, + gains=gains, + diagnostics=diagnostics, + ) + + +def periodic_poisson(gy: np.ndarray, gx: np.ndarray) -> np.ndarray: + rows, columns = gy.shape + ky = np.fft.fftfreq(rows)[:, None] + kx = np.fft.fftfreq(columns)[None, :] + my = np.exp(2j * np.pi * ky) - 1.0 + mx = np.exp(2j * np.pi * kx) - 1.0 + denominator = np.abs(my) ** 2 + np.abs(mx) ** 2 + denominator[0, 0] = 1.0 + numerator = np.conj(my) * np.fft.fft2(gy) + np.conj(mx) * np.fft.fft2(gx) + potential = np.fft.ifft2(numerator / denominator).real + return potential - potential.mean() + + +def fit_log_gradient( + raw: np.ndarray, + grid: TileGrid, + *, + n: int = TILE_N, + smooth_sigma: float = 2.0, + field_smooth: float = 1.5, + per_tile_gain: bool = True, +) -> PeriodicFieldModel: + source = np.asarray(raw) + array = as_plane(source).astype(np.float32) + positive = array[array > 0] + if positive.size == 0: + raise ValueError("Folded log-gradient requires positive image intensities") + floor = max(float(np.percentile(positive, 1)) * 0.05, 1.0) + log_image = np.log(np.maximum(array, floor)) + if smooth_sigma: + log_image = ndimage.gaussian_filter(log_image, float(smooth_sigma)) + + gy = np.empty_like(log_image) + gx = np.empty_like(log_image) + gy[:-1] = np.diff(log_image, axis=0) + gy[-1] = gy[-2] + gx[:, :-1] = np.diff(log_image, axis=1) + gx[:, -1] = gx[:, -2] + + if np.issubdtype(source.dtype, np.integer): + saturation = float(np.iinfo(source.dtype).max) + bad = array >= saturation + if bad.any(): + bad = ndimage.binary_dilation(bad, np.ones((3, 3), dtype=bool)) + gy = np.where(bad, np.nan, gy) + gx = np.where(bad, np.nan, gx) + + with np.errstate(invalid="ignore"): + folded_gy = np.nanmedian(tile_stack(gy, grid, n=n), axis=0) + folded_gx = np.nanmedian(tile_stack(gx, grid, n=n), axis=0) + folded_gy = np.nan_to_num(folded_gy) * (grid.pitch_y / n) + folded_gx = np.nan_to_num(folded_gx) * (grid.pitch_x / n) + + log_field = periodic_poisson(folded_gy, folded_gx) + if field_smooth: + log_field = smooth_periodic(log_field, float(field_smooth)) + flat_tile = normalize_flat(np.exp(log_field).astype(np.float32)) + + gains = None + if per_tile_gain: + stack = tile_stack(array, grid, n=128) + sampled_flat = ndimage.zoom( + flat_tile, + (128 / flat_tile.shape[0], 128 / flat_tile.shape[1]), + order=1, + mode="wrap", + ) + sampled_flat = sampled_flat[:128, :128] + gains = _estimate_tile_gains(stack / np.maximum(sampled_flat[None], 1e-3)) + + diagnostics = { + "per_tile_gain": bool(per_tile_gain), + "tile_n": int(n), + "log_field_range": float(np.exp(log_field.max() - log_field.min())), + "tile_gain_sd_pct": ( + float(100.0 * np.std(gains)) if gains is not None else 0.0 + ), + } + return PeriodicFieldModel( + "fold_log_gradient", + grid, + flat_tile, + dark=0.0, + gains=gains, + diagnostics=diagnostics, + ) + + +LOCATION_QUANTILES = (0.10, 0.25, 0.50) +PROFILE_QUANTILES = (0.10, 0.25, 0.50, 0.75) +BASE_SCALES = np.asarray((3.0, 9.0, 27.0, 81.0, 243.0)) + + +@dataclass +class _Profiles: + location_y: np.ndarray + location_x: np.ndarray + spread_y: np.ndarray + spread_x: np.ndarray + + +def _extract_profiles( + stack: np.ndarray, + gain_y: np.ndarray | None = None, + gain_x: np.ndarray | None = None, +) -> _Profiles: + values = np.asarray(stack, dtype=np.float32).copy() + if gain_y is not None and gain_x is not None: + values /= np.maximum(gain_y[None, :, None] * gain_x[None, None, :], 1e-6) + values -= np.median(values, axis=(1, 2), keepdims=True) + yq = np.quantile(values, PROFILE_QUANTILES, axis=2).transpose(1, 0, 2) + xq = np.quantile(values, PROFILE_QUANTILES, axis=1).transpose(1, 0, 2) + return _Profiles( + location_y=yq[:, : len(LOCATION_QUANTILES)], + location_x=xq[:, : len(LOCATION_QUANTILES)], + spread_y=yq[:, 3] - yq[:, 1], + spread_x=xq[:, 3] - xq[:, 1], + ) + + +def _aggregate_location( + profiles: np.ndarray, indices: np.ndarray | slice = slice(None) +) -> np.ndarray: + quantiles = np.median(profiles[indices], axis=0) + quantiles -= np.mean(quantiles, axis=1, keepdims=True) + result = np.median(quantiles, axis=0) + return result - np.mean(result) + + +def _aggregate_log_spread( + profiles: np.ndarray, indices: np.ndarray | slice = slice(None) +) -> np.ndarray: + spread = np.maximum(np.median(profiles[indices], axis=0), 1e-6) + result = np.log(spread) + return result - np.mean(result) + + +def _split_indices(count: int, split_count: int, seed: int): + if count < 4: + raise ValueError("At least four tile cells are required") + rng = np.random.default_rng(seed) + first_size = (count + 1) // 2 + return [ + (order[:first_size], order[first_size:]) + for order in (rng.permutation(count) for _ in range(split_count)) + ] + + +def _concordance(left: np.ndarray, right: np.ndarray) -> float: + left = left - np.mean(left) + right = right - np.mean(right) + denominator = np.mean(left * left) + np.mean(right * right) + if denominator <= np.finfo(float).eps: + return 0.0 + return float(np.clip(2.0 * np.mean(left * right) / denominator, 0.0, 1.0)) + + +def _multiscale_bands(curve: np.ndarray, scales: np.ndarray) -> list[np.ndarray]: + smoothed = [ndimage.gaussian_filter1d(curve, s, mode="wrap") for s in scales] + return [ + smoothed[-1], + *[ + smoothed[index] - smoothed[index + 1] + for index in range(len(smoothed) - 2, -1, -1) + ], + ] + + +def _denoise_reproducible_curve( + profiles: np.ndarray, + aggregate: Callable, + splits, + scales: np.ndarray, +) -> tuple[np.ndarray, list[dict]]: + full = _multiscale_bands(aggregate(profiles), scales) + halves = [ + ( + _multiscale_bands(aggregate(profiles, left), scales), + _multiscale_bands(aggregate(profiles, right), scales), + ) + for left, right in splits + ] + result = np.zeros_like(full[0]) + diagnostics = [] + for index, band in enumerate(full): + reliability = float( + np.median([_concordance(a[index], b[index]) for a, b in halves]) + ) + if reliability < 0.05: + reliability = 0.0 + shrinkage = 2 * reliability / (1 + reliability) if reliability else 0.0 + result += shrinkage * band + diagnostics.append( + { + "half_sample_reliability": reliability, + "full_sample_shrinkage": shrinkage, + } + ) + return result - np.mean(result), diagnostics + + +def fit_split_half_affine( + raw: np.ndarray, + grid: TileGrid, + *, + profile_size: int = TILE_N, + split_count: int = 16, + seed: int = 2026, + max_gain: float = 2.0, +) -> SplitHalfAffineModel: + stack = tile_stack(raw, grid, n=profile_size) + splits = _split_indices(stack.shape[0], split_count, seed) + scales = np.maximum(0.5, BASE_SCALES * profile_size / 1024.0) + profiles = _extract_profiles(stack) + + log_gain_y, gain_y_diag = _denoise_reproducible_curve( + profiles.spread_y, _aggregate_log_spread, splits, scales + ) + log_gain_x, gain_x_diag = _denoise_reproducible_curve( + profiles.spread_x, _aggregate_log_spread, splits, scales + ) + limit = math.log(max_gain) + gain_y = np.exp(np.clip(log_gain_y, -limit / 2, limit / 2)) + gain_x = np.exp(np.clip(log_gain_x, -limit / 2, limit / 2)) + normalization = math.sqrt(float(np.mean(np.outer(gain_y, gain_x)))) + gain_y /= max(normalization, 1e-9) + gain_x /= max(normalization, 1e-9) + + corrected_profiles = _extract_profiles(stack, gain_y, gain_x) + offset_y, offset_y_diag = _denoise_reproducible_curve( + corrected_profiles.location_y, _aggregate_location, splits, scales + ) + offset_x, offset_x_diag = _denoise_reproducible_curve( + corrected_profiles.location_x, _aggregate_location, splits, scales + ) + center = float(np.median(offset_y[:, None] + offset_x[None, :])) + offset_y -= center / 2 + offset_x -= center / 2 + + diagnostics = { + "profile_size": int(profile_size), + "split_count": int(split_count), + "seed": int(seed), + "gain_range": [ + float(np.min(np.outer(gain_y, gain_x))), + float(np.max(np.outer(gain_y, gain_x))), + ], + "offset_range": [ + float(np.min(offset_y[:, None] + offset_x[None, :])), + float(np.max(offset_y[:, None] + offset_x[None, :])), + ], + "gain_y_bands": gain_y_diag, + "gain_x_bands": gain_x_diag, + "offset_y_bands": offset_y_diag, + "offset_x_bands": offset_x_diag, + } + return SplitHalfAffineModel( + grid, gain_y, gain_x, offset_y, offset_x, diagnostics + ) + + +# Metric constants are fixed to the values used by the illumination study. +BG_PCTL = 10 +UV_POLY_DEGREE = 5 +N_HARMONICS = 4 +PSF_SIGMA = 1.3 +SPOT_K = 5.0 +MIN_SPOT_COUNT = 20 +HF_CUTOFF = 0.10 +HF_CROPS = 4 +HF_CROP = 2048 +HF_LOCAL_SIGMA = 32.0 +MAD_TO_SIGMA = 1.4826 +A1_BINS = 32 + + +def _metric_block(grid: TileGrid, preferred: int = 64) -> int: + return max(4, min(preferred, int(min(grid.pitch_y, grid.pitch_x) / 6))) + + +def block_percentile( + image: np.ndarray, block: int, percentile: float = BG_PCTL +) -> np.ndarray: + height, width = image.shape + block = max(2, min(int(block), height, width)) + hh, ww = (height // block) * block, (width // block) * block + if hh == 0 or ww == 0: + raise ValueError("Image is too small for block-percentile metrics") + cells = image[:hh, :ww].reshape(hh // block, block, ww // block, block) + return np.percentile(cells.astype(np.float32), percentile, axis=(1, 3)) + + +def _poly_basis(u: np.ndarray, v: np.ndarray, degree: int) -> np.ndarray: + return np.column_stack( + [ + (u**du) * (v**dv) + for du in range(degree + 1) + for dv in range(degree + 1 - du) + ] + ) + + +def a1_fold_amplitude(image: np.ndarray, grid: TileGrid) -> dict: + row, column = median_profiles(image) + scale = max(abs(float(np.median(image))), 1e-6) + amplitudes = {} + for name, profile, coordinate in ( + ("y", row, grid.u_of(np.arange(row.size))), + ("x", column, grid.v_of(np.arange(column.size))), + ): + labels = np.clip((coordinate * A1_BINS).astype(int), 0, A1_BINS - 1) + values = np.asarray( + [ + np.median(profile[labels == index]) + if np.any(labels == index) + else np.nan + for index in range(A1_BINS) + ] + ) + amplitudes[name] = float(np.nanmax(values) - np.nanmin(values)) + return { + "A1_fold_amp_y": amplitudes["y"], + "A1_fold_amp_x": amplitudes["x"], + "A1_fold_amp_rel_pct": 100 * max(amplitudes.values()) / scale, + } + + +def a2_harmonic_power(image: np.ndarray, grid: TileGrid) -> dict: + row, column = median_profiles(image) + output = {} + for name, profile, pitch in ( + ("y", row, grid.pitch_y), + ("x", column, grid.pitch_x), + ): + trend = ndimage.gaussian_filter1d(profile, pitch, mode="nearest") + relative = profile / np.maximum(trend, 1e-6) - 1.0 + spectrum = np.fft.rfft(relative) + bins: list[int] = [] + for harmonic in range(1, N_HARMONICS + 1): + center = int(round(harmonic * relative.size / pitch)) + if center < spectrum.size - 2: + bins.extend(range(max(center - 2, 1), min(center + 3, spectrum.size))) + bins = sorted(set(bins)) + if not bins: + output[f"A2_harmonic_mod_pct_{name}"] = float("nan") + continue + selected = np.zeros_like(spectrum) + selected[bins] = spectrum[bins] + output[f"A2_harmonic_mod_pct_{name}"] = float( + 100 * np.std(np.fft.irfft(selected, n=relative.size)) + ) + values = [v for v in output.values() if np.isfinite(v)] + output["A2_harmonic_mod_pct_max"] = max(values) if values else float("nan") + return output + + +def _block_coordinates(grid: TileGrid, shape: tuple[int, int], block: int): + y = (np.arange(shape[0]) + 0.5) * block + x = (np.arange(shape[1]) + 0.5) * block + return grid.u_of(y), grid.v_of(x), grid.tile_of_y(y), grid.tile_of_x(x) + + +def _blocks_spanning_seams(grid: TileGrid, shape: tuple[int, int], block: int): + y = np.arange(shape[0]) * block + x = np.arange(shape[1]) * block + return ( + grid.tile_of_y(y) != grid.tile_of_y(y + block - 1), + grid.tile_of_x(x) != grid.tile_of_x(x + block - 1), + ) + + +def a3_uv_dependence(image: np.ndarray, grid: TileGrid) -> dict: + block = _metric_block(grid) + background = block_percentile(image, block) + u, v, tile_y, tile_x = _block_coordinates(grid, background.shape, block) + span_y, span_x = _blocks_spanning_seams(grid, background.shape, block) + keep = ~(span_y[:, None] | span_x[None, :]) + U = np.broadcast_to(u[:, None], background.shape)[keep] + V = np.broadcast_to(v[None, :], background.shape)[keep] + parity = (np.add.outer(tile_y, tile_x) % 2).astype(bool)[keep] + values = np.log(np.maximum(background, 1e-6))[keep] + basis = _poly_basis(U, V, UV_POLY_DEGREE) + train = ~parity + if train.sum() < basis.shape[1] * 3 or parity.sum() < 10: + return { + "A3_uv_var_explained": float("nan"), + "A3_uv_modulation_pct": float("nan"), + } + coefficients, *_ = np.linalg.lstsq(basis[train], values[train], rcond=None) + predicted = basis @ coefficients + residual = values[parity] - predicted[parity] + variance = float(np.var(values[parity])) + return { + "A3_uv_var_explained": float( + 1 - np.var(residual) / max(variance, 1e-12) + ), + "A3_uv_modulation_pct": float(100 * np.std(predicted[parity])), + } + + +def a5_background_range(image: np.ndarray, grid: TileGrid) -> dict: + background = block_percentile(image, _metric_block(grid)) + low, high = np.percentile(background, [1, 99]) + return { + "A5_bg_range_ratio": float(high / max(low, 1e-6)), + "A5_bg_p1": float(low), + "A5_bg_p99": float(high), + } + + +def a6_tile_level(image: np.ndarray, grid: TileGrid) -> dict: + block = _metric_block(grid, preferred=32) + background = block_percentile(image, block) + u, v, tile_y, tile_x = _block_coordinates(grid, background.shape, block) + span_y, span_x = _blocks_spanning_seams(grid, background.shape, block) + keep_y = (u > 0.2) & (u < 0.8) & ~span_y + keep_x = (v > 0.2) & (v < 0.8) & ~span_x + levels = {} + if keep_y.sum() >= 2 and keep_x.sum() >= 2: + subset = background[np.ix_(keep_y, keep_x)] + iy, ix = tile_y[keep_y], tile_x[keep_x] + for y_index in np.unique(iy): + for x_index in np.unique(ix): + cell = subset[np.ix_(iy == y_index, ix == x_index)] + if cell.size >= 4: + levels[(int(y_index), int(x_index))] = float(np.median(cell)) + if len(levels) < 10: + return {"A6_tile_level_sd_pct": float("nan"), "A6_n_tiles": len(levels)} + + keys = np.asarray(list(levels)) + values = np.asarray([levels[tuple(key)] for key in keys]) + if np.any(values <= 0): + return {"A6_tile_level_sd_pct": float("nan"), "A6_n_tiles": len(levels)} + y = keys[:, 0].astype(float) + x = keys[:, 1].astype(float) + y = (y - y.mean()) / max(y.std(), 1e-9) + x = (x - x.mean()) / max(x.std(), 1e-9) + basis = np.column_stack( + [y**a * x**b for a in range(3) for b in range(3 - a)] + ) + log_values = np.log(values) + coefficients, *_ = np.linalg.lstsq(basis, log_values, rcond=None) + residual = log_values - basis @ coefficients + return { + "A6_tile_level_sd_pct": float(100 * np.std(residual)), + "A6_n_tiles": int(values.size), + } + + +def _robust_sigma(values: np.ndarray) -> float: + return float( + MAD_TO_SIGMA * np.median(np.abs(values - np.median(values))) + 1e-12 + ) + + +def p1_spots(image: np.ndarray, grid: TileGrid) -> dict: + values = np.asarray(image, dtype=np.float32) + response = -ndimage.gaussian_laplace(values, PSF_SIGMA) + threshold = SPOT_K * _robust_sigma(response[::4, ::4]) + peaks = (response >= threshold) & ( + ndimage.maximum_filter(response, size=3) == response + ) + y, x = np.nonzero(peaks) + if y.size < MIN_SPOT_COUNT: + return { + "P1_spot_count": int(y.size), + "P1_spot_uniformity": float("nan"), + "P1_applicable": False, + } + u, v = grid.u_of(y), grid.v_of(x) + inner = (np.abs(u - 0.5) < 0.25) & (np.abs(v - 0.5) < 0.25) + outer = (np.abs(u - 0.5) > 0.40) | (np.abs(v - 0.5) > 0.40) + # Jeffreys-style pseudocounts keep the ratio finite and symmetric when + # either region has no detections. A zero outer count must not look like an + # unavailable metric while a large outer count is penalized. + density_inner = (inner.sum() + 0.5) / 0.25 + density_outer = (outer.sum() + 0.5) / 0.36 + return { + "P1_spot_count": int(y.size), + "P1_spot_uniformity": float(density_outer / max(density_inner, 1e-9)), + "P1_applicable": True, + } + + +def build_object_mask( + raw: np.ndarray, max_objects: int = 2000, min_area: int = 50 +) -> tuple[np.ndarray, int]: + values = np.asarray(raw, dtype=np.float32) + best = (np.zeros(values.shape, dtype=np.int32), 0) + for percentile in (99.5, 99.0, 98.0, 95.0, 90.0): + labels, count = ndimage.label(values > np.percentile(values, percentile)) + if count == 0: + continue + areas = np.bincount(labels.ravel()) + areas[0] = 0 + keep = np.nonzero(areas >= min_area)[0] + if keep.size > max_objects: + keep = keep[np.argsort(areas[keep])[::-1][:max_objects]] + remap = np.zeros(areas.size, dtype=np.int32) + remap[keep] = np.arange(1, keep.size + 1) + best = remap[labels], int(keep.size) + if keep.size >= 10: + break + return best + + +def p2_object_intensity( + raw: np.ndarray, + corrected: np.ndarray, + labels: np.ndarray, + count: int, +) -> dict: + if count < 10: + return { + "P2_n_objects": count, + "P2_spearman": float("nan"), + "P2_applicable": False, + } + indices = np.arange(1, count + 1) + raw_sum = ndimage.sum_labels(raw.astype(np.float64), labels, indices) + corrected_sum = ndimage.sum_labels( + corrected.astype(np.float64), labels, indices + ) + keep = (raw_sum > 0) & (corrected_sum > 0) + if keep.sum() < 10: + return { + "P2_n_objects": int(keep.sum()), + "P2_spearman": float("nan"), + "P2_applicable": False, + } + from scipy.stats import spearmanr + + return { + "P2_n_objects": int(keep.sum()), + "P2_spearman": float(spearmanr(raw_sum[keep], corrected_sum[keep]).statistic), + "P2_applicable": True, + } + + +def _local_normalize(image: np.ndarray) -> np.ndarray: + smooth = ndimage.gaussian_filter(image, HF_LOCAL_SIGMA, mode="nearest") + return image / np.maximum(smooth, 1e-6) + + +def p3_high_frequency(raw: np.ndarray, corrected: np.ndarray) -> dict: + height, width = raw.shape + crop = min(HF_CROP, height, width) + if crop < 8: + return { + "P3_hf_power_ratio": float("nan"), + "P3_applicable": False, + "P3_n_crops": 0, + } + fy = np.fft.fftfreq(crop) + fx = np.fft.rfftfreq(crop) + mask = np.sqrt(fy[:, None] ** 2 + fx[None, :] ** 2) > HF_CUTOFF + rng = np.random.default_rng(0) + ratios = [] + for _ in range(HF_CROPS): + y = int(rng.integers(0, max(height - crop + 1, 1))) + x = int(rng.integers(0, max(width - crop + 1, 1))) + before = _local_normalize( + raw[y : y + crop, x : x + crop].astype(np.float32) + ) + after = _local_normalize( + corrected[y : y + crop, x : x + crop].astype(np.float32) + ) + power_before = float((np.abs(np.fft.rfft2(before)) ** 2)[mask].sum()) + power_after = float((np.abs(np.fft.rfft2(after)) ** 2)[mask].sum()) + if np.isfinite(power_before) and power_before > 1e-20: + ratios.append(power_after / power_before) + if not ratios: + return { + "P3_hf_power_ratio": float("nan"), + "P3_applicable": False, + "P3_n_crops": 0, + } + return { + "P3_hf_power_ratio": float(np.mean(ratios)), + "P3_applicable": True, + "P3_n_crops": len(ratios), + } + + +def p5_range(raw: np.ndarray, corrected: np.ndarray) -> dict: + raw_values = np.asarray(raw) + corrected_values = np.asarray(corrected) + finite = np.isfinite(corrected_values) + source_finite = np.isfinite(raw_values) + finite_values = corrected_values[finite] + return { + "P5_frac_source_nonfinite": float( + np.count_nonzero(~source_finite) / raw_values.size + ), + "P5_frac_nonfinite": float( + np.count_nonzero(~finite) / corrected_values.size + ), + "P5_frac_nonpositive": float( + np.count_nonzero(finite & (corrected_values <= 0)) + / corrected_values.size + ), + "P5_frac_new_nonpositive": float( + np.count_nonzero( + source_finite + & (raw_values > 0) + & finite + & (corrected_values <= 0) + ) + / corrected_values.size + ), + "P5_min": ( + float(np.min(finite_values)) if finite_values.size else float("nan") + ), + "P5_max": ( + float(np.max(finite_values)) if finite_values.size else float("nan") + ), + } + + +GUARDRAILS = { + "P2_spearman": ("min", 0.98), + "P3_hf_power_ratio": ("min", 0.90), + "P5_frac_source_nonfinite": ("max", 0.0), + "P5_frac_nonfinite": ("max", 0.0), + "P5_frac_new_nonpositive": ("max", 1e-4), +} +CONDITIONAL_GUARDRAILS = { + "P2_spearman": "P2_applicable", + "P3_hf_power_ratio": "P3_applicable", +} + + +def check_guardrails(metrics: dict) -> list[str]: + violations = [] + for key, (kind, limit) in GUARDRAILS.items(): + applicable_key = CONDITIONAL_GUARDRAILS.get(key) + if applicable_key is not None and metrics.get(applicable_key) is False: + continue + value = metrics.get(key) + if value is None or not np.isfinite(value): + violations.append(f"{key} is unavailable or non-finite") + continue + if kind == "min" and value < limit: + violations.append(f"{key}={value:.4g} < {limit}") + if kind == "max" and value > limit: + violations.append(f"{key}={value:.4g} > {limit}") + return violations + + +def unavailable_guardrails(metrics: dict) -> list[str]: + return [ + key + for key, applicable_key in CONDITIONAL_GUARDRAILS.items() + if metrics.get(applicable_key) is False + ] + + +def preservation_metrics( + raw: np.ndarray, + corrected: np.ndarray, + labels: np.ndarray | None = None, + count: int = 0, +) -> dict: + if labels is None: + labels, count = build_object_mask(raw) + metrics = {} + metrics.update(p2_object_intensity(raw, corrected, labels, count)) + metrics.update(p3_high_frequency(raw, corrected)) + metrics.update(p5_range(raw, corrected)) + metrics["guardrail_violations"] = check_guardrails(metrics) + metrics["guardrail_unavailable"] = unavailable_guardrails(metrics) + return metrics + + +def evaluate( + corrected: np.ndarray, + raw: np.ndarray, + grid: TileGrid, + labels: np.ndarray | None = None, + count: int = 0, +) -> dict: + if labels is None: + labels, count = build_object_mask(raw) + metrics = {} + metrics.update(a1_fold_amplitude(corrected, grid)) + metrics.update(a2_harmonic_power(corrected, grid)) + metrics.update(a3_uv_dependence(corrected, grid)) + metrics.update(a5_background_range(corrected, grid)) + metrics.update(a6_tile_level(corrected, grid)) + metrics.update(p1_spots(corrected, grid)) + metrics.update(preservation_metrics(raw, corrected, labels, count)) + return metrics + + +ARTIFACT_KEYS = ( + "A1_fold_amp_rel_pct", + "A2_harmonic_mod_pct_max", + "A3_uv_modulation_pct", + "A5_bg_range_ratio", + "A6_tile_level_sd_pct", +) +NOISE_FLOORS = { + "A1_fold_amp_rel_pct": 0.30, + "A2_harmonic_mod_pct_max": 0.61, + "A3_uv_modulation_pct": 0.93, + "A6_tile_level_sd_pct": 0.30, +} + + +def artifact_ratios(metrics: dict, baseline: dict) -> dict[str, float]: + ratios = {} + for key in ARTIFACT_KEYS: + value, raw_value = metrics.get(key), baseline.get(key) + if value is None or raw_value is None: + continue + if not np.isfinite(value) or not np.isfinite(raw_value): + continue + if key == "A5_bg_range_ratio": + value, raw_value = value - 1.0, raw_value - 1.0 + if raw_value <= 1e-9: + continue + ratios[key] = float( + max(value, NOISE_FLOORS.get(key, 1e-9)) / raw_value + ) + return ratios + + +def artifact_index(metrics: dict, baseline: dict) -> float: + ratios = list(artifact_ratios(metrics, baseline).values()) + if not ratios: + return float("nan") + return float(np.exp(np.mean(np.log(ratios)))) + + +def physics_violations(model, raw: np.ndarray) -> list[str]: + violations = [] + if isinstance(model, PeriodicFieldModel): + flat = model.flat_tile + if not np.isfinite(flat).all() or float(np.min(flat)) <= 0: + violations.append("flatfield is non-finite or nonpositive") + elif float(np.max(flat) / max(np.min(flat), 1e-9)) > 20: + violations.append("flatfield dynamic range exceeds 20x") + if np.ndim(model.dark) > 0 and np.any(np.asarray(model.dark)): + dark_mean = float(np.mean(model.dark)) + image_floor = float(np.percentile(raw, 1)) + if dark_mean < 0: + violations.append("darkfield mean is negative") + if dark_mean >= image_floor: + violations.append( + f"darkfield mean {dark_mean:.4g} reaches the image floor " + f"{image_floor:.4g}" + ) + return violations + + +@dataclass +class CandidateResult: + name: str + model: object + metrics: dict + artifact_index: float + violations: list[str] + physics_violations: list[str] + complexity: int + fit_seconds: float = 0.0 + selection_score: float = float("inf") + alternatives: list[dict] = field(default_factory=list) + artifact_ratios: dict[str, float] = field(default_factory=dict) + selection_samples: list[dict] = field(default_factory=list) + score_log_se: float = 0.0 + pareto_optimal: bool = False + fit_metrics: dict = field(default_factory=dict) + validation_reports: list[dict] = field(default_factory=list) + + @property + def valid(self) -> bool: + return not self.violations and not self.physics_violations + + def summary(self) -> dict: + return { + "name": self.name, + "artifact_index": float(self.artifact_index), + "selection_score": float(self.selection_score), + "fit_seconds": float(self.fit_seconds), + "valid": bool(self.valid), + "pareto_optimal": bool(self.pareto_optimal), + "score_log_se": float(self.score_log_se), + "guardrail_violations": list(self.violations), + "physics_violations": list(self.physics_violations), + "artifact_ratios": dict(self.artifact_ratios), + "metrics": dict(self.metrics), + "fit_metrics": dict(self.fit_metrics), + "validation_reports": list(self.validation_reports), + } + + +def _selection_score(index: float, spot_uniformity, use_spot_uniformity: bool): + if not np.isfinite(index) or index <= 0: + return float("inf") + spot_penalty = 0.0 + if ( + use_spot_uniformity + and spot_uniformity is not None + and np.isfinite(spot_uniformity) + ): + bounded = max(float(spot_uniformity), 1e-6) + spot_penalty = 0.5 * abs(math.log(bounded)) + return float(index * math.exp(spot_penalty)) + + +def _dominates(left: CandidateResult, right: CandidateResult) -> bool: + keys = sorted(left.artifact_ratios) + if not keys or set(keys) != set(right.artifact_ratios): + return False + no_worse = all( + left.artifact_ratios[key] <= right.artifact_ratios[key] for key in keys + ) + strictly_better = any( + left.artifact_ratios[key] < right.artifact_ratios[key] for key in keys + ) + return no_worse and strictly_better + + +def rank_candidates( + candidates: Iterable[CandidateResult], + tie_fraction: float = 0.05, + use_spot_uniformity: bool = False, +) -> tuple[CandidateResult, list[CandidateResult]]: + valid = [candidate for candidate in candidates if candidate.valid] + if not valid: + raise ValueError( + "Every correction candidate failed a preservation or physics check" + ) + for candidate in valid: + candidate.selection_score = _selection_score( + float(candidate.artifact_index), + candidate.metrics.get("P1_spot_uniformity"), + use_spot_uniformity, + ) + sample_scores = [ + _selection_score( + float(sample.get("artifact_index", float("nan"))), + sample.get("P1_spot_uniformity"), + use_spot_uniformity, + ) + for sample in candidate.selection_samples + ] + sample_scores = [score for score in sample_scores if np.isfinite(score)] + if len(sample_scores) >= 2: + candidate.score_log_se = float( + np.std(np.log(sample_scores), ddof=1) / math.sqrt(len(sample_scores)) + ) + + require_complete_panel = any(candidate.name == "identity" for candidate in valid) + scorable = [ + candidate + for candidate in valid + if np.isfinite(candidate.selection_score) + and ( + not require_complete_panel + or len(candidate.artifact_ratios) >= 3 + or not candidate.artifact_ratios + ) + ] + if not scorable: + raise ValueError("No safe candidate produced a finite artifact score") + + frontier = [ + candidate + for candidate in scorable + if not any( + other is not candidate and _dominates(other, candidate) + for other in scorable + ) + ] + for candidate in frontier: + candidate.pareto_optimal = True + + ranked = sorted( + scorable, + key=lambda c: ( + not c.pareto_optimal, + c.selection_score, + c.complexity, + c.name, + ), + ) + selectable = sorted( + frontier, key=lambda c: (c.selection_score, c.complexity, c.name) + ) + best = selectable[0] + best_score = best.selection_score + base_margin = math.log1p(tie_fraction) + tied = [ + candidate + for candidate in selectable + if math.log(candidate.selection_score / best_score) + <= max( + base_margin, + 1.96 + * math.hypot(best.score_log_se, candidate.score_log_se), + ) + ] + selected = min(tied, key=lambda c: (c.complexity, c.selection_score, c.name)) + return selected, ranked + + +def _candidate_specs(algorithm: str, darkfield_mode: str): + if algorithm == ALGORITHM_OPTIONS[0]: + return [ + ("identity", -1), + ("basic_darkfield_off", 2), + ("basic_darkfield_on", 3), + ("fold_log_gradient", 1), + ("split_half_affine", 0), + ] + if algorithm == ALGORITHM_OPTIONS[1]: + if darkfield_mode == DARKFIELD_OPTIONS[1]: + return [("basic_darkfield_on", 3)] + if darkfield_mode == DARKFIELD_OPTIONS[2]: + return [("basic_darkfield_off", 2)] + return [("basic_darkfield_off", 2), ("basic_darkfield_on", 3)] + if algorithm == ALGORITHM_OPTIONS[2]: + return [("fold_log_gradient", 1)] + if algorithm == ALGORITHM_OPTIONS[3]: + return [("split_half_affine", 0)] + raise ValueError(f"Unknown illumination-correction algorithm: {algorithm}") + + +def _aggregate_ratios(reports: list[dict]) -> dict[str, float]: + output = {} + for key in ARTIFACT_KEYS: + values = [ + report["artifact_ratios"].get(key) + for report in reports + if report["artifact_ratios"].get(key) is not None + and np.isfinite(report["artifact_ratios"].get(key)) + and report["artifact_ratios"].get(key) > 0 + ] + if values: + output[key] = float(np.exp(np.mean(np.log(values)))) + return output + + +def _aggregate_selection_metrics(reports: list[dict], basis: str) -> dict: + output = {"selection_basis": basis, "selection_plane_count": len(reports)} + keys = sorted(set().union(*(report["metrics"] for report in reports))) + for key in keys: + values = [report["metrics"].get(key) for report in reports] + numeric = [ + float(value) + for value in values + if isinstance(value, (int, float, np.integer, np.floating)) + and not isinstance(value, (bool, np.bool_)) + and np.isfinite(value) + ] + if numeric: + output[key] = float(np.median(numeric)) + output["guardrail_violations"] = sorted( + { + violation + for report in reports + for violation in report["metrics"].get("guardrail_violations", []) + } + ) + output["guardrail_unavailable"] = sorted( + { + unavailable + for report in reports + for unavailable in report["metrics"].get("guardrail_unavailable", []) + } + ) + return output + + +def _evaluate_model_plane( + model, raw: np.ndarray, grid: TileGrid, label: str +) -> dict: + plane = as_plane(raw) + if not np.isfinite(plane).all(): + raise ValueError(f"{label} contains non-finite source pixels") + labels, count = build_object_mask(plane) + baseline = evaluate(plane, plane, grid, labels, count) + corrected = model.apply(plane) + metrics = evaluate(corrected, plane, grid, labels, count) + ratios = artifact_ratios(metrics, baseline) + index = artifact_index(metrics, baseline) + return { + "label": label, + "metrics": metrics, + "artifact_ratios": ratios, + "artifact_index": index, + } + + +def select_model( + raw: np.ndarray, + grid: TileGrid, + algorithm: str, + darkfield_mode: str, + per_tile_gain: bool, + progress: Callable[[float, str, str], None] | None = None, + validation_source: Callable[ + [], Iterable[tuple[str, np.ndarray]] + ] + | None = None, + use_spot_uniformity: bool = False, +) -> CandidateResult: + """Fit on one plane and select on independent planes when available.""" + raw = as_plane(raw) + if not np.isfinite(raw).all(): + raise ValueError("The reference plane contains non-finite source pixels") + labels, count = build_object_mask(raw) + baseline = evaluate(raw, raw, grid, labels, count) + validation_planes = ( + list(validation_source()) if validation_source is not None else [] + ) + + if algorithm == ALGORITHM_OPTIONS[0] and not validation_planes: + metrics = dict(baseline) + metrics["selection_basis"] = "identity_without_holdout" + model = IdentityModel( + grid, + diagnostics={ + "selection_basis": "identity_without_holdout", + "reason": ( + "Automatic correction requires an independent Z plane; " + "the channel was left unchanged" + ), + }, + ) + identity = CandidateResult( + name="identity", + model=model, + metrics=metrics, + artifact_index=1.0, + violations=list(metrics["guardrail_violations"]), + physics_violations=[], + complexity=-1, + artifact_ratios={key: 1.0 for key in ARTIFACT_KEYS}, + fit_metrics=dict(metrics), + pareto_optimal=True, + selection_score=1.0, + ) + identity.alternatives = [identity.summary()] + return identity + + specs = _candidate_specs(algorithm, darkfield_mode) + candidates: list[CandidateResult] = [] + failures = [] + + for index, (name, complexity) in enumerate(specs): + if progress: + progress( + index / max(len(specs), 1), + "Illumination correction", + f"Fitting {name.replace('_', ' ')}", + ) + started = time.monotonic() + try: + if name == "identity": + model = IdentityModel(grid) + elif name == "basic_darkfield_off": + model = fit_basic( + raw, grid, darkfield=False, per_tile_gain=per_tile_gain + ) + elif name == "basic_darkfield_on": + model = fit_basic( + raw, grid, darkfield=True, per_tile_gain=per_tile_gain + ) + elif name == "fold_log_gradient": + model = fit_log_gradient(raw, grid, per_tile_gain=per_tile_gain) + else: + model = fit_split_half_affine(raw, grid) + fit_report = _evaluate_model_plane(model, raw, grid, "fit plane") + validation_reports = ( + [ + _evaluate_model_plane(model, plane, grid, label) + for label, plane in validation_planes + ] + if validation_planes + else [] + ) + selection_reports = validation_reports or [fit_report] + ratios = _aggregate_ratios(selection_reports) + index_value = ( + float(np.exp(np.mean(np.log(list(ratios.values()))))) + if ratios + else float("nan") + ) + basis = "held_out_z" if validation_reports else "fit_plane" + metrics = _aggregate_selection_metrics(selection_reports, basis) + violations = [ + f"{fit_report['label']}: {violation}" + for violation in fit_report["metrics"]["guardrail_violations"] + ] + violations.extend( + f"{report['label']}: {violation}" + for report in validation_reports + for violation in report["metrics"]["guardrail_violations"] + ) + candidates.append( + CandidateResult( + name=name, + model=model, + metrics=metrics, + artifact_index=index_value, + violations=violations, + physics_violations=physics_violations(model, raw), + complexity=complexity, + fit_seconds=time.monotonic() - started, + artifact_ratios=ratios, + selection_samples=[ + { + "artifact_index": report["artifact_index"], + "P1_spot_uniformity": report["metrics"].get( + "P1_spot_uniformity" + ), + } + for report in selection_reports + ], + fit_metrics=fit_report["metrics"], + validation_reports=validation_reports, + ) + ) + except Exception as exc: + failures.append(f"{name}: {exc}") + + if not candidates: + raise ValueError( + "No correction candidate could be fitted. " + "; ".join(failures) + ) + try: + selected, _ = rank_candidates( + candidates, use_spot_uniformity=use_spot_uniformity + ) + except ValueError as exc: + detail = "; ".join( + f"{candidate.name}: " + + ", ".join(candidate.violations + candidate.physics_violations) + for candidate in candidates + ) + raise ValueError(f"{exc}. {detail}") from exc + + selected.alternatives = [candidate.summary() for candidate in candidates] + if failures: + selected.alternatives.extend( + {"name": failure.split(":", 1)[0], "valid": False, "error": failure} + for failure in failures + ) + return selected diff --git a/workers/annotations/illumination_correction/tests/Dockerfile_Test b/workers/annotations/illumination_correction/tests/Dockerfile_Test new file mode 100644 index 0000000..70ee3a9 --- /dev/null +++ b/workers/annotations/illumination_correction/tests/Dockerfile_Test @@ -0,0 +1,11 @@ +FROM annotations/illumination_correction:latest AS test + +SHELL ["conda", "run", "-n", "worker", "/bin/bash", "-c"] +RUN pip install pytest pytest-mock + +RUN mkdir -p /tests +COPY ./workers/annotations/illumination_correction/tests/*.py /tests +WORKDIR /tests +ENV PYTHONPATH=/ + +ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "worker", "python3", "-m", "pytest", "-v"] diff --git a/workers/annotations/illumination_correction/tests/__init__.py b/workers/annotations/illumination_correction/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/workers/annotations/illumination_correction/tests/test_entrypoint.py b/workers/annotations/illumination_correction/tests/test_entrypoint.py new file mode 100644 index 0000000..428b0b6 --- /dev/null +++ b/workers/annotations/illumination_correction/tests/test_entrypoint.py @@ -0,0 +1,442 @@ +import sys +import types +from pathlib import Path +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + + +WORKER_DIR = Path(__file__).resolve().parents[1] +if str(WORKER_DIR) not in sys.path: + sys.path.insert(0, str(WORKER_DIR)) + + +# Native development runs do not install the NimbusImage client packages. The +# Docker test image does, so only provide lightweight import shims when absent. +try: + import annotation_client.tiles # noqa: F401 +except (ImportError, ModuleNotFoundError): + annotation_client = types.ModuleType("annotation_client") + annotation_client.__path__ = [] + tiles_module = types.ModuleType("annotation_client.tiles") + workers_module = types.ModuleType("annotation_client.workers") + utils_module = types.ModuleType("annotation_client.utils") + tiles_module.UPennContrastDataset = MagicMock() + workers_module.UPennContrastWorkerPreviewClient = MagicMock() + utils_module.sendError = MagicMock() + utils_module.sendProgress = MagicMock() + utils_module.sendWarning = MagicMock() + sys.modules.update( + { + "annotation_client": annotation_client, + "annotation_client.tiles": tiles_module, + "annotation_client.workers": workers_module, + "annotation_client.utils": utils_module, + } + ) + +try: + import annotation_utilities.annotation_tools # noqa: F401 +except (ImportError, ModuleNotFoundError): + annotation_utilities = types.ModuleType("annotation_utilities") + annotation_utilities.__path__ = [] + annotation_tools = types.ModuleType("annotation_utilities.annotation_tools") + + def get_selected_channels(value, field_name="channel selection"): + if value in (None, "", {}): + return [] + if not isinstance(value, dict): + raise ValueError(f"{field_name} must be a mapping") + return sorted(int(key) for key, selected in value.items() if selected) + + def get_required_select(value, field_name, allowed_values=None): + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"The '{field_name}' setting has no valid value") + if allowed_values is not None and value not in allowed_values: + raise ValueError(f"The '{field_name}' setting is stale: {value!r}") + return value + + def split_channel_selection(selected_channels, num_channels): + present = sorted( + {value for value in selected_channels if 0 <= value < num_channels} + ) + missing = sorted({value for value in selected_channels if value not in present}) + return present, missing + + def get_frame_index(frame, dimension, default=0): + key = dimension if dimension.startswith("Index") else f"Index{dimension}" + return frame.get(key, default) + + def frame_to_large_image_params(frame): + return { + key.lower()[5:]: value + for key, value in frame.items() + if key.startswith("Index") and len(key) > 5 + } + + annotation_tools.get_selected_channels = get_selected_channels + annotation_tools.get_required_select = get_required_select + annotation_tools.split_channel_selection = split_channel_selection + annotation_tools.get_frame_index = get_frame_index + annotation_tools.frame_to_large_image_params = frame_to_large_image_params + sys.modules.update( + { + "annotation_utilities": annotation_utilities, + "annotation_utilities.annotation_tools": annotation_tools, + } + ) + +from illumination import CandidateResult, TileGrid # noqa: E402 +from entrypoint import compute, interface # noqa: E402 + + +def _grid(shape=(64, 64)): + return TileGrid( + pitch_y=16.0, + pitch_x=16.0, + seam_y=0.0, + seam_x=0.0, + height=shape[0], + width=shape[1], + seams_y=(0.0, 16.0, 32.0, 48.0, 64.0), + seams_x=(0.0, 16.0, 32.0, 48.0, 64.0), + seam_residual_y=0.0, + seam_residual_x=0.0, + ) + + +class _AddOneModel: + name = "fold_log_gradient" + diagnostics = {"test": True} + + def apply(self, image): + return np.asarray(image, dtype=np.float32) + 1.0 + + +def _selection(): + return CandidateResult( + name="fold_log_gradient", + model=_AddOneModel(), + metrics={ + "A1_fold_amp_rel_pct": 1.0, + "P1_spot_uniformity": 1.0, + "P2_spearman": 0.99, + "P3_hf_power_ratio": 1.0, + "P5_frac_nonpositive": 0.0, + }, + artifact_index=0.25, + violations=[], + physics_violations=[], + complexity=1, + ) + + +def _params(channels=None): + return { + "channel": 0, + "tile": {"XY": 0, "Z": 0, "Time": 0}, + "workerInterface": { + "Channels to correct": channels or {"0": True, "1": False}, + "Punctate channels for spot metric": {}, + "Algorithm": "Automatic (recommended)", + "Reference channel mode": "Automatically choose best channel", + "Reference channel": 0, + "Reference XY": "", + "Reference Z": "", + "Reference Time": "", + "BaSiC darkfield": "Automatic", + "Per-tile gain correction": False, + "Output type": "Float32 (recommended)", + "Validate every corrected plane": False, + "Minimum tile pitch": 10, + "Maximum tile pitch": 30, + }, + } + + +def _tile_client(frames=True): + client = MagicMock() + client.tiles = { + "IndexRange": {"IndexXY": 1, "IndexZ": 1, "IndexT": 1, "IndexC": 2}, + "channels": ["DAPI", "YFP"], + "mm_x": 0.001, + "mm_y": 0.001, + "magnification": 20, + "dtype": np.dtype("uint16"), + } + if frames: + client.tiles["frames"] = [ + {"IndexXY": 0, "IndexZ": 0, "IndexT": 0, "IndexC": 0}, + {"IndexXY": 0, "IndexZ": 0, "IndexT": 0, "IndexC": 1}, + ] + client.coordinatesToFrameIndex.side_effect = lambda xy, z, time, channel: channel + client.getRegion.side_effect = lambda dataset_id, frame: np.full( + (64, 64, 1), 100 + frame, dtype=np.uint16 + ) + client.client.uploadFileToFolder.return_value = {"itemId": "output-id"} + return client + + +def test_interface_exposes_auto_and_manual_controls(): + with patch("entrypoint.workers.UPennContrastWorkerPreviewClient") as preview_cls: + interface("image-id", "http://api", "token") + + values = preview_cls.return_value.setWorkerImageInterface.call_args.args[1] + assert values["Illumination correction"]["type"] == "notes" + assert values["Channels to correct"]["type"] == "channelCheckboxes" + assert values["Algorithm"]["items"] == [ + "Automatic (recommended)", + "BaSiC", + "Folded log-gradient", + "Split-half affine", + ] + assert values["Reference channel mode"]["type"] == "select" + assert values["Reference channel"]["type"] == "channel" + assert values["Punctate channels for spot metric"]["type"] == "channelCheckboxes" + assert values["Per-tile gain correction"]["default"] is False + assert values["Output type"]["default"] == "Float32 (recommended)" + assert sorted(v["displayOrder"] for v in values.values()) == list( + range(len(values)) + ) + + +def test_compute_corrects_selected_channels_and_uploads_tiff(): + tile_client = _tile_client() + sink = MagicMock() + fake_large_image = types.SimpleNamespace(new=MagicMock(return_value=sink)) + reference = (_grid(), 1, [{"channel": 1, "quality_score": 0.1}]) + + with ( + patch("entrypoint.tiles.UPennContrastDataset", return_value=tile_client), + patch("entrypoint.correction.choose_reference_grid", return_value=reference), + patch( + "entrypoint.correction.select_model", return_value=_selection() + ) as select_model, + patch.dict(sys.modules, {"large_image": fake_large_image}), + ): + compute("dataset-id", "http://api", "token", _params()) + + select_model.assert_called_once() + assert sink.addTile.call_count == 2 + corrected = sink.addTile.call_args_list[0].args[0] + untouched = sink.addTile.call_args_list[1].args[0] + assert corrected.dtype == np.float32 + assert np.all(corrected == 101.0) + assert untouched.dtype == np.float32 + assert np.all(untouched == 101) + sink.write.assert_called_once_with("/tmp/illumination_corrected.tiff") + tile_client.client.uploadFileToFolder.assert_called_once_with( + "dataset-id", "/tmp/illumination_corrected.tiff" + ) + metadata = tile_client.client.addMetadataToItem.call_args.args[1] + assert metadata["tool"] == "Illumination correction" + assert metadata["reference_channel_zero_based"] == 1 + assert metadata["reference_channel_one_based"] == 2 + assert metadata["reference_coordinates_zero_based"] == { + "XY": 0, + "Z": 0, + "Time": 0, + } + assert metadata["correction_scope"]["Z"].startswith("all planes") + assert metadata["channel_models"]["0"]["selected"] == "fold_log_gradient" + + +def test_compute_supplies_independent_z_planes_for_automatic_selection(): + tile_client = _tile_client() + tile_client.tiles["IndexRange"]["IndexZ"] = 2 + tile_client.coordinatesToFrameIndex.side_effect = ( + lambda xy, z, time, channel: z * 2 + channel + ) + sink = MagicMock() + fake_large_image = types.SimpleNamespace(new=MagicMock(return_value=sink)) + reference = (_grid(), 0, [{"channel": 0, "quality_score": 0.1}]) + + def inspect_validation(*args, validation_source=None, **kwargs): + reports = list(validation_source()) + assert [label for label, _ in reports] == ["held-out Z 2"] + assert np.all(reports[0][1] == 102) + return _selection() + + with ( + patch("entrypoint.tiles.UPennContrastDataset", return_value=tile_client), + patch("entrypoint.correction.choose_reference_grid", return_value=reference), + patch("entrypoint.correction.select_model", side_effect=inspect_validation), + patch.dict(sys.modules, {"large_image": fake_large_image}), + ): + compute("dataset-id", "http://api", "token", _params()) + + tile_client.client.uploadFileToFolder.assert_called_once() + + +def test_compute_rejects_malformed_channel_selection(): + tile_client = _tile_client() + with ( + patch("entrypoint.tiles.UPennContrastDataset", return_value=tile_client), + patch("entrypoint.sendError") as send_error, + pytest.raises(ValueError, match="mapping"), + ): + compute("dataset-id", "http://api", "token", _params(channels=[0])) + + send_error.assert_called_once() + assert "channel selection" in send_error.call_args.args[0].lower() + tile_client.client.uploadFileToFolder.assert_not_called() + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("Algorithm", None), + ("Algorithm", "Removed algorithm"), + ("Reference channel mode", None), + ("BaSiC darkfield", "Maybe"), + ("Output type", None), + ], +) +def test_compute_rejects_missing_or_stale_select_values(field, value): + params = _params() + params["workerInterface"][field] = value + + with ( + patch("entrypoint.tiles.UPennContrastDataset") as dataset_client, + patch("entrypoint.sendError") as send_error, + pytest.raises(ValueError, match=field), + ): + compute("dataset-id", "http://api", "token", params) + + send_error.assert_called_once() + dataset_client.assert_not_called() + + +def test_compute_supports_single_frame_without_index_range(): + tile_client = _tile_client(frames=False) + tile_client.tiles.pop("IndexRange") + tile_client.tiles["channels"] = ["DAPI"] + sink = MagicMock() + fake_large_image = types.SimpleNamespace(new=MagicMock(return_value=sink)) + reference = (_grid(), 0, [{"channel": 0, "quality_score": 0.1}]) + params = _params(channels={"0": True}) + + with ( + patch("entrypoint.tiles.UPennContrastDataset", return_value=tile_client), + patch("entrypoint.correction.choose_reference_grid", return_value=reference), + patch("entrypoint.correction.select_model", return_value=_selection()), + patch.dict(sys.modules, {"large_image": fake_large_image}), + ): + compute("dataset-id", "http://api", "token", params) + + sink.addTile.assert_called_once() + tile_client.client.uploadFileToFolder.assert_called_once() + + +def test_compute_rejects_selected_channels_outside_dataset(): + tile_client = _tile_client() + params = _params(channels={"3": True}) + + with ( + patch("entrypoint.tiles.UPennContrastDataset", return_value=tile_client), + patch("entrypoint.correction.choose_reference_grid") as choose_grid, + patch("entrypoint.sendError") as send_error, + pytest.raises(ValueError, match="do not exist"), + ): + compute("dataset-id", "http://api", "token", params) + + choose_grid.assert_not_called() + send_error.assert_called_once() + + +def test_compute_rejects_reference_coordinate_outside_dataset(): + tile_client = _tile_client() + params = _params() + params["workerInterface"]["Reference Z"] = "2" + + with ( + patch("entrypoint.tiles.UPennContrastDataset", return_value=tile_client), + patch("entrypoint.correction.choose_reference_grid") as choose_grid, + patch("entrypoint.sendError") as send_error, + pytest.raises(ValueError, match="Reference Z"), + ): + compute("dataset-id", "http://api", "token", params) + + choose_grid.assert_not_called() + send_error.assert_called_once() + + +def test_automatic_reference_ignores_malformed_manual_channel_value(): + tile_client = _tile_client() + sink = MagicMock() + fake_large_image = types.SimpleNamespace(new=MagicMock(return_value=sink)) + reference = (_grid(), 0, [{"channel": 0, "quality_score": 0.1}]) + params = _params() + params["workerInterface"]["Reference channel"] = None + + with ( + patch("entrypoint.tiles.UPennContrastDataset", return_value=tile_client), + patch("entrypoint.correction.choose_reference_grid", return_value=reference), + patch("entrypoint.correction.select_model", return_value=_selection()), + patch.dict(sys.modules, {"large_image": fake_large_image}), + ): + compute("dataset-id", "http://api", "token", params) + + tile_client.client.uploadFileToFolder.assert_called_once() + + +def test_compute_only_applies_model_to_reference_xy_and_time(): + tile_client = _tile_client() + tile_client.tiles["IndexRange"]["IndexXY"] = 2 + tile_client.tiles["frames"] = [ + {"IndexXY": 0, "IndexZ": 0, "IndexT": 0, "IndexC": 0}, + {"IndexXY": 1, "IndexZ": 0, "IndexT": 0, "IndexC": 0}, + ] + tile_client.getRegion.side_effect = lambda dataset_id, frame: np.full( + (64, 64, 1), 100 + frame, dtype=np.uint16 + ) + sink = MagicMock() + fake_large_image = types.SimpleNamespace(new=MagicMock(return_value=sink)) + reference = (_grid(), 0, [{"channel": 0, "quality_score": 0.1}]) + params = _params(channels={"0": True}) + + with ( + patch("entrypoint.tiles.UPennContrastDataset", return_value=tile_client), + patch("entrypoint.correction.choose_reference_grid", return_value=reference), + patch("entrypoint.correction.select_model", return_value=_selection()), + patch.dict(sys.modules, {"large_image": fake_large_image}), + ): + compute("dataset-id", "http://api", "token", params) + + reference_output = sink.addTile.call_args_list[0].args[0] + other_xy_output = sink.addTile.call_args_list[1].args[0] + assert np.all(reference_output == 101) + assert np.all(other_xy_output == 101) + + +class _OverflowModel: + name = "fold_log_gradient" + diagnostics = {} + + def apply(self, image): + return np.full(np.asarray(image).squeeze().shape, 70000.0, dtype=np.float32) + + +def test_preserve_dtype_rejects_material_clipping(): + tile_client = _tile_client() + sink = MagicMock() + fake_large_image = types.SimpleNamespace(new=MagicMock(return_value=sink)) + reference = (_grid(), 0, [{"channel": 0, "quality_score": 0.1}]) + selection = _selection() + selection.model = _OverflowModel() + params = _params(channels={"0": True}) + params["workerInterface"]["Output type"] = "Preserve source dtype" + + with ( + patch("entrypoint.tiles.UPennContrastDataset", return_value=tile_client), + patch("entrypoint.correction.choose_reference_grid", return_value=reference), + patch("entrypoint.correction.select_model", return_value=selection), + patch("entrypoint.sendError") as send_error, + patch.dict(sys.modules, {"large_image": fake_large_image}), + pytest.raises(ValueError, match="clipped"), + ): + compute("dataset-id", "http://api", "token", params) + + send_error.assert_called() + sink.write.assert_not_called() diff --git a/workers/annotations/illumination_correction/tests/test_illumination.py b/workers/annotations/illumination_correction/tests/test_illumination.py new file mode 100644 index 0000000..90b4a7a --- /dev/null +++ b/workers/annotations/illumination_correction/tests/test_illumination.py @@ -0,0 +1,302 @@ +import sys +from pathlib import Path + +import numpy as np +import pytest + + +WORKER_DIR = Path(__file__).resolve().parents[1] +if str(WORKER_DIR) not in sys.path: + sys.path.insert(0, str(WORKER_DIR)) + +from illumination import ( # noqa: E402 + CandidateResult, + IdentityModel, + TileGrid, + a1_fold_amplitude, + choose_reference_grid, + fit_basic, + fit_grid, + fit_log_gradient, + fit_split_half_affine, + normalize_flat, + preservation_metrics, + rank_candidates, + select_model, +) + + +def _synthetic_mosaic(seed=3, pitch=32, tiles=6): + rng = np.random.default_rng(seed) + size = pitch * tiles + y, x = np.indices((size, size), dtype=np.float32) + uy = (y % pitch) / pitch + ux = (x % pitch) / pitch + + # A smooth, periodic illumination field with a pronounced seam falloff. + flat = ( + 0.72 + + 0.20 * np.sin(np.pi * uy) ** 2 + + 0.16 * np.sin(np.pi * ux) ** 2 + + 0.06 * np.sin(2 * np.pi * uy) * np.sin(2 * np.pi * ux) + ) + + # Independent texture in every physical tile prevents the estimator from + # learning one repeated biological pattern. + biology = 900.0 + rng.normal(0, 35, (size, size)).astype(np.float32) + for tile_y in range(tiles): + for tile_x in range(tiles): + y0, x0 = tile_y * pitch, tile_x * pitch + for _ in range(3): + cy = y0 + int(rng.integers(4, pitch - 4)) + cx = x0 + int(rng.integers(4, pitch - 4)) + biology[cy - 2 : cy + 3, cx - 2 : cx + 3] += rng.uniform(250, 700) + + raw = 180.0 + flat * biology + grid = TileGrid( + pitch_y=float(pitch), + pitch_x=float(pitch), + seam_y=0.0, + seam_x=0.0, + height=size, + width=size, + seams_y=tuple(float(v) for v in range(0, size + 1, pitch)), + seams_x=tuple(float(v) for v in range(0, size + 1, pitch)), + seam_residual_y=0.0, + seam_residual_x=0.0, + ) + return biology.astype(np.float32), raw.astype(np.float32), grid + + +def test_fit_grid_recovers_periodic_geometry(): + _, raw, _ = _synthetic_mosaic() + fitted = fit_grid(raw, pitch_min=24, pitch_max=40) + + assert abs(fitted.pitch_y - 32) < 2.0 + assert abs(fitted.pitch_x - 32) < 2.0 + assert fitted.is_valid + assert len(fitted.seams_y) >= 4 + assert len(fitted.seams_x) >= 4 + + +def test_auto_reference_uses_best_grid_in_dominant_channel_cluster(monkeypatch): + class TileClient: + tiles = {"frames": [{"IndexC": 0}, {"IndexC": 1}, {"IndexC": 2}]} + + @staticmethod + def coordinatesToFrameIndex(xy, z, time, channel): + return channel + + @staticmethod + def getRegion(dataset_id, frame): + return np.full((64, 64), frame, dtype=np.float32) + + def candidate_grid(image, pitch_min, pitch_max): + channel = int(image[0, 0]) + pitches = ((32.0, 32.0), (32.5, 31.8), (20.0, 20.0)) + residuals = (1.0, 0.2, 0.01) + prominences = (3.0, 8.0, 50.0) + py, px = pitches[channel] + return TileGrid( + pitch_y=py, + pitch_x=px, + seam_y=0.0, + seam_x=0.0, + height=64, + width=64, + seams_y=(0.0, py, 2 * py, 3 * py), + seams_x=(0.0, px, 2 * px, 3 * px), + seam_residual_y=residuals[channel], + seam_residual_x=residuals[channel], + prominence_y=prominences[channel], + prominence_x=prominences[channel], + ) + + monkeypatch.setattr("illumination.fit_grid", candidate_grid) + selected, channel, reports = choose_reference_grid( + TileClient(), + "dataset-id", + {"XY": 0, "Z": 0, "Time": 0}, + "Automatically choose best channel", + 0, + 10, + 40, + ) + + assert channel == 1 + assert selected.pitch_y == 32.5 + assert [report["cross_channel_agreement"] for report in reports] == [2, 2, 1] + + +def test_folded_log_gradient_reduces_position_locked_artifact(): + _, raw, grid = _synthetic_mosaic(seed=7) + before = a1_fold_amplitude(raw, grid)["A1_fold_amp_rel_pct"] + + model = fit_log_gradient(raw, grid, n=64, per_tile_gain=True) + corrected = model.apply(raw) + after = a1_fold_amplitude(corrected, grid)["A1_fold_amp_rel_pct"] + + assert corrected.dtype == np.float32 + assert np.isfinite(corrected).all() + assert after < before * 0.65 + + +def test_split_half_affine_reduces_position_locked_artifact(): + _, raw, grid = _synthetic_mosaic(seed=11) + before = a1_fold_amplitude(raw, grid)["A1_fold_amp_rel_pct"] + + model = fit_split_half_affine( + raw, + grid, + profile_size=64, + split_count=6, + seed=11, + ) + corrected = model.apply(raw) + after = a1_fold_amplitude(corrected, grid)["A1_fold_amp_rel_pct"] + + assert np.isfinite(corrected).all() + assert after < before * 0.75 + + +def _candidate(name, artifact_index, p1=1.0, violations=(), complexity=0): + return CandidateResult( + name=name, + model=object(), + metrics={"P1_spot_uniformity": p1}, + artifact_index=artifact_index, + violations=list(violations), + physics_violations=[], + complexity=complexity, + ) + + +def test_rank_candidates_rejects_guardrail_failures_and_penalizes_spot_bias(): + unsafe = _candidate( + "basic_darkfield_on", + artifact_index=0.18, + p1=1.0, + violations=("P2_spearman=0.96 < 0.98",), + ) + biased = _candidate( + "basic_darkfield_off", artifact_index=0.39, p1=1.55, complexity=2 + ) + stable = _candidate("fold_log_gradient", artifact_index=0.41, p1=1.02, complexity=1) + + selected, ranked = rank_candidates( + [unsafe, biased, stable], use_spot_uniformity=True + ) + + assert selected.name == "fold_log_gradient" + assert unsafe not in ranked + + +def test_rank_candidates_prefers_simpler_method_inside_tie_margin(): + complex_candidate = _candidate("basic_darkfield_on", 0.200, complexity=3) + simple_candidate = _candidate("split_half_affine", 0.207, complexity=0) + + selected, _ = rank_candidates( + [complex_candidate, simple_candidate], tie_fraction=0.05 + ) + + assert selected.name == "split_half_affine" + + +def test_automatic_selection_can_keep_identity_when_corrections_are_worse(): + identity = _candidate("identity", 1.0, complexity=-1) + worse = _candidate("fold_log_gradient", 1.25, complexity=1) + + selected, _ = rank_candidates([identity, worse]) + + assert selected.name == "identity" + + +def test_rank_candidates_rejects_undefined_artifact_scores(): + undefined = _candidate("split_half_affine", float("nan"), complexity=0) + + with pytest.raises(ValueError, match="finite artifact score"): + rank_candidates([undefined]) + + +def test_spot_penalty_treats_zero_as_extreme_bias(): + zero_outer = _candidate("zero_outer", 0.20, p1=0.0, complexity=0) + balanced = _candidate("balanced", 0.25, p1=1.0, complexity=1) + + selected, _ = rank_candidates( + [zero_outer, balanced], use_spot_uniformity=True + ) + + assert selected.name == "balanced" + + +def test_normalize_flat_rejects_nonfinite_and_nonpositive_fields(): + with pytest.raises(ValueError, match="finite and strictly positive"): + normalize_flat(np.full((8, 8), np.nan, dtype=np.float32)) + with pytest.raises(ValueError, match="finite and strictly positive"): + normalize_flat(np.zeros((8, 8), dtype=np.float32)) + + +def test_range_guardrail_ignores_preexisting_zeros_but_rejects_new_ones(): + raw = np.full((64, 64), 10.0, dtype=np.float32) + raw[:8] = 0.0 + unchanged = preservation_metrics(raw, raw.copy()) + + damaged = raw.copy() + damaged[8:16] = 0.0 + damaged_metrics = preservation_metrics(raw, damaged) + + assert unchanged["P5_frac_new_nonpositive"] == 0.0 + assert not any( + "P5_frac_new_nonpositive" in item + for item in unchanged["guardrail_violations"] + ) + assert any( + "P5_frac_new_nonpositive" in item + for item in damaged_metrics["guardrail_violations"] + ) + + +def test_range_guardrail_rejects_nonfinite_output(): + raw = np.full((64, 64), 10.0, dtype=np.float32) + corrected = raw.copy() + corrected[0, 0] = np.nan + + metrics = preservation_metrics(raw, corrected) + + assert metrics["P5_frac_nonfinite"] > 0 + assert any( + "P5_frac_nonfinite" in item for item in metrics["guardrail_violations"] + ) + + +def test_automatic_selection_without_held_out_plane_returns_identity(): + _, raw, grid = _synthetic_mosaic(seed=17) + + selected = select_model( + raw, + grid, + "Automatic (recommended)", + "Automatic", + per_tile_gain=False, + ) + + assert isinstance(selected.model, IdentityModel) + assert selected.name == "identity" + assert selected.metrics["selection_basis"] == "identity_without_holdout" + + +@pytest.mark.parametrize("darkfield", [False, True]) +def test_basic_fit_smoke_uses_production_dependency_path(darkfield): + _, raw, grid = _synthetic_mosaic(seed=19, pitch=32, tiles=4) + + model = fit_basic( + raw, + grid, + darkfield=darkfield, + per_tile_gain=False, + tile_n=32, + ) + corrected = model.apply(raw) + + assert corrected.shape == raw.shape + assert np.isfinite(corrected).all()