diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..406eb48 --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +__pycache__/ +*.py[cod] +*.egg-info/ +.pytest_cache/ diff --git a/FLUOR_CORRECTION_WORKERS_SPEC.md b/FLUOR_CORRECTION_WORKERS_SPEC.md new file mode 100644 index 0000000..e639edc --- /dev/null +++ b/FLUOR_CORRECTION_WORKERS_SPEC.md @@ -0,0 +1,490 @@ +# Spec: Fluorescence Illumination-Correction and Image-Restoration Workers + +Status: design spec for two new NimbusImage image-processing workers. +Author: initial design pass; implementation delegated per worker. + +## Goal + +Add two new **image-processing** workers (category `Image Processing`, live under +`workers/annotations/`, following the `histogram_matching` / `deconwolf` / +`rolling_ball` / `registration` convention of reading all frames, processing, +and uploading a new TIFF to Girder): + +1. **`illumination_correction`** — corrects spatial illumination / shading / + vignetting / background bias / striping. Single worker, user picks the + algorithm from a `select` dropdown. Implements **all** algorithms from the + ChatGPT illumination shortlist. +2. **`image_restoration`** — denoises / deblurs / restores. Single worker, user + picks the algorithm. Implements the best **4** methods, deliberately biased + toward **reference-free / self-supervised / pretrained** methods (we usually + have no paired clean ground truth), and **includes FluoResFM** per explicit + request. + +Both workers share the same skeleton as `histogram_matching`/`deconwolf`: +`interface()` + `compute()`, `argparse` main, iterate `tileClient.tiles['frames']`, +build a `large_image` sink, copy metadata (`channelNames`, `mm_x`, `mm_y`, +`magnification`), write `/tmp/.tiff`, `gc.uploadFileToFolder`, then +`gc.addMetadataToItem` with a provenance dict (`tool`, `method`, key params). + +--- + +## Design principles (apply to both workers) + +1. **One worker, `Method` `select` dropdown.** The interface exposes ALL + parameters for ALL methods; each parameter's tooltip notes which method(s) + it applies to. `compute()` dispatches on `Method` to a per-algorithm + function. (There is no conditional UI in this framework, so unused params + for the chosen method are simply ignored — document this.) + +2. **Per-channel processing.** A `channelCheckboxes` field ("Channels to + correct" / "Channels to restore") selects channels. Unselected channels + pass through **unchanged** (exactly like `histogram_matching` and + `deconwolf`). Each selected channel is corrected/restored independently. + +3. **Retrospective estimators operate on a per-channel image collection.** + For BaSiC / CIDRE / CellProfiler (illumination) the "collection" is the set + of all frames of that channel across XY / Z / Time. Gather that stack + (`N, Y, X`), fit the model on it, then transform each frame. Warn (via + `sendWarning`) if the collection is very small (e.g. < 3 images) since + retrospective methods are unreliable then. + +4. **Testability is mandatory and shapes the code structure.** Each algorithm + MUST be a standalone function (e.g. `correct_basic(stack, opts) -> stack`, + `restore_n2v(stack, opts) -> stack`). Heavy third-party imports (`basicpy`, + `careamics`, `cellpose`, torch, the FluoResFM/ZS-DeconvNet modules) MUST be + imported **lazily inside those functions**, never at module top level. This + keeps the `interface` path fast (see `todo/worker-startup-latency.md`) AND + lets tests run natively without the heavy deps installed. `compute()`'s + dispatch, frame iteration, sink assembly, metadata copy, channel filtering, + and error paths are what the pytest suite exercises — with the per-algorithm + functions **mocked** (patch `entrypoint.correct_basic`, etc.), mirroring how + `histogram_matching/tests/test_histogram_matching.py` patches + `entrypoint.match_histograms`. + +5. **Graceful degradation.** If a method's optional dependency or model weights + are missing at runtime, call `sendError(...)` with a clear, actionable + message (what's missing, how to enable it) and return — do not crash with a + raw traceback. GPU methods must fall back to CPU with a `sendWarning` when + CUDA is unavailable (pattern: `deconwolf`'s OpenCL→CPU fallback, and + `torch.cuda.is_available()`). + +6. **Metadata provenance.** Always `addMetadataToItem` with at least + `{'tool': , 'method': , 'channels': , ...method params}`. + +7. **dtype handling.** Preserve input dtype where reasonable. Many of these + algorithms work in float internally; when writing back, clip and cast to the + source dtype (read from `tileClient.tiles.get('dtype')` or the frame's own + dtype) to avoid silent 16-bit→float bloat in the output TIFF. Document the + cast. Guard against NaN/inf from the algorithms (`np.nan_to_num`). + +8. **Progress.** Call `sendProgress(fraction, title, info)` regularly (per frame + or per stage) so long ML jobs show life. + +--- + +## Worker 1: `illumination_correction` + +Path: `workers/annotations/illumination_correction/` +Interface name: `Illumination Correction` • category: `Image Processing` +Default tool name: `Illumination Correction` +Base image: **CPU** — `nimbusimage/image-processing-base:latest` +(none of these methods require a GPU; BaSiCPy runs fine on CPU via torch). +No `Dockerfile_M1` needed unless a dep is x86-only — use the same Dockerfile for both (it inherits the multi-arch base). If a `MAC_DEVELOPMENT_MODE` variant is trivial, skip it; the base is already arch-aware. + +### Methods (`Method` select) — implement ALL of these + +| Value | Label | Kind | Dependency | Needs ref frames? | +|-------|-------|------|-----------|-------------------| +| `basic` | BaSiC (shading + darkfield, recommended) | retrospective | `basicpy` (pip) | no | +| `cidre` | CIDRE-style retrospective | retrospective | in-house numpy/scipy | no | +| `cellprofiler` | CellProfiler-style illumination function | retrospective or per-image | in-house numpy/scipy/skimage | no | +| `flatfield` | Flat/Dark-field (reference-based) | reference | in-house numpy | **yes** | +| `destripe` | Stripe / tiling-seam correction (classical, fast, no weights) | classical destriping | `pystripe` (pip) | no | +| `sscor` | SSCOR deep-learning stripe self-correction | DL (subprocess to vendored repo) | vendored `lxxcontinue/SSCOR` + runtime checkpoint | no | + +Default `Method` = `basic`. + +Plus an **EVEN-style QC** toggle (see below) — not a correction method, an +optional post-correction quality report. + +#### Method details + +**`basic` — BaSiC / BaSiCPy** (flagship retrospective, Peng et al. 2017; BaSiCPy 2.x, PyTorch backend). +- `pip install basicpy`. API: + ```python + from basicpy import BaSiC + basic = BaSiC(get_darkfield=, + smoothness_flatfield=, # default 1.0 + smoothness_darkfield=, # default 1.0 + fitting_mode='approximate') # or 'ladmap' + basic.fit(stack) # stack: (N, Y, X) float + corrected = basic.transform(stack) + # timelapse baseline drift: basic.baseline holds per-image baseline after fit + ``` +- Interface params used: `Estimate darkfield` (checkbox, default True), + `Flatfield smoothness` (number, default 1.0), `Darkfield smoothness` + (number, default 1.0), `Correct timelapse baseline drift` (checkbox, default + False — when True, sort the collection by Time and let BaSiC estimate/subtract + the temporal baseline; store fitted flatfield/darkfield in item metadata). +- Fit per channel on the full (N,Y,X) collection; `transform` all frames of + that channel. Save the estimated flatfield (and darkfield) mean values in + metadata for provenance. + +**`cidre` — CIDRE-style retrospective** (Smith et al. 2015, Nature Methods). +There is no maintained pip CIDRE. Implement a faithful **CIDRE-style** in-house +estimator (document clearly it is a lightweight reimplementation of CIDRE's +retrospective gain/offset model, not the original MATLAB code): +- Stack all frames of the channel `(N, Y, X)`. +- Estimate an **offset (dark) surface** `z(x,y)` ≈ robust per-pixel low quantile + (e.g. 2nd percentile) across the collection, then heavily Gaussian-smoothed. +- Estimate a **gain (flat) surface** `v(x,y)` ≈ per-pixel robust central tendency + (e.g. per-pixel median of `frame - z`), normalized so `mean(v) == 1`, then + Gaussian-smoothed with a large sigma (spatial regularization stand-in for + CIDRE's energy minimization). +- Correct: `corrected = (frame - z) / v`. +- Interface params: reuse a `Smoothing sigma (retrospective)` number + (default e.g. 0.25 × image width, or a pixel value like 50) controlling the + Gaussian regularization; `Dark quantile` number (default 0.02). +- Document the simplification honestly in the docs and a code comment. + +**`cellprofiler` — CellProfiler-style illumination function**. +Reimplement CellProfiler's `CorrectIlluminationCalculate` core (no CellProfiler +dependency): +- Two modes via a `CellProfiler mode` select: `regular` (across-batch: + average all frames of the channel, then smooth → illumination function) and + `background` (per-image: heavily smooth each frame to estimate its own + background). Default `regular`. +- Smoothing method: large-kernel Gaussian or median (expose + `Smoothing sigma` reused from above). Rescale the illumination function so its + mean (or its min, CellProfiler-style) is 1.0, then `corrected = frame / illum`. +- For `regular`, the illumination function is computed once per channel from the + collection average and applied to every frame; for `background`, computed and + subtracted/divided per frame. + +**`flatfield` — Flat/Dark-field reference-based** (classical, gold standard when +calibration frames exist): `corrected = (raw - dark) / (flat - dark)`, rescaled +to preserve mean intensity. +- Reference frames are supplied by **XY coordinate** (1-indexed text inputs, + like `histogram_matching`'s reference coords): `Flat-field XY coordinate` and + `Dark-field XY coordinate` (both text; empty = disabled). Per selected channel, + the flat/dark reference is the frame at (that XY, Z=0, T=0, that channel). If a + dark ref is empty, treat dark as 0 (or a scalar `Dark-field constant` number, + default 0). If flat ref is empty for this method, `sendError` (flatfield is + meaningless without a flat reference) and return. +- Normalize `flat` so its mean is 1 after dark subtraction to preserve overall + intensity scale. Clip negatives to 0. + +**`destripe` — stripe / tiling-seam correction** (stand-in for SSCOR, which is a +non-packaged DL/GAN method): use **`pystripe`** (`pip install pystripe`, +wavelet-FFT destriping, originally for light-sheet): + ```python + from pystripe import core as pystripe_core + filtered = pystripe_core.filter_streaks(frame, sigma=[sigma1, sigma2], + level=level, wavelet=wavelet) + ``` +- Interface params: `Destripe sigma` (number, default 128 — controls + band-pass), `Destripe wavelet` (select: `db3`,`db5`,`haar`, default `db3`), + `Destripe level` (number, default 0 = auto). Applied per frame. +- Docs must state this is a **classical destriping** alternative and that the + DL method SSCOR is not vendored (no packaged weights). + +#### EVEN-style QC toggle (not a correction method) + +`Report correction quality (QC)` checkbox (default False). When True, after +correction, compute simple flat-field-quality metrics per selected channel on +the corrected collection and emit them via `sendWarning`/print and store in item +metadata: e.g. coefficient of variation of a smoothed mean image, residual +vignetting (corner-vs-center ratio), inter-frame intensity CV. Document that +full **EVEN** (Nat. Commun. 2026; an ML/LDA evaluation-and-optimization +framework) is not vendored — this is a lightweight quantitative QC stand-in that +lets a user compare methods. + +### Interface (`illumination_correction`) + +```python +interface = { + 'Method': {'type': 'select', + 'items': ['basic','cidre','cellprofiler','flatfield','destripe'], + 'default': 'basic', 'displayOrder': 0, + 'tooltip': 'Illumination-correction algorithm. basic (BaSiC) recommended when no calibration frames.'}, + 'Channels to correct': {'type': 'channelCheckboxes', 'displayOrder': 1, + 'tooltip': 'Process selected channels; others pass through unchanged.'}, + # --- BaSiC --- + 'Estimate darkfield': {'type': 'checkbox', 'default': True, 'displayOrder': 2, + 'tooltip': 'BaSiC: also estimate a darkfield (offset) term.'}, + 'Flatfield smoothness': {'type': 'number', 'min': 0, 'max': 100, 'default': 1.0, 'displayOrder': 3, + 'tooltip': 'BaSiC: smoothness regularization of the flatfield.'}, + 'Darkfield smoothness': {'type': 'number', 'min': 0, 'max': 100, 'default': 1.0, 'displayOrder': 4, + 'tooltip': 'BaSiC: smoothness regularization of the darkfield.'}, + 'Correct timelapse baseline drift': {'type': 'checkbox', 'default': False, 'displayOrder': 5, + 'tooltip': 'BaSiC: estimate and remove temporal background/bleaching drift across Time.'}, + # --- CIDRE / CellProfiler shared --- + 'Smoothing sigma': {'type': 'number', 'min': 1, 'max': 2000, 'default': 50, 'displayOrder': 6, + 'tooltip': 'CIDRE/CellProfiler: Gaussian smoothing sigma (px) for the illumination surface.'}, + 'Dark quantile': {'type': 'number', 'min': 0, 'max': 0.5, 'default': 0.02, 'displayOrder': 7, + 'tooltip': 'CIDRE: percentile for the dark/offset surface estimate.'}, + 'CellProfiler mode': {'type': 'select', 'items': ['regular','background'], 'default': 'regular', 'displayOrder': 8, + 'tooltip': 'CellProfiler: regular=across-batch average; background=per-image background.'}, + # --- flatfield reference --- + 'Flat-field XY coordinate': {'type': 'text', 'displayOrder': 9, + 'vueAttrs': {...placeholder 'ex. 1'...}, + 'tooltip': 'flatfield: 1-indexed XY position of the flat-field reference image.'}, + 'Dark-field XY coordinate': {'type': 'text', 'displayOrder': 10, + 'tooltip': 'flatfield: 1-indexed XY position of the dark-field reference (empty = use constant).'}, + 'Dark-field constant': {'type': 'number', 'min': 0, 'max': 65535, 'default': 0, 'displayOrder': 11, + 'tooltip': 'flatfield: constant dark offset if no dark reference frame.'}, + # --- destripe --- + 'Destripe sigma': {'type': 'number', 'min': 1, 'max': 2000, 'default': 128, 'displayOrder': 12, + 'tooltip': 'destripe: band-pass sigma for wavelet-FFT stripe removal.'}, + 'Destripe wavelet': {'type': 'select', 'items': ['db3','db5','haar','sym4'], 'default': 'db3', 'displayOrder': 13}, + 'Destripe level': {'type': 'number', 'min': 0, 'max': 12, 'default': 0, 'displayOrder': 14, + 'tooltip': 'destripe: DWT decomposition level (0 = auto).'}, + # --- QC --- + 'Report correction quality (QC)': {'type': 'checkbox', 'default': False, 'displayOrder': 15}, +} +``` +(Confirm `select` uses `items` list — check how existing workers e.g. cellposesam +or sample_interface declare `select` options and match that exact schema.) + +### Dependencies / Dockerfile (`illumination_correction`) + +`environment.yml` (conda-forge based, mirror `histogram_matching/environment.yml`): +```yaml +name: worker +channels: [conda-forge, defaults] +dependencies: + - python=3.11 + - pip + - numpy>=2.0 + - scipy + - scikit-image + - shapely>=2.0.6 + - libtiff + - openslide + - libvips + - tifffile + - pip: + - basicpy + - pystripe +``` +Dockerfile: `FROM nimbusimage/image-processing-base:latest`, `COPY entrypoint.py`, +install `large-image-source-*` wheels (see deconwolf), pip-install +`basicpy`/`pystripe` (if not via env), LABELs (`interfaceName="Illumination Correction"`, +`interfaceCategory="Image Processing"`, `isAnnotationWorker=""`, +`description="Corrects illumination/shading/striping (BaSiC, CIDRE, CellProfiler, flat-field, destripe)"`), +and the `run_worker.sh` entrypoint. **Note:** verify `basicpy` installs cleanly +on the base's Python 3.11; if `basicpy` pins an old jax, prefer BaSiCPy 2.x +(torch). If a conda install of torch is needed, add it to env. + +--- + +## Worker 2: `image_restoration` + +Path: `workers/annotations/image_restoration/` +Interface name: `Image Restoration` • category: `Image Processing` +Default tool name: `Image Restoration` +Base image: **GPU** — `nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04` with the +Miniforge+conda pattern from `deconwolf`/`cellposesam` Dockerfiles, **plus a CPU +fallback at runtime** (`torch.cuda.is_available()` → CPU with `sendWarning`). +Provide a `Dockerfile_M1` (CPU-only, `FROM nimbusimage/image-processing-base`) +for Mac dev, like `deconwolf`. + +### Methods (`Method` select) — implement these 4 (all reference-free / pretrained) + +| Value | Label | Kind | Dependency | Paired data? | +|-------|-------|------|-----------|--------------| +| `n2v` | Noise2Void / N2V2 (self-supervised denoise) | trains on your data, no clean target | `careamics` (pip, PyTorch) | no | +| `cellpose3` | Cellpose3 restoration (denoise/deblur/upsample) | pretrained | `cellpose>=3` (pip) | no | +| `zs_deconvnet` | ZS-DeconvNet (zero-shot denoise + deconv) | zero-shot, trains on the single input | vendored repo | no | +| `fluoresfm` | FluoResFM (foundation model, pretrained) | pretrained, text-prompted | vendored repo + weights | no | + +Default `Method` = `n2v`. + +**Rationale for the selection (put in docs):** we usually lack paired clean +ground truth, so supervised CARE/CSBDeep and 3D-RCAN are intentionally excluded. +All four chosen methods work without paired references: N2V/N2V2 is +self-supervised (trains on the noisy data itself), Cellpose3 restoration ships +pretrained models, ZS-DeconvNet is zero-shot (trains on the single input volume), +and FluoResFM is a pretrained foundation model (included per user request; treat +as experimental and validate before quantitative use). + +#### Method details + +**`n2v` — Noise2Void / N2V2 via CAREamics** (`pip install careamics`). +- API (verify against careamics docs at build time): + ```python + from careamics import CAREamist + from careamics.config import create_n2v_configuration + cfg = create_n2v_configuration(experiment_name='n2v', data_type='array', + axes='YX', patch_size=[64,64], batch_size=16, + num_epochs=, use_n2v2=) + engine = CAREamist(cfg) + engine.train(train_source=stack) # self-supervised on the (N,Y,X) stack + restored = engine.predict(source=stack) + ``` +- Trains **per channel** on that channel's collection (self-supervised). +- Params: `Epochs` (number, default 20), `Use N2V2` (checkbox, default True — + N2V2 reduces checkerboard artifacts), `Patch size` (number, default 64). +- GPU strongly recommended; CPU fallback allowed but warn it is slow. + +**`cellpose3` — Cellpose3 restoration** (`pip install cellpose>=3`). +- API: + ```python + from cellpose import denoise + model = denoise.DenoiseModel(model_type=, gpu=) + restored = model.eval(frame, channels=[0,0]) + ``` +- `Cellpose3 model` select: `denoise_cyto3`, `deblur_cyto3`, `upsample_cyto3`, + `denoise_nuclei`, `deblur_nuclei`, `oneclick_cyto3` (verify exact model names + against installed cellpose version; expose a reasonable subset). Default + `denoise_cyto3`. +- Pretrained; applied per frame. Note in docs: best used as **segmentation + preprocessing**, not for quantitative intensity restoration. Cellpose weights + download on first use — do it at build time via a small `download_models.py` + (pattern: cellposesam) so runtime has no network dependency. + +**`zs_deconvnet` — ZS-DeconvNet** (zero-shot; Nat. Commun. 2024; +project https://tristazeng.github.io/ZS-DeconvNet-page/, code on GitHub). +- Vendor by `git clone` in the Dockerfile (pin a commit). ZS-DeconvNet trains a + small network **on the single input image/stack** (no external training data) + using its physics-consistency loss, then infers. Wrap its training+inference + entry so `restore_zs_deconvnet(stack, opts)` runs the zero-shot pipeline per + frame (2D) or per Z-stack (3D). +- Params: `ZS iterations` (number, default per repo), `ZS upsampling` + (checkbox — SR mode vs denoise-only), and a PSF spec (reuse NA / wavelength / + pixel-size numbers like deconwolf, OR let it estimate). Keep params minimal; + default to denoise+deconv 2D. +- **Highest integration risk.** If the repo's API cannot be driven cleanly, + implement its published 2D zero-shot training loop directly (dual-stage + denoise→deconv with the Richardson-Lucy-consistency loss) in a + `zs_deconvnet.py` module. If truly infeasible in scope, the method must still + be selectable and `sendError` a clear "ZS-DeconvNet unavailable in this build" + message rather than crash — but real integration is the goal. Document + clearly what was implemented. + +**`fluoresfm` — FluoResFM** (foundation model; Nat. Commun. 2026; napari plugin +`napari-fluoresfm` by Qiqi Lu; depends on PyTorch + triton; pretrained `.pt` +checkpoint; text-prompt conditioned). +- Vendor the inference code (`git clone` the FluoResFM repo — find the canonical + repo, likely under the same author as UNiFMIR `github.com/cxm12`; the napari + plugin `napari-fluoresfm` is the reference). Weights: download the pretrained + `.pt` at build time via `download_models.py` if a stable URL (Zenodo/HF) is + known; otherwise load from a path configurable by env var + `FLUORESFM_WEIGHTS` and, if absent at runtime, `sendError` with instructions + (where to obtain weights, which env var / mount to set). Do NOT fail the build + if weights can't be fetched — make weights runtime-optional with a clear error. +- Params: `FluoResFM task` (select: `denoise`,`deconvolution`,`super-resolution`, + default `denoise`), `FluoResFM text prompt` (text, optional — the model is + text-conditioned; provide a sensible default prompt describing the structure, + e.g. "fluorescence microscopy, "), applied per frame. +- Mark **experimental** in docs; warn that restored intensities may be + hallucinated and should be validated before quantitative use (measure on + illumination-corrected raw, segment on restored). + +### Interface (`image_restoration`) + +Mirror the illumination worker's structure: `Method` select (default `n2v`), +`Channels to restore` channelCheckboxes, then the per-method params above, each +tooltip-tagged with its method. Include a `Use GPU` checkbox (default True) with +runtime `torch.cuda.is_available()` fallback + `sendWarning`. + +### Dependencies / Dockerfile (`image_restoration`) + +GPU Dockerfile modeled on `cellposesam`/`deconwolf`: +- `FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04`, Miniforge, `conda env + create` from `environment.yml`, install NimbusImage annotation_client (clone + `https://github.com/arjunrajlaboratory/NimbusImage/`, pip install its + annotation_client), install `annotation_utilities`, install large_image + wheels. +- pip: `torch` (CUDA build via the cuda base), `careamics`, `cellpose>=3`, + `triton` (for fluoresfm), `tifffile`. Clone ZS-DeconvNet and FluoResFM repos + (pinned commits). Run `download_models.py` for cellpose + (best-effort) + fluoresfm weights. +- LABELs: `isUPennContrastWorker=True`, `isAnnotationWorker=True` (image-processing + workers use the annotation-worker slot here — match `deconwolf`'s labels which + use `isAnnotationWorker=""`), `interfaceName="Image Restoration"`, + `interfaceCategory="Image Processing"`, + `description="Restores images (Noise2Void/N2V2, Cellpose3, ZS-DeconvNet, FluoResFM)"`, + `defaultToolName="Image Restoration"`. Include the + `com.nvidia.volumes.needed`/`NVIDIA_*` env from the GPU workers. +- `run_worker.sh` entrypoint. +- `Dockerfile_M1`: CPU-only `FROM nimbusimage/image-processing-base:latest`, + same pip installs minus CUDA torch (use CPU torch), no fluoresfm/zs weights + required to build. + +--- + +## Tests (both workers) + +Follow `histogram_matching/tests/` exactly: +- `tests/__init__.py`, `tests/Dockerfile_Test` (same template: `FROM + annotations/:latest AS test`, `pip install pytest pytest-mock`, copy + tests, pytest entrypoint), `tests/test_.py`. +- Tests run **natively without heavy deps** by mocking: + - `annotation_client.tiles.UPennContrastDataset` (fixture with `.tiles`, + `.getRegion`, `.coordinatesToFrameIndex`, `.client`), + - `large_image.new` (mock sink), + - `annotation_client.workers.UPennContrastWorkerPreviewClient`, + - each per-algorithm function (`entrypoint.correct_basic`, + `entrypoint.correct_cidre`, ... / `entrypoint.restore_n2v`, ...) — patched so + heavy libs are never imported. +- Required test cases per worker: + 1. `test_interface` — interface set once; asserts `Method` select exists with + expected items, channel field present, key params present. + 2. Dispatch: for each `Method` value, `compute` calls the right per-algorithm + function (assert the mock for the selected method is called and others are + not). + 3. Channel filtering: only selected channels processed; unselected frames + copied unchanged to the sink. + 4. Output plumbing: `sink.write('/tmp/.tiff')`, `uploadFileToFolder`, + `addMetadataToItem` with `tool` + `method`. + 5. Metadata preservation: `channelNames`, `mm_x`, `mm_y`, `magnification`. + 6. Error paths: no channels selected → `sendError`; single frame / no `frames` + → handled; `flatfield` with no flat reference → `sendError`. + 7. Progress reporting emitted. +- Register tests in the registry's "Tests" column (Yes). + +## Docs (both workers) + +`ILLUMINATION_CORRECTION.md` and `IMAGE_RESTORATION.md` in each worker dir, +following the template in CLAUDE.md / `DECONWOLF.md`: title + 1–2 sentence +description, How It Works, **Interface Parameters table** (every param, type, +default, which method it applies to), Methods table (algorithm, kind, dependency, +reference-free?), Implementation Details (coordinate/dtype handling, retrospective +collection semantics, CIDRE/CellProfiler simplifications, SSCOR→pystripe and +EVEN→QC substitutions honestly noted), Notes (GPU requirements, FluoResFM/ZS +experimental caveats, "segment on restored, measure on illumination-corrected +raw" guidance). + +## Registry + +Add both workers to `REGISTRY.md` under "Annotation Workers" (that's where the +other Image Processing workers like Deconwolf, Histogram Matching, Rolling Ball +live). Columns: name, description, GPU (Image Restoration = Yes, Illumination +Correction = blank), Tests = Yes, Docs link. Do NOT run +`generate_worker_docs.py --force` (it clobbers hand-written docs); edit +`REGISTRY.md` by hand. + +## Out of scope / honest substitutions (state in docs) + +- **SSCOR** (DL stripe correction): now implemented as the `sscor` method, + vendored from `github.com/lxxcontinue/SSCOR` (pinned commit) and driven via + subprocess to its CLI scripts. Two modes via `SSCOR mode`: + - `pretrained`: runs SSCOR's `restore.py` inference stage with a user-supplied + trained generator checkpoint (env `SSCOR_WEIGHTS`, staged internally as the + `latest_net_G_A.pth` the CycleGAN loader expects). + - `self-train`: the faithful upstream per-image method, needs no checkpoint — + for each frame it samples striped/stripe-free patches from the image itself + (`sample_stripe.py`/`sample_stripe_2.py`, controlled by stripe + direction/counts), trains a fresh CycleGAN (`train.py`), then restores. + Trains per frame → slow, GPU strongly recommended. + Operates in 8-bit (lossy round-trip for >8-bit sources). The classical + `pystripe` `destripe` method is retained as the fast, no-weights alternative. +- **EVEN** (ML evaluation/optimization framework): not vendored → optional + lightweight QC metrics report instead. +- **CIDRE**: no maintained pip package → faithful in-house gain/offset + reimplementation (documented as such). +- **Supervised restoration** (CARE/CSBDeep, 3D-RCAN) and time-lapse-specific + denoisers (DeepCAD-RT/FAST/TeD): excluded from the restoration worker because + they need paired data or are niche; the 4 chosen methods cover the + reference-free case. Note them as possible future additions. diff --git a/REGISTRY.md b/REGISTRY.md index 57de4e9..e525f9c 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 | 26 | +| Annotation Workers | 28 | | Property Workers -- Blobs | 10 | | Property Workers -- Points | 8 | | Property Workers -- Lines | 3 | | Property Workers -- Connections | 2 | | Test / Sample Workers | 7 | -| **Total** | **56** | +| **Total** | **58** | ## Annotation Workers @@ -36,6 +36,8 @@ 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 | Corrects illumination/shading/striping (BaSiC, CIDRE, CellProfiler, flat-field, destripe, SSCOR) | | Yes | [docs](workers/annotations/illumination_correction/ILLUMINATION_CORRECTION.md) | +| Image Restoration | Restores images (Noise2Void/N2V2, Cellpose3, ZS-DeconvNet, FluoResFM) | Yes | Yes | [docs](workers/annotations/image_restoration/IMAGE_RESTORATION.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..1cf8de4 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -224,6 +224,24 @@ 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 + + image_restoration: + build: + context: . + dockerfile: ./workers/annotations/image_restoration/${DOCKERFILE:-Dockerfile} + image: annotations/image_restoration:latest + profiles: ["worker", "image-processing"] + depends_on: + - image-processing-base + # AI workers ai_analysis: build: @@ -473,3 +491,21 @@ services: depends_on: - deconwolf 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"] + + image_restoration_test: + build: + context: . + dockerfile: ./workers/annotations/image_restoration/tests/Dockerfile_Test + image: annotations/image_restoration:test + depends_on: + - image_restoration + profiles: ["test"] diff --git a/workers/annotations/illumination_correction/Dockerfile b/workers/annotations/illumination_correction/Dockerfile new file mode 100644 index 0000000..f8357f3 --- /dev/null +++ b/workers/annotations/illumination_correction/Dockerfile @@ -0,0 +1,68 @@ +FROM nimbusimage/image-processing-base:latest + +# The base image's "worker" conda env already has numpy>=2.0 / scipy / +# scikit-image and large_image; activate it to install this worker's extra +# dependencies. +SHELL ["conda", "run", "-n", "worker", "/bin/bash", "-c"] + +# large_image wheels (mirrors deconwolf's pattern; already present in the base +# image, pinned again here for explicitness/self-containment). +RUN pip install large-image-source-zarr large-image-source-tiff large-image-converter large-image --find-links https://girder.github.io/large_image_wheels + +# CPU-only torch/torchvision. The base is a plain (non-CUDA) image, so pull the +# CPU wheels explicitly -- the default PyPI torch wheel bundles a multi-GB CUDA +# runtime that would never be used here. Installing torch first also lets us +# satisfy basicpy's/SSCOR's torch dependency with the CPU build. +RUN pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu + +# BaSiCPy (BaSiC shading/darkfield correction) for the `basic` method. +# Installed with --no-deps, then its runtime deps by hand: BaSiCPy 2.0.0 +# hard-pins scipy<1.13, but the base env is numpy>=2.0 (which needs scipy>=1.13, +# already installed and linked by the conda scikit-image/large_image binaries). +# Letting pip honor basicpy's pin would downgrade scipy and break those +# numpy-2-ABI C extensions at runtime, so we keep the base scipy and install +# only basicpy's *other* runtime deps (torch is the CPU build installed above). +RUN pip install --no-deps basicpy && \ + pip install "hyperactive>=4.4.0,<5" torch-dct pooch "pydantic>=2.0.0" dask tqdm + +# pystripe (wavelet-FFT destriping) for the `destripe` method. +RUN pip install pystripe + +# --- SSCOR (deep-learning stripe self-correction) --- +# Vendored from https://github.com/lxxcontinue/SSCOR, pinned to commit +# 985479cd79bcf1359e3d9ba44bacd5f372eb2e60. It's a pytorch-CycleGAN-and-pix2pix +# -style codebase with no clean importable inference API, so `correct_sscor` in +# entrypoint.py shells out to its CLI scripts (restore.py / train.py / sampling; +# see SSCOR_REPO_PATH). +RUN git clone https://github.com/lxxcontinue/SSCOR.git /sscor && \ + cd /sscor && git checkout 985479cd79bcf1359e3d9ba44bacd5f372eb2e60 + +# Extra deps SSCOR's pix2pix-style pipeline needs (torch/torchvision are the CPU +# builds installed above). +RUN pip install dominate visdom opencv-python matplotlib wandb + +# NOTE: SSCOR weights are NOT baked into this image -- the upstream repo +# distributes trained generator checkpoints via Google Drive, not a stable +# direct-download URL, and producing one requires an image-specific offline +# self-training stage (see ILLUMINATION_CORRECTION.md). For SSCOR mode +# "pretrained", supply a checkpoint at runtime by mounting it into the container +# and setting the SSCOR_WEIGHTS environment variable to its path (staged +# internally as the `latest_net_G_A.pth` the CycleGAN loader expects); the +# `sscor` Method sendError()s with instructions if it is unset/missing. SSCOR +# mode "self-train" needs no checkpoint (it trains one per frame). + +# entrypoint.py is copied last so routine edits don't invalidate the (slow) +# dependency/clone layers above. +COPY ./workers/annotations/illumination_correction/entrypoint.py / + +LABEL isUPennContrastWorker="" \ + isAnnotationWorker="" \ + interfaceName="Illumination Correction" \ + interfaceCategory="Image Processing" \ + description="Corrects illumination/shading/striping (BaSiC, CIDRE, CellProfiler, flat-field, destripe, SSCOR)" \ + hasPreview="False" \ + advancedOptionsPanel="False" \ + annotationConfigurationPanel="False" \ + defaultToolName="Illumination Correction" + +# ENTRYPOINT inherited from image-processing-base (run_worker.sh /entrypoint.py — see todo/worker-startup-latency.md) diff --git a/workers/annotations/illumination_correction/ILLUMINATION_CORRECTION.md b/workers/annotations/illumination_correction/ILLUMINATION_CORRECTION.md new file mode 100644 index 0000000..bcc4910 --- /dev/null +++ b/workers/annotations/illumination_correction/ILLUMINATION_CORRECTION.md @@ -0,0 +1,301 @@ +# Illumination Correction Worker + +This worker corrects spatial illumination artifacts in fluorescence microscopy images -- +shading, vignetting, background/dark bias, and scanner/tiling stripes -- using a single +`Method` dropdown that selects between six algorithms: **BaSiC**, a **CIDRE-style** +retrospective estimator, a **CellProfiler-style** illumination function, classical +**flat/dark-field** reference correction, **destripe** (wavelet-FFT stripe removal), and +**SSCOR** (deep-learning stripe self-correction). An optional lightweight QC report can be +generated after correction. + +## How It Works + +1. **Channel selection**: The user selects which channels to correct via "Channels to + correct" (a `channelCheckboxes` field). Unselected channels pass through to the + output completely unchanged. +2. **Collection gathering**: For each selected channel, the worker gathers every frame + of that channel across XY / Z / Time into a `(N, Y, X)` stack ("collection"). This is + the input to the retrospective methods (`basic`, `cidre`, `cellprofiler`); the + reference-based (`flatfield`) and per-frame classical (`destripe`) methods also + receive this same stack but process each frame independently. +3. **Correction**: `compute()` dispatches on `Method` to one standalone function + (`correct_basic`, `correct_cidre`, `correct_cellprofiler`, `correct_flatfield`, + `correct_destripe`, `correct_sscor`), each with the signature + `f(stack, opts) -> (corrected_stack, diagnostics)`. `sscor` is the one exception to the + "no upfront validation" rule below the fold, and its upfront check is conditional on + **SSCOR mode**: in `pretrained` mode (the default) it requires an externally-supplied + checkpoint, so `compute()` resolves `SSCOR_WEIGHTS` (via `resolve_sscor_checkpoint`) + *before* the channel loop and `sendError`-and-returns if no checkpoint is available; in + `self-train` mode no checkpoint is required at all (a `sendWarning` fires once instead, + noting that self-training is slow). GPU availability is always resolved before the + channel loop for `sscor`, in either mode. The resolved checkpoint path (or `None` for + `self-train`) and GPU choice are injected into `method_opts` -- mirroring how the + `flatfield` method's flat-reference check works. +4. **QC (optional)**: If "Report correction quality (QC)" is enabled, `compute_qc_metrics` + computes a few simple flat-field-quality numbers on the corrected collection for each + selected channel. +5. **dtype handling**: Corrected values are `np.nan_to_num`'d and, if the source dtype is + known (`tileClient.tiles['dtype']`), clipped and cast back to that dtype (e.g. + `uint16`) before being written, to avoid silently bloating the output TIFF to float. +6. **Output assembly**: A `large_image` sink is built frame-by-frame -- corrected frames + for selected channels, untouched frames for everything else -- preserving the original + `channelNames` / `mm_x` / `mm_y` / `magnification`, written to + `/tmp/illumination_corrected.tiff`, and uploaded back to Girder with a provenance + metadata dict (`tool`, `method`, `channels`, method-specific params, per-channel + `diagnostics`, and `qc` if requested). + +There is no conditional UI in this framework, so **all** parameters for **all** methods +are always shown; each parameter's tooltip states which method(s) it applies to, and +unused parameters for the currently-selected method are simply ignored. + +## Methods + +| Value | Label | Kind | Dependency | Needs reference frames? | +|-------|-------|------|-----------|--------------------------| +| `basic` | BaSiC (shading + darkfield) | retrospective | `basicpy` (BaSiCPy 2.x, PyTorch) | no | +| `cidre` | CIDRE-style retrospective | retrospective | in-house numpy/scipy | no | +| `cellprofiler` | CellProfiler-style illumination function | retrospective or per-image | in-house numpy/scipy | no | +| `flatfield` | Flat/Dark-field (reference-based) | reference | in-house numpy | **yes** | +| `destripe` | Stripe / tiling-seam correction | classical destriping | `pystripe` (wavelet-FFT) | no | +| `sscor` | Deep-learning stripe self-correction | pix2pix-style GAN, inference or per-frame self-training | vendored `SSCOR` repo; `pretrained` mode needs `SSCOR_WEIGHTS`, `self-train` mode needs none | no (`pretrained`: needs a trained checkpoint instead; `self-train`: samples/trains from the image itself) | + +Default `Method` is `basic`. + +## Interface Parameters + +| Parameter | Type | Default | Applies to | Description | +|-----------|------|---------|------------|-------------| +| **Method** | select | `basic` | all | Illumination-correction algorithm: `basic`, `cidre`, `cellprofiler`, `flatfield`, `destripe`, `sscor`. | +| **Channels to correct** | channelCheckboxes | -- | all | Channels to process; unselected channels pass through unchanged. | +| **Estimate darkfield** | checkbox | `True` | basic | Also estimate a darkfield (offset) term. | +| **Flatfield smoothness** | number (0-100) | `1.0` | basic | Smoothness regularization of the fitted flatfield. | +| **Darkfield smoothness** | number (0-100) | `1.0` | basic | Smoothness regularization of the fitted darkfield. | +| **Correct timelapse baseline drift** | checkbox | `False` | basic | Sort the channel's collection by Time and let BaSiC's transform correct for temporal drift; the per-frame baseline (if produced) is recorded in metadata. | +| **Smoothing sigma** | number (1-2000) | `50` | cidre, cellprofiler | Gaussian smoothing sigma (px) applied to the estimated illumination surface. | +| **Dark quantile** | number (0-0.5) | `0.02` | cidre | Per-pixel quantile used to estimate the dark/offset surface. | +| **CellProfiler mode** | select | `regular` | cellprofiler | `regular` = one illumination function from the across-batch average, applied to every frame; `background` = each frame gets its own smoothed background estimate. | +| **Flat-field XY coordinate** | text | -- | flatfield | 1-indexed XY position of the flat-field reference image. Required for `flatfield`. | +| **Dark-field XY coordinate** | text | -- | flatfield | 1-indexed XY position of the dark-field reference image; empty = use the Dark-field constant instead. | +| **Dark-field constant** | number (0-65535) | `0` | flatfield | Constant dark offset used when no dark-field reference frame is given. | +| **Destripe sigma** | number (1-2000) | `128` | destripe | Band-pass sigma for pystripe's wavelet-FFT stripe removal. | +| **Destripe wavelet** | select | `db3` | destripe | Wavelet family: `db3`, `db5`, `haar`, `sym4`. | +| **Destripe level** | number (0-12) | `0` | destripe | DWT decomposition level (`0` = auto). | +| **SSCOR mode** | select | `pretrained` | sscor | `pretrained` = use a checkpoint trained offline, supplied via `SSCOR_WEIGHTS`; `self-train` = no checkpoint needed -- sample stripe/clean patches from the image itself and train a fresh CycleGAN per frame before restoring. | +| **SSCOR stripe direction** | select | `horizontal` | sscor (self-train) | Stripe orientation to sample training patches for: `horizontal`, `vertical`, or `grid` (both, via `sample_stripe_2.py`). | +| **SSCOR horizontal stripe count** | number (1-4096) | `1` | sscor (self-train) | Number of horizontal stripe bands (`--h_n`) used when sampling training patches. | +| **SSCOR vertical stripe count** | number (1-4096) | `1` | sscor (self-train) | Number of vertical stripe bands (`--v_n`) used when sampling training patches. | +| **SSCOR grid direction** | number (0-3) | `0` | sscor (self-train, direction=grid) | Junction direction (0=Upper Left, 1=Upper Right, 2=Lower Left, 3=Lower Right) passed as `sample_stripe_2.py`'s `--direction`. | +| **SSCOR training epochs** | number (1-300) | `30` | sscor (self-train) | Number of training epochs for the per-frame CycleGAN self-training pass. | +| **SSCOR patch size** | number (1-4096) | `256` | sscor | Sliding-window patch size (px) fed to the generator. | +| **SSCOR offset size** | number (1-4096) | `100` | sscor | Sliding-window step/offset size (px) between patches. | +| **SSCOR repeat** | number (1-5) | `1` | sscor | Number of repeated passes (with shifted offsets) combined via max-projection, per upstream `restore.py`. | +| **SSCOR dark threshold** | number (0-255) | `10` | sscor | 8-bit intensity threshold below which the original (uncorrected) pixel is kept instead of the restored value. | +| **Report correction quality (QC)** | checkbox | `False` | all | Computes lightweight EVEN-style QC metrics per corrected channel and stores them in the output item's metadata. | + +## Implementation Details + +### Retrospective collection semantics + +For `basic`, `cidre`, and `cellprofiler`, the "collection" fit on is every frame of the +selected channel across XY, Z, and Time -- gathered into a `(N, Y, X)` stack, fit once, +then applied (`transform`) to every frame in that stack. If a channel's collection has +fewer than 3 frames, `sendWarning` fires: retrospective illumination estimation is +unreliable with very few images. + +### Method notes and honest simplifications + +- **`basic`**: Uses BaSiCPy 2.x (`from basicpy import BaSiC`), PyTorch backend, CPU-only + in this worker (no GPU dependency). `get_darkfield`, `smoothness_flatfield`, + `smoothness_darkfield` map directly to the corresponding interface parameters; + `fitting_mode='approximate'`. When "Correct timelapse baseline drift" is enabled, the + collection is sorted by `IndexT` before fitting/transforming and un-sorted afterward; + the fitted flatfield/darkfield means (and baseline, if produced) are stored per-channel + under `diagnostics` in the output item's metadata for provenance. +- **`cidre`**: **This is a lightweight, in-house reimplementation of CIDRE's + retrospective gain/offset model (Smith et al. 2015, Nature Methods) -- not the + original MATLAB implementation and its full energy-minimization solve.** The offset + (dark) surface is estimated as a robust per-pixel low quantile (`Dark quantile`, + default 2%) across the collection, then heavily Gaussian-smoothed (`Smoothing sigma`). + The gain (flat) surface is the per-pixel median of `frame - offset`, normalized to + mean 1 and smoothed the same way, as a spatial-regularization stand-in for CIDRE's + energy minimization. `corrected = (frame - offset) / gain`. +- **`cellprofiler`**: Reimplements the core of CellProfiler's + `CorrectIlluminationCalculate` module without a CellProfiler dependency. `regular` + mode averages the whole collection, Gaussian-smooths it, and normalizes to mean 1 to + get one illumination function applied to every frame; `background` mode does the same + per-frame, independently. `corrected = frame / illumination_function`. +- **`flatfield`**: Classical `corrected = (raw - dark) / (flat - dark)`, with the gain + `(flat - dark)` normalized so its mean is 1 (preserving the original intensity scale). + The flat/dark references are read at `(that XY, Z=0, T=0, that channel)` for the + channel currently being corrected. If no dark reference is given, dark is treated as + the constant `Dark-field constant` (default 0). `flatfield` requires a flat-field + reference; if `Flat-field XY coordinate` is empty, the worker calls `sendError` and + exits rather than silently producing a meaningless correction. +- **`destripe`**: Uses `pystripe` (`pystripe.core.filter_streaks`), a wavelet-FFT + destriping method originally built for light-sheet microscopy, applied independently + per frame. This is the fast, classical, **no-weights-needed** stripe-removal option -- + prefer it when a trained `sscor` checkpoint isn't available. +- **`sscor`**: Deep-learning stripe self-correction, vendored from + [`lxxcontinue/SSCOR`](https://github.com/lxxcontinue/SSCOR) (pinned commit + `985479cd79bcf1359e3d9ba44bacd5f372eb2e60`). See "SSCOR integration details" below. +- **QC report (`compute_qc_metrics`)**: **This is a lightweight quantitative stand-in + for EVEN** (Nat. Commun. 2026, an ML/LDA-based illumination evaluation-and-optimization + framework) -- EVEN itself is not vendored here. The QC report computes, per corrected + channel: `cv_mean_image` (coefficient of variation of a smoothed mean image -- lower is + flatter), `corner_center_ratio` (mean corner vs. mean center intensity of the mean + image -- closer to 1.0 means less residual vignetting), and `interframe_cv` + (coefficient of variation of per-frame mean intensity across the collection). These are + meant only to let a user quickly compare methods against each other, not as a + publication-grade quality metric. + +### SSCOR integration details + +SSCOR ([`lxxcontinue/SSCOR`](https://github.com/lxxcontinue/SSCOR), pinned commit +`985479cd79bcf1359e3d9ba44bacd5f372eb2e60`) is a pytorch-CycleGAN-and-pix2pix-style +codebase with no clean importable inference API, so it is integrated exactly like +`deconwolf` shells out to the `dw` binary: via `subprocess`, driving the repo's CLI scripts +directly (see `correct_sscor` in `entrypoint.py`). Two modes are available via **SSCOR +mode**, both dispatched from the same `correct_sscor` function: + +#### `pretrained` mode + +Runs ONLY the repo's inference stage, using a checkpoint trained offline. + +- **Env-gated weights (`SSCOR_WEIGHTS`)**: SSCOR's trained generator checkpoints are not + baked into the Docker image -- the upstream repo distributes them via Google Drive, not + a stable direct-download URL, and producing one requires an image-specific **offline + self-training stage** (proximity sampling + adversarial training with stripe-orientation + parameters tuned to the dataset) -- see `self-train` mode below for a version of that + stage this worker runs itself. `resolve_sscor_checkpoint` reads the `SSCOR_WEIGHTS` + environment variable; if unset or not a file, `compute()` `sendError`s with actionable + instructions (download a checkpoint per the upstream README, mount it into the + container, set `SSCOR_WEIGHTS` to its path, use `self-train` mode instead, or choose a + different Method) **before the channel loop runs**, mirroring the `flatfield` method's + upfront flat-reference check. +- **Subprocess-to-`restore.py` integration**: For each frame, `correct_sscor` writes an + 8-bit RGB PNG to a temp "dataroot" directory, copies the resolved checkpoint to + `/sscor/latest_net_G_A.pth`, and invokes + `restore.py --dataroot --name sscor --model sscor --image_name --offset_size + --patch_size --repeat + --dark_threshold --checkpoints_dir --gpu_ids + <0 or -1> --eval` with `cwd` set to the vendored repo (`SSCOR_REPO_PATH`, default + `/sscor`). The restored image is read back from `/result/restore-` and + converted from RGB back to a single intensity channel (mean over channels). A non-zero + return code raises `RuntimeError` with the captured `stderr`, matching how `deconwolf` + surfaces `dw` failures. +- **Checkpoint filename (`latest_net_G_A.pth`)**: at inference time + `SSCORModel.model_names == ['G_A']` (only the A->B generator is needed for restoration), + and `BaseModel.load_networks` loads `'%s_net_%s.pth' % (epoch, name)` with `epoch='latest'` + -- i.e. it looks for `latest_net_G_A.pth`, never `latest_net_G.pth`. `correct_sscor` + therefore stages the user-supplied `SSCOR_WEIGHTS` file under the corrected name + `latest_net_G_A.pth` regardless of its original filename. (An earlier version of this + integration staged it as `latest_net_G.pth`, which `restore.py` would silently fail to + find/load correctly -- this has been fixed.) + +#### `self-train` mode + +Runs the **faithful upstream self-training pipeline** end-to-end, per frame, with no +pre-supplied checkpoint at all -- SSCOR's real design point is per-image self-training, +not a globally pretrained model. For each frame, `correct_sscor`: + +1. **Samples** stripe-free "clean" patches (`trainB`) and striped patches (`trainA`) out + of the frame itself, via the upstream `sample/sample_stripe.py` (`--h`/`--v` for a pure + horizontal or vertical stripe pattern) or `sample/sample_stripe_2.py` (`--h_n`/`--v_n`/ + `--direction` for a grid of both, used when **SSCOR stripe direction**=`grid`), + controlled by **SSCOR stripe direction**, **SSCOR horizontal stripe count**, **SSCOR + vertical stripe count**, **SSCOR grid direction**, and **SSCOR patch size**. If no + patches are produced (frame too small for the requested patch size, or the stripe + direction/count doesn't fit the image), `correct_sscor` raises a `RuntimeError` naming + the frame and suggesting a smaller patch size or a different stripe direction/count. +2. **Trains** a fresh CycleGAN from scratch on those self-sampled patches via `train.py + --dataroot --name sscor_selftrain --model sscor --checkpoints_dir + --gpu_ids <0 or -1> --display_id 0 --no_html --load_size --crop_size + --n_epochs --n_epochs_decay 0 --save_epoch_freq + `. Setting `--save_epoch_freq` equal to the epoch count + guarantees `model.save_networks('latest')` (-> `latest_net_G_A.pth`) is written exactly + once, at the final epoch. `--display_id 0` avoids a visdom dependency; `--no_html` + avoids writing intermediate HTML previews. +3. **Restores** the frame with the just-trained generator via the same `restore.py` + invocation as `pretrained` mode, pointed at the freshly trained + `--name sscor_selftrain --checkpoints_dir `. + +Each frame gets its own fresh temporary directories for sampling output, checkpoints, and +data, so nothing leaks between frames' independent self-training runs. Because a full +CycleGAN is trained from scratch per frame, this mode is **much slower** than `pretrained` +mode and a GPU is strongly recommended; `compute()` sends a one-time `sendWarning` when +`self-train` is selected. + +#### Shared across both modes + +- **8-bit operation (lossy)**: SSCOR's generator only supports 8-bit RGB I/O (PIL-loaded + input, `tensor2im`-produced `uint8` output). Each frame is therefore rescaled to `uint8` + `[0, 255]` via per-frame min/max before being handed to the SSCOR scripts, and the 8-bit + result is rescaled back to the frame's original `[min, max]` range afterward. **This + round trip is inherently lossy** (8-bit quantization) -- it is a limitation of SSCOR's + design, not an artifact of this integration. +- **GPU strongly recommended**: `compute()` resolves GPU availability once per run (a + lazy, best-effort `torch.cuda.is_available()` check; if `torch` isn't importable, it + falls back to CPU) and passes `--gpu_ids 0` or `--gpu_ids -1` accordingly, for both + modes. If falling back to CPU, `sendWarning` fires once: SSCOR on CPU is very slow (and + for `self-train`, which also trains a model per frame, doubly so). + +### Numerical stability safeguard + +Every mean-normalized gain/illumination surface (`cidre`'s gain, `cellprofiler`'s +illumination function, `flatfield`'s gain) is floored at 5% of its own mean +(`_MIN_GAIN_FRACTION` in `entrypoint.py`) before dividing. The "real" CIDRE and +CellProfiler algorithms solve a regularized (well-posed) model that cannot collapse to +near-zero gain; this worker's lightweight per-pixel-statistics-plus-smoothing stand-ins +have no such built-in guarantee, so without a floor a handful of near-zero pixels (heavy +vignetting at the extreme corners, or noise-dominated low-signal regions) can amplify +noise to enormous corrected values. The 5% floor bounds local amplification to at most +20x while barely affecting well-behaved regions -- this was verified empirically against +synthetic shading/noise stacks during development. + +### dtype and coordinate handling + +Corrected stacks are computed in `float64` internally. Before being written to the +output sink, they are passed through `np.nan_to_num` (guarding against NaN/inf from any +of the algorithms) and, if `tileClient.tiles['dtype']` is available, clipped to that +dtype's valid range (`np.iinfo` for integer types) and cast back -- this avoids silently +bloating a 16-bit input into a float64 output TIFF. Frame iteration and the +`large_image` sink assembly follow the same `Index*` -> lowercase-key convention used by +`histogram_matching` / `deconwolf` (`IndexXY` -> `xy`, `IndexZ` -> `z`, etc.). + +### Heavy imports are lazy + +`basicpy` (inside `correct_basic`), `pystripe` (inside `correct_destripe`), and +`subprocess`/`PIL`/`torch` (inside `correct_sscor` and the upfront GPU-detection check in +`compute()`) are imported **inside** those functions only, never at module top level. This +keeps the `interface()` preview path fast (see `todo/worker-startup-latency.md`) and lets +the test suite run natively without any of those packages installed -- tests patch +`entrypoint.correct_basic` / `entrypoint.correct_cidre` / `entrypoint.correct_sscor` / etc. +directly, mirroring `histogram_matching/tests/test_histogram_matching.py`'s pattern of +patching `entrypoint.match_histograms`. (`entrypoint.resolve_sscor_checkpoint` is likewise +patched directly in tests rather than exercising the real `SSCOR_WEIGHTS` env-var check.) + +## Notes + +- **No GPU required for five of the six methods.** `basic`, `cidre`, `cellprofiler`, + `flatfield`, and `destripe` all run on CPU; BaSiCPy's PyTorch backend runs fine without + CUDA for the collection sizes typical of a single dataset. `sscor` is the exception -- + see below. +- **Small collections**: retrospective methods (`basic`, `cidre`, `cellprofiler`) warn + when a channel's collection has fewer than 3 frames; results are unreliable in that + regime regardless of method. +- **`flatfield` requires calibration frames** in the dataset itself (identified by XY + coordinate); if you don't have dedicated flat/dark calibration images, use `basic`, + `cidre`, or `cellprofiler` instead. +- **`sscor` has two modes, and a GPU is strongly recommended for both.** `pretrained` mode + (the default) needs a user-supplied generator checkpoint (`SSCOR_WEIGHTS`, see "SSCOR + integration details" above) produced offline via SSCOR's self-training procedure -- + there is no bundled/default checkpoint. `self-train` mode needs no checkpoint at all -- + it runs that same self-training procedure itself, from scratch, per frame -- but is + correspondingly much slower (a full CycleGAN trained from zero for every frame). If + you don't have a checkpoint and don't want to wait for `self-train`, use `destripe` + (`pystripe`) instead: it needs no weights and is much faster, at the cost of being a + classical (non-deep-learning) method. Inference-only (`pretrained`) also works on CPU + but is very slow; `self-train` on CPU is slower still. +- Related workers: `histogram_matching` (post-hoc intensity histogram matching across + frames), `deconwolf` (deconvolution, a different kind of optical correction), + `rolling_ball` (simple background subtraction). diff --git a/workers/annotations/illumination_correction/entrypoint.py b/workers/annotations/illumination_correction/entrypoint.py new file mode 100644 index 0000000..b5e0d62 --- /dev/null +++ b/workers/annotations/illumination_correction/entrypoint.py @@ -0,0 +1,1166 @@ +import argparse +import json +import os +import sys + +import numpy as np + +import annotation_client.tiles as tiles +import annotation_client.workers as workers + +from annotation_client.utils import sendProgress, sendWarning, sendError + + +# Floor applied to every mean-normalized gain/illumination surface (which is +# normalized so its mean is 1.0) before dividing. Real CIDRE/CellProfiler solve +# a regularized (well-posed) model that can't collapse to near-zero gain; our +# lightweight per-pixel-statistics-plus-smoothing stand-in has no such +# guarantee, so without a floor a handful of near-zero pixels (e.g. heavy +# vignetting at the extreme corners, or noise-dominated low-signal regions) +# can blow up to enormous corrected values. Capping the gain at 5% of its own +# mean keeps the correction bounded (<=20x local amplification) while barely +# affecting well-behaved regions. +_MIN_GAIN_FRACTION = 0.05 + + +# --------------------------------------------------------------------------- +# Interface +# --------------------------------------------------------------------------- + +def interface(image, apiUrl, token): + client = workers.UPennContrastWorkerPreviewClient( + apiUrl=apiUrl, token=token) + + interface = { + 'Method': { + 'type': 'select', + 'items': ['basic', 'cidre', 'cellprofiler', 'flatfield', 'destripe', 'sscor'], + 'default': 'basic', + 'tooltip': 'Illumination-correction algorithm. basic (BaSiC) recommended when no ' + 'calibration frames are available. sscor is a deep-learning stripe ' + 'self-correction method that requires a user-supplied trained checkpoint ' + '(see SSCOR_WEIGHTS). Every parameter below is shown regardless of Method; ' + 'each tooltip notes which method(s) it applies to -- unused parameters for ' + 'the chosen method are simply ignored.', + 'displayOrder': 0, + }, + 'Channels to correct': { + 'type': 'channelCheckboxes', + 'tooltip': 'Process selected channels; unselected channels pass through unchanged.', + 'displayOrder': 1, + }, + # --- BaSiC --- + 'Estimate darkfield': { + 'type': 'checkbox', + 'default': True, + 'tooltip': 'basic: also estimate a darkfield (offset) term.', + 'displayOrder': 2, + }, + 'Flatfield smoothness': { + 'type': 'number', + 'min': 0, + 'max': 100, + 'default': 1.0, + 'tooltip': 'basic: smoothness regularization of the flatfield.', + 'displayOrder': 3, + }, + 'Darkfield smoothness': { + 'type': 'number', + 'min': 0, + 'max': 100, + 'default': 1.0, + 'tooltip': 'basic: smoothness regularization of the darkfield.', + 'displayOrder': 4, + }, + 'Correct timelapse baseline drift': { + 'type': 'checkbox', + 'default': False, + 'tooltip': 'basic: estimate and remove temporal background/bleaching drift across Time.', + 'displayOrder': 5, + }, + # --- CIDRE / CellProfiler shared --- + 'Smoothing sigma': { + 'type': 'number', + 'min': 1, + 'max': 2000, + 'default': 50, + 'tooltip': 'cidre/cellprofiler: Gaussian smoothing sigma (px) for the illumination surface.', + 'displayOrder': 6, + }, + 'Dark quantile': { + 'type': 'number', + 'min': 0, + 'max': 0.5, + 'default': 0.02, + 'tooltip': 'cidre: percentile for the dark/offset surface estimate.', + 'displayOrder': 7, + }, + 'CellProfiler mode': { + 'type': 'select', + 'items': ['regular', 'background'], + 'default': 'regular', + 'tooltip': 'cellprofiler: regular=across-batch average; background=per-image background.', + 'displayOrder': 8, + }, + # --- flatfield reference --- + 'Flat-field XY coordinate': { + 'type': 'text', + 'vueAttrs': { + 'placeholder': 'ex. 1', + 'label': 'Flat-field XY coordinate', + 'persistentPlaceholder': True, + 'filled': True, + }, + 'tooltip': 'flatfield: 1-indexed XY position of the flat-field reference image. ' + 'Required for the flatfield method.', + 'displayOrder': 9, + }, + 'Dark-field XY coordinate': { + 'type': 'text', + 'vueAttrs': { + 'placeholder': 'ex. 1', + 'label': 'Dark-field XY coordinate', + 'persistentPlaceholder': True, + 'filled': True, + }, + 'tooltip': 'flatfield: 1-indexed XY position of the dark-field reference ' + '(empty = use the Dark-field constant instead).', + 'displayOrder': 10, + }, + 'Dark-field constant': { + 'type': 'number', + 'min': 0, + 'max': 65535, + 'default': 0, + 'tooltip': 'flatfield: constant dark offset used if no dark-field reference frame is given.', + 'displayOrder': 11, + }, + # --- destripe --- + 'Destripe sigma': { + 'type': 'number', + 'min': 1, + 'max': 2000, + 'default': 128, + 'tooltip': 'destripe: band-pass sigma for wavelet-FFT stripe removal.', + 'displayOrder': 12, + }, + 'Destripe wavelet': { + 'type': 'select', + 'items': ['db3', 'db5', 'haar', 'sym4'], + 'default': 'db3', + 'tooltip': 'destripe: wavelet used for the wavelet-FFT decomposition.', + 'displayOrder': 13, + }, + 'Destripe level': { + 'type': 'number', + 'min': 0, + 'max': 12, + 'default': 0, + 'tooltip': 'destripe: DWT decomposition level (0 = auto).', + 'displayOrder': 14, + }, + # --- SSCOR --- + 'SSCOR mode': { + 'type': 'select', + 'items': ['pretrained', 'self-train'], + 'default': 'pretrained', + 'tooltip': 'sscor: how the generator checkpoint is obtained. pretrained = use a ' + 'checkpoint trained offline and supplied via the SSCOR_WEIGHTS ' + 'environment variable (fast, but you must already have a trained ' + 'checkpoint). self-train = no checkpoint needed -- SSCOR samples ' + 'stripe-free "clean" patches and striped patches from the image itself ' + 'and trains a small CycleGAN on them before restoring, faithfully ' + 'reproducing the upstream per-image self-training method. self-train ' + 'trains a fresh model PER FRAME on a GPU and is slow.', + 'displayOrder': 15, + }, + 'SSCOR stripe direction': { + 'type': 'select', + 'items': ['horizontal', 'vertical', 'grid'], + 'default': 'horizontal', + 'tooltip': 'sscor (self-train): stripe orientation to sample training patches for -- ' + 'horizontal, vertical, or grid (both directions, using the upstream ' + 'sample_stripe_2.py corner/junction sampling). Trains a fresh model PER ' + 'FRAME on a GPU and is slow.', + 'displayOrder': 16, + }, + 'SSCOR horizontal stripe count': { + 'type': 'number', + 'min': 1, + 'max': 4096, + 'default': 1, + 'tooltip': 'sscor (self-train): number of horizontal stripe bands (--h_n) used when ' + 'sampling training patches (horizontal or grid stripe direction). Trains ' + 'a fresh model PER FRAME on a GPU and is slow.', + 'displayOrder': 17, + }, + 'SSCOR vertical stripe count': { + 'type': 'number', + 'min': 1, + 'max': 4096, + 'default': 1, + 'tooltip': 'sscor (self-train): number of vertical stripe bands (--v_n) used when ' + 'sampling training patches (vertical or grid stripe direction). Trains a ' + 'fresh model PER FRAME on a GPU and is slow.', + 'displayOrder': 18, + }, + 'SSCOR grid direction': { + 'type': 'number', + 'min': 0, + 'max': 3, + 'default': 0, + 'tooltip': 'sscor (self-train): junction direction (0=Upper Left, 1=Upper Right, ' + '2=Lower Left, 3=Lower Right) passed as sample_stripe_2.py\'s --direction; ' + 'only used when SSCOR stripe direction=grid. Trains a fresh model PER ' + 'FRAME on a GPU and is slow.', + 'displayOrder': 19, + }, + 'SSCOR training epochs': { + 'type': 'number', + 'min': 1, + 'max': 300, + 'default': 30, + 'tooltip': 'sscor (self-train): number of training epochs for the per-frame CycleGAN ' + 'self-training pass. Trains a fresh model PER FRAME on a GPU and is slow.', + 'displayOrder': 20, + }, + 'SSCOR patch size': { + 'type': 'number', + 'min': 1, + 'max': 4096, + 'default': 256, + 'tooltip': 'sscor: sliding-window patch size (px) fed to the generator.', + 'displayOrder': 21, + }, + 'SSCOR offset size': { + 'type': 'number', + 'min': 1, + 'max': 4096, + 'default': 100, + 'tooltip': 'sscor: sliding-window step/offset size (px) between patches.', + 'displayOrder': 22, + }, + 'SSCOR repeat': { + 'type': 'number', + 'min': 1, + 'max': 5, + 'default': 1, + 'tooltip': 'sscor: number of repeated passes (with shifted offsets) to combine via ' + 'max-projection, per the upstream restore.py.', + 'displayOrder': 23, + }, + 'SSCOR dark threshold': { + 'type': 'number', + 'min': 0, + 'max': 255, + 'default': 10, + 'tooltip': 'sscor: 8-bit intensity threshold below which the original (uncorrected) ' + 'pixel is kept instead of the restored value.', + 'displayOrder': 24, + }, + # --- QC --- + 'Report correction quality (QC)': { + 'type': 'checkbox', + 'default': False, + 'tooltip': 'Compute lightweight EVEN-style flat-field-quality metrics per corrected ' + 'channel and store them in the output item metadata. Not the full EVEN ' + 'framework -- a quick quantitative stand-in for comparing methods.', + 'displayOrder': 25, + }, + } + # Send the interface object to the server + client.setWorkerImageInterface(image, interface) + + +# --------------------------------------------------------------------------- +# Per-algorithm functions +# +# Each function has the signature `f(stack, opts) -> (corrected_stack, diagnostics)` +# where `stack` is a (N, Y, X) float array containing every frame of ONE channel +# (gathered across XY/Z/Time for the retrospective methods; a simple per-frame +# stack for flatfield/destripe), and `diagnostics` is a small JSON-serializable +# dict recorded for provenance. +# +# Heavy third-party imports (basicpy, pystripe) are done LAZILY inside these +# functions so that `interface()` stays fast and so tests can run natively +# without those packages installed (the tests patch these functions directly). +# --------------------------------------------------------------------------- + +def correct_basic(stack, opts): + """BaSiC retrospective shading (+ darkfield) correction (Peng et al. 2017; + BaSiCPy 2.x, PyTorch backend). + + stack: (N, Y, X) array -- every frame of one channel's collection. + opts: dict with 'estimate_darkfield' (bool), 'flatfield_smoothness' (float), + 'darkfield_smoothness' (float), 'baseline_drift' (bool), and optionally + 'time_order' (list[int] aligned with stack) used only to sort frames + chronologically before fitting when baseline_drift is requested. + """ + from basicpy import BaSiC + + stack = np.asarray(stack, dtype=np.float64) + + order = None + if opts.get('baseline_drift') and opts.get('time_order') is not None: + order = np.argsort(np.asarray(opts['time_order'])) + fit_stack = stack[order] + else: + fit_stack = stack + + basic = BaSiC( + get_darkfield=bool(opts.get('estimate_darkfield', True)), + smoothness_flatfield=float(opts.get('flatfield_smoothness', 1.0)), + smoothness_darkfield=float(opts.get('darkfield_smoothness', 1.0)), + fitting_mode='approximate', + ) + basic.fit(fit_stack) + corrected_fit_order = np.asarray(basic.transform(fit_stack)) + + if order is not None: + corrected = np.empty_like(corrected_fit_order) + corrected[order] = corrected_fit_order + else: + corrected = corrected_fit_order + + diagnostics = { + 'flatfield_mean': float(np.mean(basic.flatfield)), + 'darkfield_mean': float(np.mean(basic.darkfield)) if opts.get('estimate_darkfield', True) and getattr(basic, 'darkfield', None) is not None else 0.0, + } + baseline = getattr(basic, 'baseline', None) + if opts.get('baseline_drift') and baseline is not None: + diagnostics['baseline'] = [float(b) for b in np.asarray(baseline).ravel()] + + return corrected, diagnostics + + +def correct_cidre(stack, opts): + """CIDRE-style retrospective illumination correction (Smith et al. 2015, + Nature Methods). + + NOTE (honest simplification): this is a lightweight, in-house + reimplementation of CIDRE's retrospective gain/offset model, NOT the + original MATLAB implementation and its full energy-minimization solve. + The offset (dark) and gain (flat) surfaces here are estimated with robust + per-pixel statistics and then heavily Gaussian-smoothed as a spatial + regularization stand-in. + """ + from scipy.ndimage import gaussian_filter + + stack = np.asarray(stack, dtype=np.float64) + sigma = float(opts.get('smoothing_sigma', 50)) + q = float(opts.get('dark_quantile', 0.02)) + + # Offset (dark) surface: robust low-quantile per pixel across the collection. + z = np.quantile(stack, q, axis=0) + z = gaussian_filter(z, sigma=sigma) + + # Gain (flat) surface: per-pixel median of (frame - offset), normalized to mean 1. + residual = stack - z[np.newaxis, :, :] + v = np.median(residual, axis=0) + v = gaussian_filter(v, sigma=sigma) + v_mean = np.mean(v) + if v_mean <= 0: + v_mean = 1.0 + v = v / v_mean + v = np.clip(v, _MIN_GAIN_FRACTION, None) + + corrected = (stack - z[np.newaxis, :, :]) / v[np.newaxis, :, :] + + diagnostics = { + 'offset_mean': float(np.mean(z)), + 'gain_mean': float(np.mean(v)), + } + return corrected, diagnostics + + +def correct_cellprofiler(stack, opts): + """CellProfiler-style illumination function, reimplemented without a + CellProfiler dependency (core logic of `CorrectIlluminationCalculate`). + + mode='regular' (default): one illumination function is fit from the + across-batch mean image and applied to every frame. + mode='background': each frame supplies its own heavily-smoothed background + estimate, applied per frame independently. + """ + from scipy.ndimage import gaussian_filter + + stack = np.asarray(stack, dtype=np.float64) + sigma = float(opts.get('smoothing_sigma', 50)) + mode = opts.get('cellprofiler_mode', 'regular') + + if mode == 'background': + corrected = np.empty_like(stack) + illum_means = [] + for i in range(stack.shape[0]): + illum = gaussian_filter(stack[i], sigma=sigma) + illum_mean = np.mean(illum) + if illum_mean <= 0: + illum_mean = 1.0 + illum = np.clip(illum / illum_mean, _MIN_GAIN_FRACTION, None) + corrected[i] = stack[i] / illum + illum_means.append(float(illum_mean)) + diagnostics = {'mode': 'background', 'illumination_mean': illum_means} + else: + avg_image = np.mean(stack, axis=0) + illum = gaussian_filter(avg_image, sigma=sigma) + illum_mean = np.mean(illum) + if illum_mean <= 0: + illum_mean = 1.0 + illum = np.clip(illum / illum_mean, _MIN_GAIN_FRACTION, None) + corrected = stack / illum[np.newaxis, :, :] + diagnostics = {'mode': 'regular', 'illumination_mean': float(illum_mean)} + + return corrected, diagnostics + + +def correct_flatfield(stack, opts): + """Classical flat-field / dark-field reference-based correction: + corrected = (raw - dark) / (flat - dark), rescaled to preserve mean intensity. + + opts must include 'flat' (2D array) and 'dark' (2D array, same shape). + """ + stack = np.asarray(stack, dtype=np.float64) + flat = np.asarray(opts['flat'], dtype=np.float64) + dark = np.asarray(opts['dark'], dtype=np.float64) + + flat_minus_dark = flat - dark + flat_mean = np.mean(flat_minus_dark) + if flat_mean <= 0: + flat_mean = 1.0 + # Normalize gain so its mean is 1 after dark subtraction, to preserve the + # overall intensity scale. Floor at _MIN_GAIN_FRACTION to avoid blow-up from + # near-zero gain at extreme vignetting (see _MIN_GAIN_FRACTION docstring). + gain = np.clip(flat_minus_dark / flat_mean, _MIN_GAIN_FRACTION, None) + + corrected = (stack - dark[np.newaxis, :, :]) / gain[np.newaxis, :, :] + corrected = np.clip(corrected, 0, None) + + diagnostics = { + 'flat_mean': float(np.mean(flat)), + 'dark_mean': float(np.mean(dark)), + } + return corrected, diagnostics + + +def correct_destripe(stack, opts): + """Classical wavelet-FFT destriping via pystripe, applied independently per + frame. Documented stand-in for the DL method SSCOR, which has no packaged + weights available. + """ + from pystripe import core as pystripe_core + + stack = np.asarray(stack, dtype=np.float64) + sigma = float(opts.get('destripe_sigma', 128)) + wavelet = opts.get('destripe_wavelet', 'db3') + level = int(opts.get('destripe_level', 0)) + level_arg = None if level == 0 else level + + corrected = np.empty_like(stack) + for i in range(stack.shape[0]): + corrected[i] = pystripe_core.filter_streaks( + stack[i], sigma=[sigma, sigma], level=level_arg, wavelet=wavelet) + + diagnostics = {'sigma': sigma, 'wavelet': wavelet, 'level': level} + return corrected, diagnostics + + +def resolve_sscor_checkpoint(): + """Resolve the path to the SSCOR generator checkpoint (pretrained mode only). + + SSCOR weights are not bundled in the Docker image -- the upstream repo + (https://github.com/lxxcontinue/SSCOR) distributes trained models via + Google Drive, not a stable direct-download URL, and producing one requires + an image-specific self-training stage (proximity sampling + adversarial + training) that must either be run offline per the upstream README, or via + this worker's own `SSCOR mode=self-train` (see `correct_sscor`). The + runtime source of truth for the *pretrained* mode is therefore the + SSCOR_WEIGHTS environment variable, which should point to a mounted + generator checkpoint (the CycleGAN's `G_A` weights; at inference time + SSCOR loads it as `latest_net_G_A.pth`, regardless of the file's original + name -- `correct_sscor` stages it under that name). Returns the resolved + path, or None (after sending an actionable sendError) if unavailable. + """ + weights_path = os.environ.get('SSCOR_WEIGHTS', '').strip() + if not weights_path or not os.path.isfile(weights_path): + sendError( + 'SSCOR weights are not available.', + info=( + 'SSCOR weights are not bundled in this image -- they are distributed via ' + 'Google Drive from https://github.com/lxxcontinue/SSCOR. Either download a ' + 'trained generator checkpoint, mount it into the container, and set the ' + 'SSCOR_WEIGHTS environment variable to its path (the CycleGAN G_A generator ' + 'weights, staged internally as `latest_net_G_A.pth`), or set "SSCOR mode" to ' + '"self-train" to train a model from the image itself with no checkpoint, or ' + 'choose a different Method (e.g. destripe, which needs no weights).' + ), + ) + return None + return weights_path + + +def _sscor_frame_to_uint8(frame): + """Rescale one (Y, X) float frame to uint8 [0, 255] via per-frame min/max. + Returns (frame_min, span, frame_uint8); `span` is guaranteed > 0, so the + inverse rescale `restored_gray / 255.0 * span + frame_min` always applies. + """ + frame_min = float(np.min(frame)) + frame_max = float(np.max(frame)) + span = frame_max - frame_min + if span <= 0: + span = 1.0 + frame_uint8 = np.clip((frame - frame_min) / span * 255.0, 0, 255).astype(np.uint8) + return frame_min, span, frame_uint8 + + +def correct_sscor(stack, opts): + """SSCOR deep-learning stripe self-correction + (https://github.com/lxxcontinue/SSCOR, pinned commit + 985479cd79bcf1359e3d9ba44bacd5f372eb2e60). + + SSCOR is a pytorch-CycleGAN-and-pix2pix-style codebase with no clean + importable inference API, so this integration shells out to its CLI + scripts exactly like `deconwolf` shells out to the `dw` binary. Two modes + are supported, selected by `opts['sscor_mode']`: + + - 'pretrained': runs ONLY the repo's inference stage (`restore.py`), + using a user-supplied trained generator checkpoint resolved by + `resolve_sscor_checkpoint` (SSCOR_WEIGHTS). No training happens here. + - 'self-train': the faithful upstream SSCOR pipeline, run once PER FRAME + with no pre-supplied checkpoint. For each frame: (1) `sample/sample_stripe.py` + or `sample/sample_stripe_2.py` samples stripe-free "clean" patches + (trainB) and striped patches (trainA) from the image itself, using the + chosen stripe direction/count; (2) `train.py` trains a small CycleGAN + on those self-sampled patches for `sscor_epochs` epochs, saving + `latest_net_G_A.pth` once at the end (`--save_epoch_freq` == epoch + count); (3) `restore.py` restores the frame with that just-trained + generator. This is slow (a fresh model is trained from scratch per + frame) and a GPU is strongly recommended. + + SSCOR's generator only supports 8-bit RGB I/O (PIL-loaded input, + `tensor2im`-produced uint8 output), so each frame is rescaled to uint8 + [0, 255] via per-frame min/max before being handed to the SSCOR scripts, + and the 8-bit result is rescaled back to the frame's original [min, max] + range afterward. This round trip is LOSSY (8-bit quantization) and is + inherent to SSCOR's design, not an artifact of this integration. + + stack: (N, Y, X) array -- every frame of one channel's collection. + opts: dict with 'sscor_mode' ('pretrained' or 'self-train'), + 'sscor_patch_size', 'sscor_offset_size', 'sscor_repeat', + 'sscor_dark_threshold' (all int), 'sscor_gpu_ids' ('-1' for CPU, + '0' for GPU); for 'pretrained': 'sscor_weights' (path to a + generator checkpoint, staged as `latest_net_G_A.pth`); for + 'self-train': 'sscor_stripe_direction' ('horizontal'/'vertical'/ + 'grid'), 'sscor_h_n', 'sscor_v_n', 'sscor_grid_direction', + 'sscor_epochs' (all int). All injected by `compute()`. + """ + import shutil + import subprocess + import tempfile + + from PIL import Image + + stack = np.asarray(stack, dtype=np.float64) + + mode = opts.get('sscor_mode', 'pretrained') + gpu_ids = opts.get('sscor_gpu_ids', '-1') + patch_size = int(opts.get('sscor_patch_size', 256)) + offset_size = int(opts.get('sscor_offset_size', 100)) + repeat = int(opts.get('sscor_repeat', 1)) + dark_threshold = int(opts.get('sscor_dark_threshold', 10)) + + repo_path = os.environ.get('SSCOR_REPO_PATH', '/sscor') + restore_script = os.path.join(repo_path, 'restore.py') + + corrected = np.empty_like(stack) + + if mode == 'self-train': + stripe_direction = opts.get('sscor_stripe_direction', 'horizontal') + h_n = int(opts.get('sscor_h_n', 1)) + v_n = int(opts.get('sscor_v_n', 1)) + grid_direction = int(opts.get('sscor_grid_direction', 0)) + epochs = int(opts.get('sscor_epochs', 30)) + exp_name = 'sscor_selftrain' + + sample_stripe_script = os.path.join(repo_path, 'sample', 'sample_stripe.py') + sample_stripe_2_script = os.path.join(repo_path, 'sample', 'sample_stripe_2.py') + train_script = os.path.join(repo_path, 'train.py') + + for i in range(stack.shape[0]): + frame = stack[i] + frame_min, span, frame_uint8 = _sscor_frame_to_uint8(frame) + + # Fresh temp dirs per frame so self-trained checkpoints/samples never + # leak into the next frame's (independent) self-training run. + with tempfile.TemporaryDirectory() as tmpin, \ + tempfile.TemporaryDirectory() as tmpout, \ + tempfile.TemporaryDirectory() as tmpckpt: + + image_name = f'frame_{i:04d}.png' + input_path = os.path.join(tmpin, image_name) + Image.fromarray(frame_uint8).convert('RGB').save(input_path) + stem = os.path.splitext(image_name)[0] + + # 1. Sample stripe (trainA) / stripe-free (trainB) patches from the + # frame itself, per the upstream self-training method. + if stripe_direction == 'grid': + sample_cmd = [ + sys.executable, sample_stripe_2_script, + '--h_n', str(h_n), + '--v_n', str(v_n), + '--direction', str(grid_direction), + '--in_dir', tmpin, + '--img_name', image_name, + '--out_dir', tmpout, + '--patch_size', str(patch_size), + ] + elif stripe_direction == 'vertical': + sample_cmd = [ + sys.executable, sample_stripe_script, + '--v', + '--v_n', str(v_n), + '--in_dir', tmpin, + '--img_name', image_name, + '--out_dir', tmpout, + '--patch_size', str(patch_size), + ] + else: # 'horizontal' (default) + sample_cmd = [ + sys.executable, sample_stripe_script, + '--h', + '--h_n', str(h_n), + '--in_dir', tmpin, + '--img_name', image_name, + '--out_dir', tmpout, + '--patch_size', str(patch_size), + ] + + result = subprocess.run(sample_cmd, capture_output=True, text=True, cwd=repo_path) + if result.returncode != 0: + raise RuntimeError( + f"SSCOR self-train sampling failed on frame {i}: {result.stderr}") + + sample_dir = os.path.join(tmpout, f'sample_{stem}') + train_a_dir = os.path.join(sample_dir, 'trainA') + if not os.path.isdir(train_a_dir) or len(os.listdir(train_a_dir)) == 0: + raise RuntimeError( + f"SSCOR self-train sampling produced no training patches for frame {i} " + f"(patch_size={patch_size}, stripe_direction={stripe_direction}). The " + f"frame may be too small for this patch size, or the stripe direction/" + f"count doesn't fit the image -- try a smaller 'SSCOR patch size' or a " + f"different 'SSCOR stripe direction'/count.") + + # 2. Train a fresh per-frame CycleGAN on the self-sampled patches. + # save_epoch_freq == n_epochs guarantees model.save_networks('latest') + # (-> latest_net_G_A.pth) is written exactly once, at the final epoch. + train_cmd = [ + sys.executable, train_script, + '--dataroot', sample_dir, + '--name', exp_name, + '--model', 'sscor', + '--checkpoints_dir', tmpckpt, + '--gpu_ids', gpu_ids, + '--display_id', '0', + '--no_html', + '--load_size', str(patch_size + 30), + '--crop_size', str(patch_size), + '--n_epochs', str(epochs), + '--n_epochs_decay', '0', + '--save_epoch_freq', str(epochs), + ] + result = subprocess.run(train_cmd, capture_output=True, text=True, cwd=repo_path) + if result.returncode != 0: + raise RuntimeError( + f"SSCOR self-train train.py failed on frame {i}: {result.stderr}") + + # 3. Restore the frame with the just-trained generator. + restore_cmd = [ + sys.executable, restore_script, + '--dataroot', tmpin, + '--name', exp_name, + '--model', 'sscor', + '--image_name', image_name, + '--offset_size', str(offset_size), + '--patch_size', str(patch_size), + # load_size/crop_size must equal patch_size: restore.py's + # get_transform resizes each cropped patch to load_size and + # crops to crop_size (both default 256), and restore() writes + # the network output back into a patch_size-sized slice -- so + # leaving the defaults crashes with a shape mismatch whenever + # patch_size != 256. + '--load_size', str(patch_size), + '--crop_size', str(patch_size), + '--repeat', str(repeat), + '--dark_threshold', str(dark_threshold), + '--checkpoints_dir', tmpckpt, + '--gpu_ids', gpu_ids, + '--eval', + ] + result = subprocess.run(restore_cmd, capture_output=True, text=True, cwd=repo_path) + if result.returncode != 0: + raise RuntimeError( + f"SSCOR self-train restore.py failed on frame {i}: {result.stderr}") + + output_path = os.path.join(tmpin, 'result', f'restore-{image_name}') + if not os.path.exists(output_path): + raise RuntimeError( + f"SSCOR restore.py did not produce the expected output '{output_path}' " + f"for frame {i}. stdout: {result.stdout}") + + restored_rgb = np.asarray(Image.open(output_path).convert('RGB'), dtype=np.float64) + restored_gray = np.mean(restored_rgb, axis=-1) + corrected[i] = restored_gray / 255.0 * span + frame_min + + diagnostics = { + 'mode': 'self-train', + 'patch_size': patch_size, + 'offset_size': offset_size, + 'repeat': repeat, + 'dark_threshold': dark_threshold, + 'gpu_ids': gpu_ids, + 'stripe_direction': stripe_direction, + 'h_n': h_n, + 'v_n': v_n, + 'grid_direction': grid_direction, + 'epochs': epochs, + } + return corrected, diagnostics + + # mode == 'pretrained' + weights_path = opts['sscor_weights'] + + with tempfile.TemporaryDirectory() as tmpckpt, tempfile.TemporaryDirectory() as tmpdata: + # restore.py loads //_net_.pth (epoch + # defaults to 'latest'); at test time SSCORModel.model_names == ['G_A'], so the + # file MUST be named `latest_net_G_A.pth`, not `latest_net_G.pth`. + ckpt_dir = os.path.join(tmpckpt, 'sscor') + os.makedirs(ckpt_dir, exist_ok=True) + shutil.copy(weights_path, os.path.join(ckpt_dir, 'latest_net_G_A.pth')) + + for i in range(stack.shape[0]): + frame = stack[i] + frame_min, span, frame_uint8 = _sscor_frame_to_uint8(frame) + + image_name = f'frame_{i:04d}.png' + input_path = os.path.join(tmpdata, image_name) + Image.fromarray(frame_uint8).convert('RGB').save(input_path) + + cmd = [ + sys.executable, restore_script, + '--dataroot', tmpdata, + '--name', 'sscor', + '--model', 'sscor', + '--image_name', image_name, + '--offset_size', str(offset_size), + '--patch_size', str(patch_size), + # See note in the self-train branch: load_size/crop_size must + # equal patch_size or restore.py crashes for patch_size != 256. + '--load_size', str(patch_size), + '--crop_size', str(patch_size), + '--repeat', str(repeat), + '--dark_threshold', str(dark_threshold), + '--checkpoints_dir', tmpckpt, + '--gpu_ids', gpu_ids, + '--eval', + ] + result = subprocess.run(cmd, capture_output=True, text=True, cwd=repo_path) + if result.returncode != 0: + raise RuntimeError(f"SSCOR restore.py failed on frame {i}: {result.stderr}") + + output_path = os.path.join(tmpdata, 'result', f'restore-{image_name}') + if not os.path.exists(output_path): + raise RuntimeError( + f"SSCOR restore.py did not produce the expected output '{output_path}' " + f"for frame {i}. stdout: {result.stdout}") + + restored_rgb = np.asarray(Image.open(output_path).convert('RGB'), dtype=np.float64) + restored_gray = np.mean(restored_rgb, axis=-1) + + corrected[i] = restored_gray / 255.0 * span + frame_min + + diagnostics = { + 'mode': 'pretrained', + 'patch_size': patch_size, + 'offset_size': offset_size, + 'repeat': repeat, + 'dark_threshold': dark_threshold, + 'gpu_ids': gpu_ids, + 'checkpoint': 'SSCOR_WEIGHTS', + } + return corrected, diagnostics + + +def compute_qc_metrics(corrected_stack, opts=None): + """Lightweight EVEN-style QC stand-in (NOT the full EVEN ML evaluation and + optimization framework, Nat. Commun. 2026, which is not vendored here). + + Computes, on the corrected per-channel collection: + - cv_mean_image: coefficient of variation of a smoothed mean image + (lower = flatter/more uniform illumination). + - corner_center_ratio: residual vignetting, mean corner vs mean center + intensity of the mean image (closer to 1.0 = less vignetting). + - interframe_cv: coefficient of variation of per-frame mean intensity + across the collection. + """ + from scipy.ndimage import gaussian_filter + + stack = np.asarray(corrected_stack, dtype=np.float64) + mean_image = np.mean(stack, axis=0) + sigma = max(1.0, min(mean_image.shape[:2]) / 20.0) + smoothed = gaussian_filter(mean_image, sigma=sigma) + + mean_val = np.mean(smoothed) + std_val = np.std(smoothed) + cv_mean_image = float(std_val / mean_val) if mean_val != 0 else 0.0 + + h, w = mean_image.shape[0], mean_image.shape[1] + ch, cw = max(1, h // 8), max(1, w // 8) + center = mean_image[h // 2 - ch // 2: h // 2 + ch // 2, + w // 2 - cw // 2: w // 2 + cw // 2] + corners = [ + mean_image[:ch, :cw], mean_image[:ch, -cw:], + mean_image[-ch:, :cw], mean_image[-ch:, -cw:], + ] + corner_mean = float(np.mean([np.mean(c) for c in corners])) + center_mean = float(np.mean(center)) if center.size else 0.0 + corner_center_ratio = float(corner_mean / center_mean) if center_mean != 0 else 0.0 + + frame_means = np.mean(stack, axis=tuple(range(1, stack.ndim))) + frame_mean_of_means = np.mean(frame_means) + interframe_cv = float(np.std(frame_means) / frame_mean_of_means) if frame_mean_of_means != 0 else 0.0 + + return { + 'cv_mean_image': cv_mean_image, + 'corner_center_ratio': corner_center_ratio, + 'interframe_cv': interframe_cv, + } + + +def _method_params_for_metadata(method, opts, flat_xy, dark_xy): + """Return the subset of `opts` relevant to `method`, for provenance metadata.""" + if method == 'basic': + return { + 'estimate_darkfield': opts['estimate_darkfield'], + 'flatfield_smoothness': opts['flatfield_smoothness'], + 'darkfield_smoothness': opts['darkfield_smoothness'], + 'baseline_drift': opts['baseline_drift'], + } + if method == 'cidre': + return { + 'smoothing_sigma': opts['smoothing_sigma'], + 'dark_quantile': opts['dark_quantile'], + } + if method == 'cellprofiler': + return { + 'smoothing_sigma': opts['smoothing_sigma'], + 'cellprofiler_mode': opts['cellprofiler_mode'], + } + if method == 'flatfield': + return { + 'flat_xy': flat_xy, + 'dark_xy': dark_xy, + 'dark_constant': opts['dark_constant'], + } + if method == 'destripe': + return { + 'destripe_sigma': opts['destripe_sigma'], + 'destripe_wavelet': opts['destripe_wavelet'], + 'destripe_level': opts['destripe_level'], + } + if method == 'sscor': + params = { + 'sscor_mode': opts['sscor_mode'], + 'sscor_patch_size': opts['sscor_patch_size'], + 'sscor_offset_size': opts['sscor_offset_size'], + 'sscor_repeat': opts['sscor_repeat'], + 'sscor_dark_threshold': opts['sscor_dark_threshold'], + } + if opts['sscor_mode'] == 'self-train': + params.update({ + 'sscor_stripe_direction': opts['sscor_stripe_direction'], + 'sscor_h_n': opts['sscor_h_n'], + 'sscor_v_n': opts['sscor_v_n'], + 'sscor_grid_direction': opts['sscor_grid_direction'], + 'sscor_epochs': opts['sscor_epochs'], + }) + return params + return {} + + +_CORRECTION_FUNCTIONS = { + 'basic': 'correct_basic', + 'cidre': 'correct_cidre', + 'cellprofiler': 'correct_cellprofiler', + 'flatfield': 'correct_flatfield', + 'destripe': 'correct_destripe', + 'sscor': 'correct_sscor', +} + + +def compute(datasetId, apiUrl, token, params): + """ + params (could change): + configurationId, + datasetId, + description: tool description, + type: tool type, + id: tool id, + name: tool name, + image: docker image, + channel: annotation channel, + assignment: annotation assignment ({XY, Z, Time}), + tags: annotation tags (list of strings), + tile: tile position (TODO: roi) ({XY, Z, Time}), + connectTo: how new annotations should be connected + """ + + # Lazy import: keeps large_image off the interface/preview path; only needed during compute. + import large_image as li + + tileClient = tiles.UPennContrastDataset( + apiUrl=apiUrl, token=token, datasetId=datasetId) + + workerInterface = params['workerInterface'] + + method = workerInterface.get('Method', 'basic') + + allChannels = workerInterface.get('Channels to correct', {}) + channels = [int(k) for k, v in allChannels.items() if v] + if len(channels) == 0: + sendError('No channels to correct', + info='Select at least one channel in "Channels to correct".') + return + + if 'frames' not in tileClient.tiles: + sendError('Only one image; exiting', + info='Illumination correction requires a multi-frame dataset ' + '(a `frames` list) to build a per-channel image collection.') + return + + frames = tileClient.tiles['frames'] + + opts = { + 'estimate_darkfield': bool(workerInterface.get('Estimate darkfield', True)), + 'flatfield_smoothness': float(workerInterface.get('Flatfield smoothness', 1.0)), + 'darkfield_smoothness': float(workerInterface.get('Darkfield smoothness', 1.0)), + 'baseline_drift': bool(workerInterface.get('Correct timelapse baseline drift', False)), + 'smoothing_sigma': float(workerInterface.get('Smoothing sigma', 50)), + 'dark_quantile': float(workerInterface.get('Dark quantile', 0.02)), + 'cellprofiler_mode': workerInterface.get('CellProfiler mode', 'regular'), + 'dark_constant': float(workerInterface.get('Dark-field constant', 0)), + 'destripe_sigma': float(workerInterface.get('Destripe sigma', 128)), + 'destripe_wavelet': workerInterface.get('Destripe wavelet', 'db3'), + 'destripe_level': int(workerInterface.get('Destripe level', 0)), + 'sscor_patch_size': int(workerInterface.get('SSCOR patch size', 256)), + 'sscor_offset_size': int(workerInterface.get('SSCOR offset size', 100)), + 'sscor_repeat': int(workerInterface.get('SSCOR repeat', 1)), + 'sscor_dark_threshold': int(workerInterface.get('SSCOR dark threshold', 10)), + 'sscor_mode': workerInterface.get('SSCOR mode', 'pretrained'), + 'sscor_stripe_direction': workerInterface.get('SSCOR stripe direction', 'horizontal'), + 'sscor_h_n': int(workerInterface.get('SSCOR horizontal stripe count', 1)), + 'sscor_v_n': int(workerInterface.get('SSCOR vertical stripe count', 1)), + 'sscor_grid_direction': int(workerInterface.get('SSCOR grid direction', 0)), + 'sscor_epochs': int(workerInterface.get('SSCOR training epochs', 30)), + } + + qc_enabled = bool(workerInterface.get('Report correction quality (QC)', False)) + + flat_xy_str = str(workerInterface.get('Flat-field XY coordinate', '') or '').strip() + dark_xy_str = str(workerInterface.get('Dark-field XY coordinate', '') or '').strip() + + if method == 'flatfield' and flat_xy_str == '': + sendError('Flat-field reference required', + info='The flatfield method requires a "Flat-field XY coordinate". Enter the ' + '1-indexed XY position of a flat-field calibration image, or choose a ' + 'different Method.') + return + + flat_xy = int(flat_xy_str) - 1 if flat_xy_str != '' else None + dark_xy = int(dark_xy_str) - 1 if dark_xy_str != '' else None + + sscor_weights_path = None + sscor_gpu_ids = '-1' + if method == 'sscor': + if opts['sscor_mode'] == 'pretrained': + sscor_weights_path = resolve_sscor_checkpoint() + if sscor_weights_path is None: + return # resolve_sscor_checkpoint() already called sendError + else: + # self-train needs no pre-supplied checkpoint -- it trains its own, + # per frame, from patches sampled out of the image itself. + sendWarning( + 'SSCOR self-training enabled', + info='SSCOR "self-train" mode needs no pre-supplied checkpoint, but it trains a ' + 'fresh CycleGAN model PER FRAME (sampling stripe/clean patches from the ' + 'image itself, then training, then restoring); this can be slow, ' + 'especially without a GPU.') + + try: + import torch + cuda_available = torch.cuda.is_available() + except Exception: + cuda_available = False + + if cuda_available: + sscor_gpu_ids = '0' + else: + sscor_gpu_ids = '-1' + sendWarning('Running SSCOR on CPU', + info='No CUDA GPU detected; SSCOR deep-learning stripe correction on ' + 'CPU is very slow. A GPU is strongly recommended.') + + # Group frame indices by channel + channel_frame_indices = {c: [] for c in channels} + for i, frame in enumerate(frames): + c = frame.get('IndexC', 0) + if c in channel_frame_indices: + channel_frame_indices[c].append(i) + + correction_fn_name = _CORRECTION_FUNCTIONS.get(method) + if correction_fn_name is None: + sendError(f'Unknown method: {method}') + return + + source_dtype = tileClient.tiles.get('dtype', None) + + corrected_by_frame = {} + diagnostics_by_channel = {} + qc_by_channel = {} + + total_channels = len(channels) + for ci, channel in enumerate(channels): + indices = channel_frame_indices.get(channel, []) + if len(indices) == 0: + continue + + if len(indices) < 3 and method in ('basic', 'cidre', 'cellprofiler'): + sendWarning('Small image collection', + info=f'Channel {channel} has only {len(indices)} frame(s); retrospective ' + 'illumination estimation may be unreliable with fewer than 3 images.') + + stack = np.stack( + [tileClient.getRegion(datasetId, frame=i).squeeze() for i in indices], axis=0) + + method_opts = dict(opts) + + if method == 'basic': + method_opts['time_order'] = [frames[i].get('IndexT', 0) for i in indices] + + if method == 'flatfield': + flat_frame = tileClient.coordinatesToFrameIndex(flat_xy, 0, 0, channel) + flat = tileClient.getRegion(datasetId, frame=flat_frame).squeeze().astype(np.float64) + if dark_xy is not None: + dark_frame = tileClient.coordinatesToFrameIndex(dark_xy, 0, 0, channel) + dark = tileClient.getRegion(datasetId, frame=dark_frame).squeeze().astype(np.float64) + else: + dark = np.full_like(flat, method_opts['dark_constant'], dtype=np.float64) + method_opts['flat'] = flat + method_opts['dark'] = dark + + if method == 'sscor': + method_opts['sscor_weights'] = sscor_weights_path + method_opts['sscor_gpu_ids'] = sscor_gpu_ids + + correction_fn = globals()[correction_fn_name] + try: + corrected, diag = correction_fn(stack, method_opts) + except Exception as e: + # Surface library/subprocess failures (e.g. BaSiC.fit on a + # degenerate stack, or a SSCOR sample/train/restore subprocess + # error) as a structured error rather than a raw traceback -- + # see FLUOR_CORRECTION_WORKERS_SPEC.md design principle 5. + sendError(f"Illumination correction failed for channel {channel} using method '{method}'", + info=str(e)) + return + corrected = np.nan_to_num(np.asarray(corrected, dtype=np.float64)) + + diagnostics_by_channel[channel] = diag + + if qc_enabled: + qc_by_channel[channel] = compute_qc_metrics(corrected, method_opts) + + if source_dtype is not None: + dtype = np.dtype(source_dtype) + if np.issubdtype(dtype, np.integer): + info = np.iinfo(dtype) + corrected = np.clip(corrected, info.min, info.max).astype(dtype) + else: + corrected = corrected.astype(dtype) + + for pos, frame_idx in enumerate(indices): + corrected_by_frame[frame_idx] = corrected[pos] + + sendProgress((ci + 1) / total_channels, 'Illumination correction', + f"Corrected channel {channel} ({ci + 1}/{total_channels})") + + gc = tileClient.client + + sink = li.new() + + for i, frame in enumerate(frames): + # Create a parameters dictionary with only the indices that exist in frame + # The len(k) > 5 is to avoid the 'Index' key that has no postfix to it + large_image_params = {f'{k.lower()[5:]}': v for k, v in frame.items( + ) if k.startswith('Index') and len(k) > 5} + + if i in corrected_by_frame: + image = corrected_by_frame[i] + else: + image = tileClient.getRegion(datasetId, frame=i).squeeze() + + sink.addTile(image, 0, 0, **large_image_params) + + sendProgress(i / len(frames), 'Illumination correction', + f"Writing frame {i + 1}/{len(frames)}") + + # Copy over the metadata + if 'channels' in tileClient.tiles: + sink.channelNames = tileClient.tiles['channels'] + + sink.mm_x = tileClient.tiles['mm_x'] + sink.mm_y = tileClient.tiles['mm_y'] + sink.magnification = tileClient.tiles['magnification'] + sink.write('/tmp/illumination_corrected.tiff') + print("Wrote to file") + + item = gc.uploadFileToFolder(datasetId, '/tmp/illumination_corrected.tiff') + + metadata = { + 'tool': 'Illumination Correction', + 'method': method, + 'channels': channels, + } + metadata.update(_method_params_for_metadata(method, opts, flat_xy, dark_xy)) + if diagnostics_by_channel: + metadata['diagnostics'] = diagnostics_by_channel + if qc_enabled and qc_by_channel: + metadata['qc'] = qc_by_channel + + gc.addMetadataToItem(item['itemId'], metadata) + print("Uploaded file") + + +if __name__ == '__main__': + # Define the command-line interface for the entry point + parser = argparse.ArgumentParser( + description='Correct illumination/shading/vignetting/striping in images') + + 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:]) + + params = json.loads(args.parameters) + datasetId = args.datasetId + apiUrl = args.apiUrl + token = args.token + + match args.request: + case 'compute': + compute(datasetId, apiUrl, token, params) + case 'interface': + interface(params['image'], apiUrl, token) diff --git a/workers/annotations/illumination_correction/environment.yml b/workers/annotations/illumination_correction/environment.yml new file mode 100644 index 0000000..42e8751 --- /dev/null +++ b/workers/annotations/illumination_correction/environment.yml @@ -0,0 +1,18 @@ +name: worker +channels: +- conda-forge +- defaults +dependencies: +- python=3.11 +- pip +- numpy>=2.0 +- scipy +- scikit-image +- shapely>=2.0.6 +- libtiff +- openslide +- libvips +- tifffile +- pip: + - basicpy + - pystripe diff --git a/workers/annotations/illumination_correction/tests/Dockerfile_Test b/workers/annotations/illumination_correction/tests/Dockerfile_Test new file mode 100644 index 0000000..bf80ee4 --- /dev/null +++ b/workers/annotations/illumination_correction/tests/Dockerfile_Test @@ -0,0 +1,14 @@ +# ./workers/annotations/illumination_correction/tests/Dockerfile_Test +FROM annotations/illumination_correction:latest AS test + +# Install test dependencies +SHELL ["conda", "run", "-n", "worker", "/bin/bash", "-c"] +RUN pip install pytest pytest-mock + +# Copy test files +RUN mkdir -p /tests +COPY ./workers/annotations/illumination_correction/tests/*.py /tests +WORKDIR /tests + +# Command to run tests +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..3e69f38 --- /dev/null +++ b/workers/annotations/illumination_correction/tests/__init__.py @@ -0,0 +1,3 @@ +import sys +import os +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) diff --git a/workers/annotations/illumination_correction/tests/test_illumination_correction.py b/workers/annotations/illumination_correction/tests/test_illumination_correction.py new file mode 100644 index 0000000..3ac3546 --- /dev/null +++ b/workers/annotations/illumination_correction/tests/test_illumination_correction.py @@ -0,0 +1,442 @@ +import pytest +from unittest.mock import patch, MagicMock +import numpy as np + +# Import your worker module +from entrypoint import compute, interface + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +def _get_region_side_effect(*args, **kwargs): + """Deterministic fake image per frame index: value = (frame_index + 1) * 10.""" + frame = kwargs.get('frame') + if frame is None and len(args) > 1: + frame = args[1] + if frame is None: + frame = 0 + return np.full((16, 16), (frame + 1) * 10, dtype=np.uint16) + + +@pytest.fixture +def mock_tile_client(): + """Mock the tiles.UPennContrastDataset""" + with patch('annotation_client.tiles.UPennContrastDataset') as mock_client: + client = mock_client.return_value + # 6 frames: 2 channels x 2 XY x (mostly) 1 Z, with some time variation. + client.tiles = { + 'frames': [ + {'IndexXY': 0, 'IndexZ': 0, 'IndexT': 0, 'IndexC': 0}, # 0 + {'IndexXY': 0, 'IndexZ': 0, 'IndexT': 0, 'IndexC': 1}, # 1 + {'IndexXY': 0, 'IndexZ': 0, 'IndexT': 1, 'IndexC': 0}, # 2 + {'IndexXY': 0, 'IndexZ': 0, 'IndexT': 1, 'IndexC': 1}, # 3 + {'IndexXY': 1, 'IndexZ': 0, 'IndexT': 0, 'IndexC': 0}, # 4 + {'IndexXY': 1, 'IndexZ': 0, 'IndexT': 0, 'IndexC': 1}, # 5 + ], + 'IndexRange': { + 'IndexXY': 2, + 'IndexZ': 1, + 'IndexT': 2, + 'IndexC': 2, + }, + 'channels': ['DAPI', 'FITC'], + 'mm_x': 0.65, + 'mm_y': 0.65, + 'magnification': 20, + 'dtype': np.uint16, + } + + client.getRegion.side_effect = _get_region_side_effect + client.coordinatesToFrameIndex.return_value = 0 + + mock_gc = MagicMock() + mock_gc.uploadFileToFolder.return_value = {'itemId': 'test_item_id'} + client.client = mock_gc + + yield client + + +@pytest.fixture +def mock_worker_preview_client(): + """Mock the UPennContrastWorkerPreviewClient""" + with patch('annotation_client.workers.UPennContrastWorkerPreviewClient') as mock_client: + yield mock_client.return_value + + +@pytest.fixture +def mock_large_image(): + """Mock large_image operations""" + with patch('large_image.new') as mock_li_new: + mock_sink = MagicMock() + mock_li_new.return_value = mock_sink + yield mock_sink + + +@pytest.fixture +def mock_corrections(mocker): + """Patch every per-algorithm function with an identity pass-through mock.""" + identity = lambda stack, opts: (np.asarray(stack, dtype=np.float64), {'note': 'mock'}) + return { + 'basic': mocker.patch('entrypoint.correct_basic', side_effect=identity), + 'cidre': mocker.patch('entrypoint.correct_cidre', side_effect=identity), + 'cellprofiler': mocker.patch('entrypoint.correct_cellprofiler', side_effect=identity), + 'flatfield': mocker.patch('entrypoint.correct_flatfield', side_effect=identity), + 'destripe': mocker.patch('entrypoint.correct_destripe', side_effect=identity), + 'sscor': mocker.patch('entrypoint.correct_sscor', side_effect=identity), + } + + +def base_worker_interface(**overrides): + interface_dict = { + 'Method': 'basic', + 'Channels to correct': {'0': True, '1': False}, + 'Estimate darkfield': True, + 'Flatfield smoothness': 1.0, + 'Darkfield smoothness': 1.0, + 'Correct timelapse baseline drift': False, + 'Smoothing sigma': 50, + 'Dark quantile': 0.02, + 'CellProfiler mode': 'regular', + 'Flat-field XY coordinate': '1', + 'Dark-field XY coordinate': '', + 'Dark-field constant': 0, + 'Destripe sigma': 128, + 'Destripe wavelet': 'db3', + 'Destripe level': 0, + 'SSCOR mode': 'pretrained', + 'SSCOR stripe direction': 'horizontal', + 'SSCOR horizontal stripe count': 1, + 'SSCOR vertical stripe count': 1, + 'SSCOR grid direction': 0, + 'SSCOR training epochs': 30, + 'SSCOR patch size': 256, + 'SSCOR offset size': 100, + 'SSCOR repeat': 1, + 'SSCOR dark threshold': 10, + 'Report correction quality (QC)': False, + } + interface_dict.update(overrides) + return {'workerInterface': interface_dict} + + +# --------------------------------------------------------------------------- +# 1. interface() +# --------------------------------------------------------------------------- + +def test_interface(mock_worker_preview_client): + interface('test_image', 'http://test-api', 'test-token') + + mock_worker_preview_client.setWorkerImageInterface.assert_called_once() + + call_args = mock_worker_preview_client.setWorkerImageInterface.call_args + image_arg = call_args[0][0] + interface_data = call_args[0][1] + + assert image_arg == 'test_image' + + # Method select + assert 'Method' in interface_data + method_iface = interface_data['Method'] + assert method_iface['type'] == 'select' + assert method_iface['items'] == [ + 'basic', 'cidre', 'cellprofiler', 'flatfield', 'destripe', 'sscor'] + assert method_iface['default'] == 'basic' + + # Channel field + assert interface_data['Channels to correct']['type'] == 'channelCheckboxes' + + # Key params for each method are present + for key in [ + 'Estimate darkfield', 'Flatfield smoothness', 'Darkfield smoothness', + 'Correct timelapse baseline drift', 'Smoothing sigma', 'Dark quantile', + 'CellProfiler mode', 'Flat-field XY coordinate', 'Dark-field XY coordinate', + 'Dark-field constant', 'Destripe sigma', 'Destripe wavelet', 'Destripe level', + 'SSCOR mode', 'SSCOR stripe direction', 'SSCOR horizontal stripe count', + 'SSCOR vertical stripe count', 'SSCOR grid direction', 'SSCOR training epochs', + 'SSCOR patch size', 'SSCOR offset size', 'SSCOR repeat', 'SSCOR dark threshold', + 'Report correction quality (QC)', + ]: + assert key in interface_data, f'{key} missing from interface' + + assert interface_data['CellProfiler mode']['items'] == ['regular', 'background'] + assert interface_data['Destripe wavelet']['items'] == ['db3', 'db5', 'haar', 'sym4'] + assert interface_data['Estimate darkfield']['default'] is True + assert interface_data['Report correction quality (QC)']['default'] is False + + assert interface_data['SSCOR mode']['type'] == 'select' + assert interface_data['SSCOR mode']['items'] == ['pretrained', 'self-train'] + assert interface_data['SSCOR mode']['default'] == 'pretrained' + assert interface_data['SSCOR stripe direction']['items'] == ['horizontal', 'vertical', 'grid'] + assert interface_data['SSCOR stripe direction']['default'] == 'horizontal' + assert interface_data['SSCOR horizontal stripe count']['default'] == 1 + assert interface_data['SSCOR vertical stripe count']['default'] == 1 + assert interface_data['SSCOR grid direction']['default'] == 0 + assert interface_data['SSCOR training epochs']['default'] == 30 + + +# --------------------------------------------------------------------------- +# 2. Dispatch +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize('method', ['basic', 'cidre', 'cellprofiler', 'flatfield', 'destripe']) +def test_dispatch_calls_only_selected_method( + method, mock_tile_client, mock_large_image, mock_corrections): + params = base_worker_interface(Method=method) + + compute('test_dataset', 'http://test-api', 'test-token', params) + + for name, mock_fn in mock_corrections.items(): + if name == method: + assert mock_fn.call_count == 1, f'{name} should have been called once' + else: + assert mock_fn.call_count == 0, f'{name} should not have been called' + + +# --------------------------------------------------------------------------- +# 3. Channel filtering +# --------------------------------------------------------------------------- + +def test_unselected_channels_pass_through_unchanged( + mock_tile_client, mock_large_image, mock_corrections): + offset_correction = lambda stack, opts: (np.asarray(stack, dtype=np.float64) + 1000, {}) + mock_corrections['basic'].side_effect = offset_correction + + params = base_worker_interface(Method='basic', **{'Channels to correct': {'0': True, '1': False}}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + # Only channel 0's collection (3 frames) should have gone through correct_basic + assert mock_corrections['basic'].call_count == 1 + stack_arg = mock_corrections['basic'].call_args[0][0] + assert stack_arg.shape[0] == 3 + + add_tile_calls = mock_large_image.addTile.call_args_list + for c in add_tile_calls: + image = c.args[0] + channel = c.kwargs.get('c') + frame_idx = _frame_index_from_kwargs(mock_tile_client, c.kwargs) + expected_raw = (frame_idx + 1) * 10 + if channel == 0: + # Processed: raw + 1000 (clipped/cast back to uint16) + assert int(image.flat[0]) == expected_raw + 1000 + else: + # Untouched: exactly the raw mock image + assert int(image.flat[0]) == expected_raw + + +def _frame_index_from_kwargs(tile_client, kwargs): + """Recover the original frame index (0..5) from the addTile xy/z/t/c kwargs.""" + for i, frame in enumerate(tile_client.tiles['frames']): + if (frame['IndexXY'] == kwargs.get('xy') and frame['IndexZ'] == kwargs.get('z') + and frame['IndexT'] == kwargs.get('t') and frame['IndexC'] == kwargs.get('c')): + return i + raise AssertionError(f'Could not match frame for kwargs {kwargs}') + + +# --------------------------------------------------------------------------- +# 4. Output plumbing +# --------------------------------------------------------------------------- + +def test_output_plumbing(mock_tile_client, mock_large_image, mock_corrections): + params = base_worker_interface(Method='basic') + compute('test_dataset', 'http://test-api', 'test-token', params) + + mock_large_image.write.assert_called_once_with('/tmp/illumination_corrected.tiff') + mock_tile_client.client.uploadFileToFolder.assert_called_once_with( + 'test_dataset', '/tmp/illumination_corrected.tiff') + + mock_tile_client.client.addMetadataToItem.assert_called_once() + call_args = mock_tile_client.client.addMetadataToItem.call_args + assert call_args[0][0] == 'test_item_id' + metadata = call_args[0][1] + assert metadata['tool'] == 'Illumination Correction' + assert metadata['method'] == 'basic' + assert metadata['channels'] == [0] + + +# --------------------------------------------------------------------------- +# 5. Metadata preservation +# --------------------------------------------------------------------------- + +def test_metadata_preservation(mock_tile_client, mock_large_image, mock_corrections): + params = base_worker_interface(Method='basic') + compute('test_dataset', 'http://test-api', 'test-token', params) + + assert mock_large_image.channelNames == ['DAPI', 'FITC'] + assert mock_large_image.mm_x == 0.65 + assert mock_large_image.mm_y == 0.65 + assert mock_large_image.magnification == 20 + + +# --------------------------------------------------------------------------- +# 6. Error paths +# --------------------------------------------------------------------------- + +def test_no_channels_selected_error(mock_tile_client, mock_large_image, mock_corrections, capsys): + params = base_worker_interface(**{'Channels to correct': {'0': False, '1': False}}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + captured = capsys.readouterr() + assert '"error": "No channels to correct"' in captured.out + assert '"type": "error"' in captured.out + mock_large_image.write.assert_not_called() + + +def test_no_frames_error(mock_tile_client, mock_large_image, mock_corrections, capsys): + del mock_tile_client.tiles['frames'] + + params = base_worker_interface() + compute('test_dataset', 'http://test-api', 'test-token', params) + + captured = capsys.readouterr() + assert '"error": "Only one image; exiting"' in captured.out + assert '"type": "error"' in captured.out + mock_large_image.write.assert_not_called() + + +def test_flatfield_without_flat_reference_error( + mock_tile_client, mock_large_image, mock_corrections, capsys): + params = base_worker_interface(Method='flatfield', **{'Flat-field XY coordinate': ''}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + captured = capsys.readouterr() + assert '"error": "Flat-field reference required"' in captured.out + mock_corrections['flatfield'].assert_not_called() + mock_large_image.write.assert_not_called() + + +def test_sscor_dispatch_with_weights( + mock_tile_client, mock_large_image, mock_corrections, mocker): + """With SSCOR_WEIGHTS resolvable, sscor dispatches to correct_sscor and writes output.""" + mocker.patch('entrypoint.resolve_sscor_checkpoint', + return_value='/fake/weights/latest_net_G.pth') + + params = base_worker_interface(Method='sscor') + compute('test_dataset', 'http://test-api', 'test-token', params) + + mock_corrections['sscor'].assert_called_once() + call_opts = mock_corrections['sscor'].call_args[0][1] + assert call_opts['sscor_weights'] == '/fake/weights/latest_net_G.pth' + assert call_opts['sscor_gpu_ids'] in ('0', '-1') + assert call_opts['sscor_patch_size'] == 256 + assert call_opts['sscor_offset_size'] == 100 + assert call_opts['sscor_repeat'] == 1 + assert call_opts['sscor_dark_threshold'] == 10 + + mock_large_image.write.assert_called_once_with('/tmp/illumination_corrected.tiff') + mock_tile_client.client.uploadFileToFolder.assert_called_once() + + +def test_sscor_without_weights_error( + mock_tile_client, mock_large_image, mock_corrections, mocker, capsys): + """With no SSCOR_WEIGHTS (simulated via resolve_sscor_checkpoint -> None, which already + sent the actionable sendError), compute() must bail out before the channel loop: no + correct_sscor call, no write, no upload.""" + mocker.patch('entrypoint.resolve_sscor_checkpoint', return_value=None) + + params = base_worker_interface(Method='sscor') + compute('test_dataset', 'http://test-api', 'test-token', params) + + mock_corrections['sscor'].assert_not_called() + mock_large_image.write.assert_not_called() + mock_tile_client.client.uploadFileToFolder.assert_not_called() + + +def test_sscor_selftrain_no_weights_needed( + mock_tile_client, mock_large_image, mock_corrections, mocker, capsys): + """'SSCOR mode'='self-train' must dispatch to correct_sscor with no SSCOR_WEIGHTS at all + and WITHOUT resolve_sscor_checkpoint being patched -- self-train trains its own model per + frame instead of requiring a pre-supplied checkpoint.""" + mocker.patch.dict('os.environ', {'SSCOR_WEIGHTS': ''}) + + resolve_spy = mocker.patch('entrypoint.resolve_sscor_checkpoint') + + params = base_worker_interface(Method='sscor', **{'SSCOR mode': 'self-train'}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + captured = capsys.readouterr() + assert '"error"' not in captured.out + resolve_spy.assert_not_called() + + mock_corrections['sscor'].assert_called_once() + call_opts = mock_corrections['sscor'].call_args[0][1] + assert call_opts['sscor_mode'] == 'self-train' + assert call_opts['sscor_weights'] is None + assert call_opts['sscor_gpu_ids'] in ('0', '-1') + assert call_opts['sscor_stripe_direction'] == 'horizontal' + assert call_opts['sscor_h_n'] == 1 + assert call_opts['sscor_v_n'] == 1 + assert call_opts['sscor_grid_direction'] == 0 + assert call_opts['sscor_epochs'] == 30 + + mock_large_image.write.assert_called_once_with('/tmp/illumination_corrected.tiff') + mock_tile_client.client.uploadFileToFolder.assert_called_once() + + +# --------------------------------------------------------------------------- +# 7. Progress reporting +# --------------------------------------------------------------------------- + +def test_progress_reporting(mock_tile_client, mock_large_image, mock_corrections, capsys): + params = base_worker_interface(Method='basic') + compute('test_dataset', 'http://test-api', 'test-token', params) + + captured = capsys.readouterr() + assert '"progress":' in captured.out + assert 'Illumination correction' in captured.out + + +# --------------------------------------------------------------------------- +# Extra coverage: dtype preservation, QC metrics, small-collection warning +# --------------------------------------------------------------------------- + +def test_dtype_preserved_on_output(mock_tile_client, mock_large_image, mock_corrections): + params = base_worker_interface(Method='basic') + compute('test_dataset', 'http://test-api', 'test-token', params) + + for c in mock_large_image.addTile.call_args_list: + image = c.args[0] + assert image.dtype == np.uint16 + + +def test_qc_metrics_added_to_metadata(mock_tile_client, mock_large_image, mock_corrections): + params = base_worker_interface(Method='basic', **{'Report correction quality (QC)': True}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + metadata = mock_tile_client.client.addMetadataToItem.call_args[0][1] + assert 'qc' in metadata + assert 0 in metadata['qc'] + qc = metadata['qc'][0] + assert 'cv_mean_image' in qc + assert 'corner_center_ratio' in qc + assert 'interframe_cv' in qc + + +def test_small_collection_warning(mock_tile_client, mock_large_image, mock_corrections, capsys): + # Reduce the dataset so channel 0 only has a single frame. + mock_tile_client.tiles['frames'] = [ + {'IndexXY': 0, 'IndexZ': 0, 'IndexT': 0, 'IndexC': 0}, + {'IndexXY': 0, 'IndexZ': 0, 'IndexT': 0, 'IndexC': 1}, + ] + + params = base_worker_interface(Method='basic') + compute('test_dataset', 'http://test-api', 'test-token', params) + + captured = capsys.readouterr() + assert '"warning": "Small image collection"' in captured.out + + +def test_flatfield_dispatch_uses_reference_frames( + mock_tile_client, mock_large_image, mock_corrections): + params = base_worker_interface( + Method='flatfield', + **{'Flat-field XY coordinate': '1', 'Dark-field XY coordinate': ''}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + mock_corrections['flatfield'].assert_called_once() + call_opts = mock_corrections['flatfield'].call_args[0][1] + assert 'flat' in call_opts + assert 'dark' in call_opts + assert call_opts['flat'].shape == (16, 16) + assert call_opts['dark'].shape == (16, 16) diff --git a/workers/annotations/image_restoration/Dockerfile b/workers/annotations/image_restoration/Dockerfile new file mode 100644 index 0000000..f743cc4 --- /dev/null +++ b/workers/annotations/image_restoration/Dockerfile @@ -0,0 +1,134 @@ +FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04 AS base +LABEL isUPennContrastWorker=True +LABEL com.nvidia.volumes.needed="nvidia_driver" + +ENV NVIDIA_VISIBLE_DEVICES all +ENV NVIDIA_DRIVER_CAPABILITIES compute,utility + +RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -qy tzdata && \ + apt-get install -qy software-properties-common python3-software-properties && \ + apt-get update && apt-get install -qy \ + build-essential \ + wget \ + python3 \ + libffi-dev \ + libssl-dev \ + libjpeg-dev \ + zlib1g-dev \ + git \ + libpython3-dev && \ + apt-get clean && rm -rf /var/lib/apt/lists/* + +ENV PATH="/root/miniforge3/bin:$PATH" +ARG PATH="/root/miniforge3/bin:$PATH" + +RUN wget \ + https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh \ + && mkdir /root/.conda \ + && bash Miniforge3-Linux-x86_64.sh -b \ + && rm -f Miniforge3-Linux-x86_64.sh + +FROM base AS build + +# Create conda environment (numpy/scipy/scikit-image/tifffile/large_image system deps) +COPY ./workers/annotations/image_restoration/environment.yml / +RUN conda env create --file /environment.yml +SHELL ["conda", "run", "-n", "worker", "/bin/bash", "-c"] + +# Install NimbusImage annotation client +RUN git clone https://github.com/arjunrajlaboratory/NimbusImage/ +RUN pip install -r /NimbusImage/devops/girder/annotation_client/requirements.txt +RUN pip install -e /NimbusImage/devops/girder/annotation_client/ + +# Install annotation_utilities +COPY ./annotation_utilities /annotation_utilities +RUN pip install /annotation_utilities + +# Install large_image for image processing (zarr source supports multi-dimensional images) +RUN pip install large-image-source-zarr large-image-source-tiff large-image-converter large-image --find-links https://girder.github.io/large_image_wheels + +# PyTorch (CUDA 12.1 wheels), then the restoration-method dependencies. torch is +# intentionally installed via pip (not conda) to match the CUDA base image's +# runtime exactly. +RUN pip install torch --index-url https://download.pytorch.org/whl/cu121 +RUN pip install careamics "cellpose>=3,<4" triton tifffile +# FluoResFM inference deps (see entrypoint.py restore_fluoresfm + the vendored +# repo's 3_0_test_it2i.py / 3_1_test_i2i.py imports): +# open_clip_torch + transformers -> BiomedCLIP text embedder (unet_sd_c backbone) +# huggingface_hub -> optional embedder download (1_1_download_embedder.py) +# torchinfo -> imported by models.care / models.dfcan / models.unifmir +# timm -> imported by models.unifmir +# flash-attn is intentionally NOT installed: the vendored CrossAttention falls +# back to plain attention when flash_attn is absent, and only uses it for +# self-attention (never for the text cross-attention this worker relies on). +RUN pip install open_clip_torch transformers huggingface_hub torchinfo timm + +# --- ZS-DeconvNet (zero-shot denoise + deconvolution) --- +# Vendored for provenance/reference and potential future use of its PSF/data +# tooling. entrypoint.py's restore_zs_deconvnet() implements a self-contained +# reimplementation of the published zero-shot algorithm rather than driving +# this repo's shell-script CLI directly -- see IMAGE_RESTORATION.md for why. +# Pinned to tag v1.0 (commit 04d2c21), the only tagged release as of this +# writing. +RUN git clone https://github.com/TristaZeng/ZS-DeconvNet.git /zs_deconvnet && \ + cd /zs_deconvnet && git checkout 04d2c21 + +# --- FluoResFM (pretrained, text-conditioned foundation model) --- +# Vendored from the canonical author repo (Qiqi Lu; see also the +# napari-fluoresfm plugin at https://github.com/qiqi-lu/napari-fluoresfm and +# the related UNiFMIR work at https://github.com/cxm12/UNiFMIR). Pinned to +# tag v1.0.1. Weights are NOT bundled here (see download_models.py) -- they +# are distributed via Google Drive/Baidu Yun and must be supplied at runtime +# via the FLUORESFM_WEIGHTS environment variable / volume mount. +# +# Real inference API (see entrypoint.py restore_fluoresfm): +# * text-conditioned FluoResFM = models.unet_sd_c.UNetModel conditioned on a +# BiomedCLIP text embedding (models.biomedclip_embedder.BiomedCLIPTextEmbedder), +# per 3_0_test_it2i.py. The checkpoint + a small example model come from the +# repo README's download links (Google Drive / Baidu Yun). +# * The BiomedCLIP embedder weights (open_clip_config.json + +# open_clip_pytorch_model.bin) come from HuggingFace +# 'microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224' +# (repo's 1_1_download_embedder.py) and are supplied at runtime via +# FLUORESFM_EMBEDDER_DIR (not baked in -- the .bin is ~0.4 GB). +RUN git clone https://github.com/qiqi-lu/fluoresfm.git /fluoresfm && \ + cd /fluoresfm && git checkout v1.0.1 +ENV FLUORESFM_REPO_PATH=/fluoresfm +# Default runtime weights location; override by mounting a checkpoint here or +# by overriding FLUORESFM_WEIGHTS entirely. Left unset/missing by default -- +# restore_fluoresfm() sendError()s with instructions if no file exists here. +ENV FLUORESFM_WEIGHTS=/opt/fluoresfm_weights/fluoresfm.pt +# Directory holding the BiomedCLIP text-embedder files (open_clip_config.json + +# open_clip_pytorch_model.bin) for the text-conditioned 'unet_sd_c' backbone. +# Not populated at build time; mount/download the two files here at runtime. +ENV FLUORESFM_EMBEDDER_DIR=/opt/fluoresfm_weights/biomedclip + +# models_config.py is the single source of truth for the built-in Cellpose3 +# restoration checkpoints, so it must be present before download_models.py runs. +COPY ./workers/annotations/image_restoration/models_config.py / +COPY ./workers/annotations/image_restoration/download_models.py / +# FLUORESFM_WEIGHTS_URL is an optional build ARG for internal mirrors; if +# unset, the FluoResFM weight download is skipped without failing the build. +ARG FLUORESFM_WEIGHTS_URL="" +ENV FLUORESFM_WEIGHTS_URL=${FLUORESFM_WEIGHTS_URL} +RUN python /download_models.py + +# Copy entrypoint +COPY ./workers/annotations/image_restoration/entrypoint.py / + +WORKDIR / + +LABEL isUPennContrastWorker="" \ + isAnnotationWorker="" \ + interfaceName="Image Restoration" \ + interfaceCategory="Image Processing" \ + description="Restores images (Noise2Void/N2V2, Cellpose3, ZS-DeconvNet, FluoResFM)" \ + hasPreview="False" \ + advancedOptionsPanel="False" \ + annotationConfigurationPanel="False" \ + defaultToolName="Image Restoration" + +# Fast env activation in place of `conda run` (see todo/worker-startup-latency.md) +COPY ./workers/base_docker_images/run_worker.sh /usr/local/bin/run_worker.sh +RUN chmod +x /usr/local/bin/run_worker.sh +ENTRYPOINT ["/usr/local/bin/run_worker.sh", "/entrypoint.py"] diff --git a/workers/annotations/image_restoration/Dockerfile_M1 b/workers/annotations/image_restoration/Dockerfile_M1 new file mode 100644 index 0000000..d0e913b --- /dev/null +++ b/workers/annotations/image_restoration/Dockerfile_M1 @@ -0,0 +1,61 @@ +# CPU-only build for Mac (arm64) development. `nimbusimage/image-processing-base` +# already provides the multi-arch "worker" conda env with annotation_client, +# annotation_utilities, worker_client, and large_image pre-installed (see +# workers/base_docker_images/Dockerfile.image_processing_worker_base) -- we +# only need to add this worker's own (heavier, ML-specific) dependencies on +# top of it, using CPU-only PyTorch instead of the CUDA build used in the +# production Dockerfile. +FROM nimbusimage/image-processing-base:latest + +SHELL ["conda", "run", "-n", "worker", "/bin/bash", "-c"] + +# CPU-only torch, then the restoration-method dependencies. +RUN pip install torch --index-url https://download.pytorch.org/whl/cpu +RUN pip install careamics "cellpose>=3,<4" tifffile +# triton (only needed for fluoresfm) is CUDA-oriented and may not have wheels +# for all CPU/arm64 build hosts; don't fail the build if it's unavailable -- +# fluoresfm will simply be unusable on such a build until weights + a +# compatible torch/triton are supplied. +RUN pip install triton || echo "triton unavailable on this platform; FluoResFM inference will be disabled in this build" +# FluoResFM inference deps (BiomedCLIP text embedder + baseline-backbone imports; +# see entrypoint.py restore_fluoresfm). flash-attn is intentionally omitted -- +# the vendored attention falls back to plain attention without it. +RUN pip install open_clip_torch transformers huggingface_hub torchinfo timm + +# --- ZS-DeconvNet (vendored for provenance; see main Dockerfile and +# IMAGE_RESTORATION.md for why entrypoint.py implements its own zero-shot +# reimplementation rather than driving this repo's CLI directly) --- +RUN git clone https://github.com/TristaZeng/ZS-DeconvNet.git /zs_deconvnet && \ + cd /zs_deconvnet && git checkout 04d2c21 + +# --- FluoResFM (vendored inference code only; weights are never required to +# build -- see download_models.py) --- +RUN git clone https://github.com/qiqi-lu/fluoresfm.git /fluoresfm && \ + cd /fluoresfm && git checkout v1.0.1 +ENV FLUORESFM_REPO_PATH=/fluoresfm +ENV FLUORESFM_WEIGHTS=/opt/fluoresfm_weights/fluoresfm.pt +# BiomedCLIP text-embedder files (open_clip_config.json + open_clip_pytorch_model.bin) +# for the text-conditioned 'unet_sd_c' backbone; supplied at runtime (see +# entrypoint.py resolve_fluoresfm_embedder + the repo's 1_1_download_embedder.py). +ENV FLUORESFM_EMBEDDER_DIR=/opt/fluoresfm_weights/biomedclip + +COPY ./workers/annotations/image_restoration/models_config.py / +COPY ./workers/annotations/image_restoration/download_models.py / +# No FLUORESFM_WEIGHTS_URL on Mac dev builds -- FluoResFM weight download is +# skipped (best-effort only, see download_models.py); Cellpose3 checkpoint +# download is also best-effort and will not fail the build. +RUN python /download_models.py + +COPY ./workers/annotations/image_restoration/entrypoint.py / + +LABEL isUPennContrastWorker="" \ + isAnnotationWorker="" \ + interfaceName="Image Restoration" \ + interfaceCategory="Image Processing" \ + description="Restores images (Noise2Void/N2V2, Cellpose3, ZS-DeconvNet, FluoResFM) -- CPU-only Mac dev build" \ + hasPreview="False" \ + advancedOptionsPanel="False" \ + annotationConfigurationPanel="False" \ + defaultToolName="Image Restoration" + +# ENTRYPOINT inherited from image-processing-base (run_worker.sh /entrypoint.py — see todo/worker-startup-latency.md) diff --git a/workers/annotations/image_restoration/IMAGE_RESTORATION.md b/workers/annotations/image_restoration/IMAGE_RESTORATION.md new file mode 100644 index 0000000..5744638 --- /dev/null +++ b/workers/annotations/image_restoration/IMAGE_RESTORATION.md @@ -0,0 +1,249 @@ +# Image Restoration Worker + +Restores (denoises/deblurs/deconvolves) fluorescence microscopy images using one of four +reference-free/pretrained algorithms: Noise2Void/N2V2, Cellpose3 restoration, ZS-DeconvNet +(zero-shot), or FluoResFM (pretrained foundation model). The user picks a channel set and +a `Method`; the chosen algorithm is applied per channel and the result is uploaded as a new +multi-frame TIFF, exactly like `histogram_matching`/`deconwolf`/`rolling_ball`. + +This worker was designed alongside `illumination_correction` per +[`FLUOR_CORRECTION_WORKERS_SPEC.md`](../../../FLUOR_CORRECTION_WORKERS_SPEC.md) ("Worker 2: +image_restoration"); see that spec for the full design rationale. + +## Why these four methods (and not CARE/3D-RCAN) + +We usually lack **paired clean ground truth** for fluorescence images, so classic supervised +restoration methods (CARE/CSBDeep, 3D-RCAN) are intentionally **excluded** -- they need a +matched low/high-quality image pair to train on, which most users don't have. All four +methods implemented here work **without** paired references: + +| Value | Label | Kind | Dependency | Paired data needed? | +|-------|-------|------|-----------|----------------------| +| `n2v` | Noise2Void / N2V2 (default) | self-supervised, trains on your data | `careamics` (pip, PyTorch) | No | +| `cellpose3` | Cellpose3 restoration | pretrained | `cellpose>=3` (pip) | No | +| `zs_deconvnet` | ZS-DeconvNet | zero-shot, trains on the single input | vendored repo (see below) | No | +| `fluoresfm` | FluoResFM | pretrained, text-prompted foundation model | vendored repo + weights | No | + +Time-lapse-specific denoisers (DeepCAD-RT/FAST/TeD) are noted here as **possible future +additions** but are out of scope for this first pass. + +## How It Works + +1. **Channel selection**: the user selects which channels to restore via `Channels to + restore` (channelCheckboxes). Unselected channels pass through to the output unchanged. +2. **Method selection**: `Method` picks one of the four algorithms below. The interface + always shows *all* per-method parameters (there is no conditional UI in this framework); + each parameter's tooltip states which method it applies to, and parameters for + non-selected methods are simply ignored by `compute()`. +3. **Per-channel collection**: for each selected channel, every frame of that channel + (across XY/Z/Time) is gathered into an `(N, Y, X)` stack via `tileClient.getRegion(...)`. + - `n2v` trains *once* per channel on that channel's whole collection, then predicts on it + (self-supervised -- no clean target is needed, but more frames generally help). + - `zs_deconvnet` is zero-shot per frame: it re-trains a small network from scratch on each + individual frame (no information is shared across frames). + - `cellpose3` and `fluoresfm` are pretrained and simply applied frame by frame. +4. **GPU/CPU**: `Use GPU` (default on) requests CUDA; `resolve_device()` checks + `torch.cuda.is_available()` and falls back to CPU with a `sendWarning` if no GPU is + available (pattern: `deconwolf`'s OpenCL->CPU fallback). CPU is slow for `n2v`, + `zs_deconvnet`, and `fluoresfm`; `cellpose3` is usually still reasonably fast on CPU. +5. **dtype handling**: restored images are computed in float internally, then + `np.nan_to_num`'d and clipped/cast back to the source dtype (read from the dataset's + `dtype` metadata) before being written, to avoid silently bloating a `uint16` source into + a `float32` output TIFF. +6. **Output assembly**: restored and pass-through frames are reassembled into a + multi-dimensional TIFF preserving channel names, pixel size (`mm_x`/`mm_y`), and + magnification, written to `/tmp/restored.tiff` and uploaded back to the dataset's folder. +7. **Provenance**: the uploaded item's metadata records `tool`, `method`, `channels`, + `gpu_requested`, `device_used`, and the method-specific parameters that were actually used. + +## Interface Parameters + +| Parameter | Type | Default | Applies to | Description | +|-----------|------|---------|------------|-------------| +| **Method** | select | `n2v` | all | Restoration algorithm: `n2v`, `cellpose3`, `zs_deconvnet`, `fluoresfm`. | +| **Channels to restore** | channelCheckboxes | -- | all | Channels to process; unselected channels pass through unchanged. | +| **Use GPU** | checkbox | true | all | Use CUDA if available; falls back to CPU with a warning otherwise. | +| **Epochs** | number | 20 | n2v | Self-supervised training epochs on the per-channel collection. | +| **Use N2V2** | checkbox | true | n2v | Use the N2V2 variant (reduces checkerboard artifacts vs. classic N2V). | +| **Patch size** | number | 64 | n2v | Training patch size in pixels (square patches; clamped to image size). | +| **Cellpose3 model** | select | `denoise_cyto3` | cellpose3 | Pretrained restoration checkpoint: `denoise_cyto3`, `deblur_cyto3`, `upsample_cyto3`, `denoise_nuclei`, `deblur_nuclei`, `oneclick_cyto3`. | +| **ZS iterations** | number | 300 | zs_deconvnet | Zero-shot self-supervised training steps, run fresh per frame. | +| **ZS upsampling** | checkbox | false | zs_deconvnet | Runs extra physics-consistency deconvolution refinement instead of the published super-resolution (pixel-grid-changing) mode -- see Implementation Details. | +| **Numerical Aperture (NA)** | number | 0.75 | zs_deconvnet | Used with wavelength/pixel size to build an approximate Gaussian PSF. | +| **Emission Wavelength (nm)** | number | 520 | zs_deconvnet | Used to build the approximate PSF. | +| **Pixel Size XY (nm)** | number | 325 | zs_deconvnet | Used to convert the PSF from physical units to pixels. | +| **FluoResFM backbone** | select | `unet_sd_c` | fluoresfm | Model architecture; **must match the `FLUORESFM_WEIGHTS` checkpoint**. `unet_sd_c` = the real text-conditioned FluoResFM foundation model (uses the prompt + BiomedCLIP embedder). `care`/`dfcan`/`unifmir` = the repo's non-text baseline backbones (prompt ignored). | +| **FluoResFM task** | select | `denoise` | fluoresfm | `denoise`, `deconvolution`, `super-resolution` -- builds the default prompt and selects the `unifmir` task index. | +| **FluoResFM text prompt** | text | (derived from task) | fluoresfm | Free-text prompt (used by `unet_sd_c` only). Left blank, defaults to `"fluorescence microscopy image, "`. | + +## Implementation Details + +### n2v (Noise2Void / N2V2 via CAREamics) + +Uses `careamics.config.create_n2v_configuration(...)` + `careamics.CAREamist` per the +documented API: builds an N2V(2) configuration from `Epochs`/`Use N2V2`/`Patch size`, trains +on the full per-channel `(N, Y, X)` collection (`axes='SYX'`, or `'YX'` for a lone frame), and +predicts on the same collection. Patch size is clamped so it never exceeds the image's own +dimensions. + +### cellpose3 (pretrained restoration) + +Uses `cellpose.denoise.DenoiseModel(model_type=..., gpu=...)` and calls `.eval(frame, +channels=[0, 0])` per frame, matching the documented Cellpose3 restoration API. Weights for +all six exposed checkpoints are pre-downloaded at Docker build time (see +`download_models.py`) so runtime has no network dependency. **Best used as segmentation +preprocessing**, not for quantitative intensity restoration -- Cellpose3's restoration models +are trained to produce clean-looking inputs for its own segmentation, not photometrically +faithful images. + +### zs_deconvnet (zero-shot denoise + deconvolution) -- simplified reimplementation + +**Honest scope note**: `restore_zs_deconvnet()` is a **good-faith, simplified +reimplementation** of the zero-shot, self-supervised idea behind ZS-DeconvNet (Qiao et al., +*Nat. Commun.* 2024; [TristaZeng/ZS-DeconvNet](https://github.com/TristaZeng/ZS-DeconvNet)), +**not** a literal port of the vendored repo's training scripts. The upstream repo's Python +pipeline (`Python_MATLAB_Codes/train_inference_python/`) is driven via shell scripts +(`train_demo_2D.sh`, `infer_demo_2D.sh`, etc.) with a CLI surface that could not be verified +end-to-end without executing it in this environment (no GPU/Docker build available here). The +Dockerfile still vendors the official repo, pinned to tag `v1.0` (commit `04d2c21`), for +provenance and as a base for a future, more faithful integration. + +What is actually implemented, per frame: + +1. **Self-supervised denoising** (`_zs_self_supervised_denoise`): a tiny 3-layer residual CNN + is trained *from scratch on that single frame* using a + [Neighbor2Neighbor](https://arxiv.org/abs/2101.02824)-style objective -- each non-overlapping + 2x2 pixel block is split into two "neighbor-subsampled" half-resolution images `g1`, `g2` + (statistically independent noisy samples of the same underlying signal), and the network is + trained so `f(g1) ~= g2`. This requires no clean reference and no external training data, + consistent with ZS-DeconvNet's "trains on the single input" framing, for `ZS iterations` + steps. +2. **Physics-consistent deconvolution** (`_richardson_lucy_gaussian`): classical + Richardson-Lucy deconvolution against an **approximate Gaussian PSF** built from `NA`, + `Emission Wavelength (nm)`, and `Pixel Size XY (nm)` via the Abbe resolution criterion + (`resolution ~= 0.21 * lambda / NA`), standing in for a full Born-Wolf PSF (c.f. + deconwolf's `dw_bw`) and for the published network's learned deconvolution stage. + +**`ZS upsampling` does not change the output's pixel dimensions.** The published method's +super-resolution mode upsamples the image grid; this worker instead runs more RL iterations +(20 vs. 10) as an in-place sharpening refinement. This is a deliberate constraint, not an +oversight: NimbusImage annotations and pixel-scale calibration are defined against the +original image grid (see CLAUDE.md's coordinate-convention notes), and this worker's contract +is "process channels in place, write back a same-shape TIFF" like every other worker in this +family. True super-resolution upsampling would require a separate resampling/re-registration +step for annotations and is out of scope here. + +### fluoresfm (foundation model, experimental) + +Vendored from [qiqi-lu/fluoresfm](https://github.com/qiqi-lu/fluoresfm) (pinned tag `v1.0.1`), +the canonical repo for FluoResFM (Lu et al., *Nat. Commun.* 2026), a text-conditioned U-Net +trained across 20+ biological structures for denoising/deconvolution/super-resolution. + +**This integration is wired against the repo's *real* inference API** (transcribed from the +repo's own test scripts), replacing an earlier stub that imported non-existent modules +(`methods.fluoresfm`, `packages.text_embedding`) and therefore always failed. The real call +sequence (per `3_0_test_it2i.py`, which the README identifies as "test the FluoResFM model"): + +1. **Text embedder** — `models.biomedclip_embedder.BiomedCLIPTextEmbedder(path_json, + path_bin, context_length=160, device)`; `embedder(prompt)` → conditioning tensor of shape + `[1, 160, 768]`. +2. **Model** — `models.unet_sd_c.UNetModel(in_channels=1, out_channels=1, channels=320, + n_res_blocks=1, attention_levels=[0,1,2,3], channel_multipliers=[1,2,4,4], n_heads=8, + tf_layers=1, d_cond=768, pixel_shuffle=False, scale_factor=4, structural_prompt=False, + n_tokens=0)` (the exact `params` block from `3_0_test_it2i.py`). +3. **Checkpoint** — `torch.load(path, weights_only=True)["model_state_dict"]`, then + `utils.optim.on_load_checkpoint(...)` to strip the `_orig_mod.` prefix that `torch.compile` + adds, then `model.load_state_dict(...)`; `.to(device).eval()`. +4. **Per frame** — clip negatives, percentile-normalize (0.03 / 0.995, identical math to + `utils.data.NormalizePercentile` but inlined to avoid that module's heavy `pydicom`/`scipy`/ + `PSFmodels` import chain), reflect-pad H/W up to a multiple of 8, run + `model(x, None, text_embed)` (forward signature is `(x, time_steps, cond)`; `time_steps` is + unused), crop the padding off, then `_rescale_to_reference()` back to the source frame's + range. + +**Backbones (`FluoResFM backbone`).** The repo also ships three *non-text* baseline backbones +(`models.care.CARE`, `models.dfcan.DFCAN`, `models.unifmir.UniModel`, per `3_1_test_i2i.py` — +the README's "test the model **without inputing the text**"). These are exposed for +completeness and are built/loaded/run faithfully (dfcan/unifmir use `scale_factor=1`/`srscale=1` +so the output grid matches the input; `unifmir` takes the task index sr=1/dn=2/iso=3/dcv=4; +`care` is reflect-padded to a multiple of 4). **They ignore the text prompt.** The default and +recommended backbone is `unet_sd_c` — the genuine text-conditioned FluoResFM. **The selected +backbone must match the architecture the `FLUORESFM_WEIGHTS` checkpoint was trained with** +(the repo ships a separate checkpoint per backbone); a mismatch surfaces as a structured +`load_state_dict` error rather than a crash. + +**Simplifications vs. the upstream script (honest scope):** +- **No sliding-window patch-stitcher.** The repo tiles large images via + `utils.data.Patch_stitcher` with a linear-ramp blend; this worker instead runs each frame + through the model whole (mirroring the repo's own non-patched branch, which triggers when the + image is smaller than the patch size). Large frames therefore need enough GPU memory to + process at full size. +- **Free-text prompt is out-of-distribution.** FluoResFM was trained on highly *structured* + prompts (`"Task: …; sample: …; structure: …; input microscope: …; input pixel size: …; + target microscope: …; target pixel size: …."`). A short free-text prompt still embeds and + runs, but for best results supply a prompt in that structured form. +- **No super-resolution grid change.** As with the other methods, output pixel dimensions are + kept identical to the input so NimbusImage annotations/scale stay aligned. + +**Weights are not baked into the Docker image** (they are large and distributed via Google +Drive / Baidu Yun, not a stable scriptable URL). Two runtime artifacts are needed for the +text-conditioned backbone: + +| Env var | Contents | Source | +|---------|----------|--------| +| `FLUORESFM_WEIGHTS` | the FluoResFM `.pt` checkpoint (per chosen backbone) | fluoresfm README's Google Drive / Baidu Yun links | +| `FLUORESFM_EMBEDDER_DIR` | dir with `open_clip_config.json` + `open_clip_pytorch_model.bin` (~0.4 GB) | HuggingFace `microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224` (repo's `1_1_download_embedder.py`) | + +`resolve_fluoresfm_weights()` and `resolve_fluoresfm_embedder()` read these; if either is +missing, `restore_fluoresfm()` returns `None` after an actionable `sendError` (naming the real +modules/paths) instead of crashing. `download_models.py` still *attempts* a best-effort +checkpoint download via the optional `FLUORESFM_WEIGHTS_URL` build arg (a mirror you control) +and never fails the build if unset. Python deps for this path (`open_clip_torch`, +`transformers`, `huggingface_hub`, `torchinfo`, `timm`) are installed in both Dockerfiles; +`flash-attn` is **not** required (the vendored attention falls back to plain attention and only +uses flash for self-attention, never the text cross-attention). + +**Experimental / untested in CI.** This path requires a GPU plus large runtime weights, so it +is **not exercised by the test suite** (the tests mock `restore_fluoresfm` and only verify the +interface, dispatch, weight/embedder resolvers, and the honest error paths). Restored +intensities from a foundation model may be partially hallucinated — follow the general +guidance below before trusting quantitative results from `fluoresfm` or `zs_deconvnet`. + +### dtype and NaN handling + +All four methods return float arrays internally. `_clip_to_dtype()` runs `np.nan_to_num` +(zeroing any NaN/inf the models might produce) and, for integer source dtypes, clips to the +dtype's representable range before casting -- so a `uint16` input never silently becomes a +`float32` output, and a model excursion above 65535 or below 0 doesn't wrap around. + +### Testability / lazy imports + +Per the project convention (see CLAUDE.md and `histogram_matching`), every heavy third-party +import (`torch`, `careamics`, `cellpose`, the vendored zs_deconvnet/fluoresfm modules) is +lazy -- imported *inside* the relevant `restore_*` function, never at module load time. This +keeps `interface()` fast and lets `tests/test_image_restoration.py` run natively, with none of +those packages installed, by patching `entrypoint.restore_n2v` / `restore_cellpose3` / +`restore_zs_deconvnet` / `restore_fluoresfm` directly. `compute()`'s dispatch table is resolved +via `globals()` *inside* `compute()` (not bound to function objects at import time), so those +patches take effect correctly. + +## Notes + +- **GPU strongly recommended** for `n2v`, `zs_deconvnet`, and `fluoresfm`; `cellpose3` is + comparatively fast even on CPU. The production Dockerfile is built on + `nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04`; `Dockerfile_M1` provides a CPU-only build + (from `nimbusimage/image-processing-base:latest`) for Mac development, with the same + restoration dependencies (CPU PyTorch instead of the CUDA wheel). +- **`zs_deconvnet` and `fluoresfm` are the highest-risk / most experimental methods** in this + worker -- see the honest scope notes above. Validate their output before any quantitative + use. +- **Segment on restored, measure on illumination-corrected raw.** All four restoration + methods can alter absolute intensities (self-supervised denoising suppresses noise + non-uniformly; pretrained/foundation models can hallucinate detail). The recommended + workflow is: use the restored image to find/segment objects (it's usually much easier to + see structure), but compute quantitative intensity properties (e.g. via the + `blob_intensity`/`point_circle_intensity` property workers) on the original or + illumination-corrected (see the sibling `illumination_correction` worker) raw image, not on + the restored one. +- Cellpose3 checkpoints are pre-downloaded at build time (`download_models.py`); ZS-DeconvNet + and FluoResFM repos are vendored via pinned `git clone`s in the Dockerfile. diff --git a/workers/annotations/image_restoration/download_models.py b/workers/annotations/image_restoration/download_models.py new file mode 100644 index 0000000..23e0db4 --- /dev/null +++ b/workers/annotations/image_restoration/download_models.py @@ -0,0 +1,61 @@ +"""Pre-download restoration model weights at build time. + +- Cellpose3 restoration checkpoints (denoise/deblur/upsample) are downloaded + unconditionally, mirroring cellposesam's download_models.py: instantiating + ``denoise.DenoiseModel(model_type=)`` fetches weights to Cellpose's + cache dir on first use, so doing it here bakes them into the image and + avoids a multi-GB download on the first job run. + +- FluoResFM weights are downloaded **best-effort only**. The FluoResFM + authors distribute the pretrained checkpoint via Google Drive / Baidu Yun + (see https://github.com/qiqi-lu/fluoresfm), not a stable, scriptable direct + -download URL, so there is no reliable way to fetch it unattended at build + time. If ``FLUORESFM_WEIGHTS_URL`` is provided as a build-time environment + variable pointing to a mirror you control (e.g. an internal artifact store + or a Zenodo/HF mirror you've set up), this script will try to fetch it into + ``/opt/fluoresfm_weights/fluoresfm.pt``. Any failure here is caught and + logged -- it must NOT fail the Docker build. At runtime, entrypoint.py reads + the checkpoint path from the ``FLUORESFM_WEIGHTS`` environment variable + (see ``resolve_fluoresfm_weights()``), which the Dockerfile can default to + this path, and gracefully ``sendError``s with instructions if it is + missing/unset. +""" + +import os +import urllib.request + +from models_config import CELLPOSE3_RESTORATION_CHECKPOINTS + +print("Downloading Cellpose3 restoration checkpoints...") +try: + from cellpose import denoise + + for checkpoint in CELLPOSE3_RESTORATION_CHECKPOINTS: + print(f" Downloading Cellpose3 restoration checkpoint: {checkpoint}") + denoise.DenoiseModel(model_type=checkpoint, gpu=False) +except Exception as e: # pragma: no cover - build-time diagnostic only + print(f" WARNING: failed to pre-download one or more Cellpose3 checkpoints: {e}") + print(" These will instead be downloaded on first use at runtime.") + +print("Attempting best-effort FluoResFM weight download...") +fluoresfm_weights_url = os.environ.get('FLUORESFM_WEIGHTS_URL', '').strip() +if not fluoresfm_weights_url: + print( + " FLUORESFM_WEIGHTS_URL not set at build time -- skipping. FluoResFM " + "weights are distributed via Google Drive/Baidu Yun " + "(https://github.com/qiqi-lu/fluoresfm) and must be supplied at " + "runtime via the FLUORESFM_WEIGHTS environment variable/volume mount. " + "This is expected and does not fail the build." + ) +else: + try: + os.makedirs('/opt/fluoresfm_weights', exist_ok=True) + dest = '/opt/fluoresfm_weights/fluoresfm.pt' + print(f" Downloading FluoResFM weights from {fluoresfm_weights_url} -> {dest}") + urllib.request.urlretrieve(fluoresfm_weights_url, dest) + print(" FluoResFM weights downloaded successfully.") + except Exception as e: # pragma: no cover - best-effort, must not fail the build + print(f" WARNING: FluoResFM weight download failed ({e}); continuing without it.") + print(" Set FLUORESFM_WEIGHTS at runtime to a mounted checkpoint instead.") + +print("download_models.py complete.") diff --git a/workers/annotations/image_restoration/entrypoint.py b/workers/annotations/image_restoration/entrypoint.py new file mode 100644 index 0000000..67d7d0a --- /dev/null +++ b/workers/annotations/image_restoration/entrypoint.py @@ -0,0 +1,1111 @@ +import argparse +import json +import os +import sys +from collections import defaultdict + +import annotation_client.tiles as tiles +import annotation_client.workers as workers +from annotation_client.utils import sendProgress, sendWarning, sendError + +import numpy as np + + +def interface(image, apiUrl, token): + """Define the worker interface shown to users. + + This worker exposes a single ``Method`` dropdown plus the union of all + per-method parameters (there is no conditional UI in this framework). + Each parameter's tooltip states which method(s) it applies to; unused + parameters for the currently-selected method are simply ignored by + ``compute()``. See FLUOR_CORRECTION_WORKERS_SPEC.md ("Worker 2: + image_restoration") for the full design rationale. + """ + client = workers.UPennContrastWorkerPreviewClient(apiUrl=apiUrl, token=token) + + interface = { + 'Method': { + 'type': 'select', + 'items': ['n2v', 'cellpose3', 'zs_deconvnet', 'fluoresfm'], + 'default': 'n2v', + 'tooltip': ( + 'Restoration algorithm. All four are reference-free/pretrained ' + '(no paired clean ground truth required): n2v (Noise2Void/N2V2, ' + 'self-supervised denoise, trains on your data), cellpose3 ' + '(pretrained denoise/deblur/upsample), zs_deconvnet (zero-shot, ' + 'trains on the single input), fluoresfm (pretrained, ' + 'text-prompted foundation model, experimental).' + ), + 'displayOrder': 0, + }, + 'Channels to restore': { + 'type': 'channelCheckboxes', + 'tooltip': 'Process selected channels; unselected channels pass through unchanged.', + 'displayOrder': 1, + }, + 'Use GPU': { + 'type': 'checkbox', + 'default': True, + 'tooltip': ( + 'Use GPU (CUDA) acceleration for all methods. Falls back to CPU ' + 'automatically with a warning if CUDA is unavailable; CPU ' + 'fallback is slow for n2v/zs_deconvnet/fluoresfm.' + ), + 'displayOrder': 2, + }, + # --- n2v (Noise2Void / N2V2 via CAREamics) --- + 'Epochs': { + 'type': 'number', + 'min': 1, + 'max': 1000, + 'default': 20, + 'tooltip': 'n2v: number of self-supervised training epochs on the per-channel image collection.', + 'displayOrder': 3, + }, + 'Use N2V2': { + 'type': 'checkbox', + 'default': True, + 'tooltip': 'n2v: use the N2V2 variant (reduces checkerboard artifacts vs. classic N2V).', + 'displayOrder': 4, + }, + 'Patch size': { + 'type': 'number', + 'min': 16, + 'max': 512, + 'default': 64, + 'tooltip': 'n2v: training patch size in pixels (square patches).', + 'displayOrder': 5, + }, + # --- cellpose3 restoration --- + 'Cellpose3 model': { + 'type': 'select', + 'items': [ + 'denoise_cyto3', + 'deblur_cyto3', + 'upsample_cyto3', + 'denoise_nuclei', + 'deblur_nuclei', + 'oneclick_cyto3', + ], + 'default': 'denoise_cyto3', + 'tooltip': ( + 'cellpose3: pretrained restoration model to apply per frame. ' + 'Best used as segmentation preprocessing, not for quantitative ' + 'intensity restoration.' + ), + 'displayOrder': 6, + }, + # --- zs_deconvnet (zero-shot denoise + deconvolution) --- + 'ZS iterations': { + 'type': 'number', + 'min': 1, + 'max': 2000, + 'default': 300, + 'tooltip': 'zs_deconvnet: number of zero-shot self-supervised training steps per frame/stack.', + 'displayOrder': 7, + }, + 'ZS upsampling': { + 'type': 'checkbox', + 'default': False, + 'tooltip': ( + 'zs_deconvnet: apply extra physics-consistency deconvolution ' + 'refinement (closer to the published super-resolution mode) ' + 'instead of denoise-only. Output pixel dimensions are always kept ' + 'identical to the input so annotations/scale stay pixel-aligned; ' + 'this does not perform true upsampling of the image grid.' + ), + 'displayOrder': 8, + }, + 'Numerical Aperture (NA)': { + 'type': 'number', + 'min': 0.1, + 'max': 1.7, + 'default': 0.75, + 'tooltip': 'zs_deconvnet: numerical aperture, used to build an approximate Gaussian PSF for the deconvolution stage.', + 'displayOrder': 9, + }, + 'Emission Wavelength (nm)': { + 'type': 'number', + 'min': 300, + 'max': 800, + 'default': 520, + 'tooltip': 'zs_deconvnet: emission wavelength in nanometers, used to build the approximate PSF.', + 'displayOrder': 10, + }, + 'Pixel Size XY (nm)': { + 'type': 'number', + 'min': 1, + 'max': 10000, + 'default': 325, + 'tooltip': 'zs_deconvnet: lateral pixel size in nanometers, used to convert the PSF to pixel units.', + 'displayOrder': 11, + }, + # --- fluoresfm (foundation model) --- + 'FluoResFM backbone': { + 'type': 'select', + 'items': ['unet_sd_c', 'care', 'dfcan', 'unifmir'], + 'default': 'unet_sd_c', + 'tooltip': ( + 'fluoresfm: model architecture. This MUST match the checkpoint ' + 'supplied via the FLUORESFM_WEIGHTS environment variable. ' + '"unet_sd_c" is the real text-conditioned FluoResFM foundation ' + 'model (uses the text prompt via a BiomedCLIP text embedder) -- ' + 'the default and recommended choice. "care"/"dfcan"/"unifmir" ' + 'are the repo\'s non-text baseline backbones (from ' + '3_1_test_i2i.py); the text prompt is IGNORED for those. See ' + 'IMAGE_RESTORATION.md.' + ), + 'displayOrder': 12, + }, + 'FluoResFM task': { + 'type': 'select', + 'items': ['denoise', 'deconvolution', 'super-resolution'], + 'default': 'denoise', + 'tooltip': ( + 'fluoresfm: restoration task. Used to build the default text ' + 'prompt and, for the "unifmir" backbone, to select the task ' + 'index. Experimental.' + ), + 'displayOrder': 13, + }, + 'FluoResFM text prompt': { + 'type': 'text', + 'vueAttrs': { + 'placeholder': 'e.g. fluorescence microscopy image, denoising', + 'label': 'FluoResFM text prompt (optional)', + 'persistentPlaceholder': True, + 'filled': True, + }, + 'tooltip': ( + 'fluoresfm: free-text prompt describing the structure/task ' + '(the "unet_sd_c" backbone is text-conditioned). Leave blank to ' + 'use a default prompt derived from the selected task. Ignored by ' + 'the non-text "care"/"dfcan"/"unifmir" backbones.' + ), + 'displayOrder': 14, + }, + } + client.setWorkerImageInterface(image, interface) + + +def resolve_device(use_gpu): + """Resolve the torch device string ('cuda' or 'cpu') for restoration methods. + + Only imports torch when GPU use is actually requested, so choosing 'Use + GPU' = False never pulls in the heavy torch dependency at all. When GPU is + requested but CUDA is unavailable, falls back to CPU with a sendWarning + (pattern: deconwolf's OpenCL->CPU fallback). + """ + if not use_gpu: + return 'cpu' + + import torch + + if torch.cuda.is_available(): + return 'cuda' + + sendWarning( + "GPU requested but CUDA is not available.", + info=( + "Falling back to CPU. ML-based restoration (n2v, zs_deconvnet, " + "fluoresfm) will be significantly slower on CPU; cellpose3 is " + "usually still reasonably fast." + ), + ) + return 'cpu' + + +def resolve_fluoresfm_weights(): + """Resolve the path to the FluoResFM pretrained checkpoint. + + FluoResFM weights are not reliably baked into the Docker image (the build + only *attempts* a best-effort download, see download_models.py, and must + not fail the build if that download fails). The runtime source of truth is + therefore the FLUORESFM_WEIGHTS environment variable, which should point + to a mounted/downloaded ``.pt`` checkpoint. Returns the resolved path, or + None (after sending an actionable sendError) if unavailable. + """ + weights_path = os.environ.get('FLUORESFM_WEIGHTS', '').strip() + if not weights_path or not os.path.isfile(weights_path): + sendError( + "FluoResFM weights are not available.", + info=( + "Set the FLUORESFM_WEIGHTS environment variable to a mounted " + "FluoResFM checkpoint (.pt) or choose a different Method. See " + "IMAGE_RESTORATION.md for download instructions -- weights are " + "distributed via Google Drive/Baidu Yun from " + "https://github.com/qiqi-lu/fluoresfm, not a stable direct-download " + "URL, so they cannot be baked into the image reliably." + ), + ) + return None + return weights_path + + +def resolve_fluoresfm_embedder(): + """Resolve the local BiomedCLIP text-embedder files for FluoResFM. + + The real text-conditioned FluoResFM (``unet_sd_c`` backbone) conditions on + a BiomedCLIP text embedding. Upstream (``models/biomedclip_embedder.py`` + + ``1_1_download_embedder.py``) loads this embedder from two local files + downloaded from HuggingFace + (``microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224``): + + * ``open_clip_config.json`` + * ``open_clip_pytorch_model.bin`` (~0.4 GB) + + Like the main checkpoint, these are supplied at runtime rather than baked + into the image. ``FLUORESFM_EMBEDDER_DIR`` points at the directory that + holds both files. Returns ``(json_path, bin_path)`` or ``None`` (after an + actionable ``sendError``) if either is missing. + """ + embedder_dir = os.environ.get('FLUORESFM_EMBEDDER_DIR', '').strip() + json_path = os.path.join(embedder_dir, 'open_clip_config.json') if embedder_dir else '' + bin_path = os.path.join(embedder_dir, 'open_clip_pytorch_model.bin') if embedder_dir else '' + if not embedder_dir or not os.path.isfile(json_path) or not os.path.isfile(bin_path): + sendError( + "FluoResFM BiomedCLIP text embedder is not available.", + info=( + "The text-conditioned FluoResFM backbone ('unet_sd_c') needs a " + "BiomedCLIP text embedder. Set FLUORESFM_EMBEDDER_DIR to a " + "directory containing 'open_clip_config.json' and " + "'open_clip_pytorch_model.bin' (download from HuggingFace " + "'microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224', " + "as in the fluoresfm repo's 1_1_download_embedder.py). " + "Alternatively pick a non-text backbone (care/dfcan/unifmir) " + "which does not require the embedder, or a different Method. " + "See IMAGE_RESTORATION.md." + ), + ) + return None + return json_path, bin_path + + +def _rescale_to_reference(output, reference): + """Linearly rescale `output`'s dynamic range to match `reference`'s. + + Pretrained restorers (Cellpose3, FluoResFM) often emit intensities in a + normalized range (e.g. ~0-1) rather than the source sensor range. Casting + such output straight into an integer source dtype via `_clip_to_dtype` + would clip almost everything to 0/1 and produce a near-black TIFF. Mapping + the output's [min, max] onto the reference frame's [min, max] preserves the + original intensity scale for downstream quantitative use. + """ + output = np.asarray(output, dtype=np.float32) + out_min = float(output.min()) + out_max = float(output.max()) + out_span = out_max - out_min + if out_span <= 0: + return np.full_like(output, float(np.asarray(reference).min())) + ref = np.asarray(reference, dtype=np.float32) + ref_min = float(ref.min()) + ref_span = float(ref.max()) - ref_min + return (output - out_min) / out_span * ref_span + ref_min + + +def _clip_to_dtype(image, dtype): + """Clip and cast a restored image back to the source dtype. + + Guards against NaN/inf that ML models can introduce, and avoids silently + bloating the output TIFF to float when the source was e.g. uint16 (see + CLAUDE.md's dtype-handling guidance). + """ + image = np.nan_to_num(image, nan=0.0, posinf=0.0, neginf=0.0) + if dtype is None: + return image + dtype = np.dtype(dtype) + if np.issubdtype(dtype, np.integer): + info = np.iinfo(dtype) + image = np.clip(image, info.min, info.max) + return image.astype(dtype) + + +def _build_method_opts(method, workerInterface): + """Extract only the parameters relevant to the selected method. + + Unused parameters for other methods are ignored (there is no conditional + UI in this framework -- all parameters are always shown; see the + docstring on ``interface()``). + """ + if method == 'n2v': + return { + 'epochs': int(workerInterface.get('Epochs', 20)), + 'use_n2v2': bool(workerInterface.get('Use N2V2', True)), + 'patch_size': int(workerInterface.get('Patch size', 64)), + } + if method == 'cellpose3': + return { + 'model_type': workerInterface.get('Cellpose3 model', 'denoise_cyto3'), + } + if method == 'zs_deconvnet': + return { + 'iterations': int(workerInterface.get('ZS iterations', 300)), + 'upsampling': bool(workerInterface.get('ZS upsampling', False)), + 'NA': float(workerInterface.get('Numerical Aperture (NA)', 0.75)), + 'wavelength': float(workerInterface.get('Emission Wavelength (nm)', 520)), + 'pixel_size_xy': float(workerInterface.get('Pixel Size XY (nm)', 325)), + } + if method == 'fluoresfm': + task = workerInterface.get('FluoResFM task', 'denoise') + backbone = workerInterface.get('FluoResFM backbone', 'unet_sd_c') + default_prompt = f"fluorescence microscopy image, {task}" + prompt = workerInterface.get('FluoResFM text prompt', '') or default_prompt + return { + 'task': task, + 'backbone': backbone, + 'prompt': prompt, + } + return {} + + +# --------------------------------------------------------------------------- +# Per-algorithm restoration functions. +# +# Each function has the signature ``restore_x(stack, opts, device) -> ndarray +# | None`` where ``stack`` is a ``(N, Y, X)`` float array (N frames of one +# channel's collection). Heavy third-party imports (careamics, cellpose, +# torch, the vendored zs_deconvnet/fluoresfm code) are imported lazily, +# *inside* these functions only, so that: +# 1. the interface() path stays fast (see todo/worker-startup-latency.md), +# 2. the pytest suite can run natively without any of these installed, by +# mocking these functions directly (mirrors +# histogram_matching/tests/test_histogram_matching.py's +# entrypoint.match_histograms patch). +# +# On a missing dependency or missing weights, these functions call +# sendError(...) with an actionable message and return None; compute() checks +# for None and aborts cleanly instead of crashing. +# --------------------------------------------------------------------------- + + +def restore_n2v(stack, opts, device): + """Self-supervised denoising with Noise2Void / N2V2 via CAREamics. + + Trains on the channel's own collection of frames (no clean target needed) + then predicts on the same collection. GPU is strongly recommended; CPU + works but is slow. + """ + try: + from careamics import CAREamist + from careamics.config import create_n2v_configuration + except ImportError: + sendError( + "Noise2Void (careamics) is not installed in this worker image.", + info="Install with `pip install careamics` and rebuild the image, or choose a different Method.", + ) + return None + + stack = np.asarray(stack, dtype=np.float32) + if stack.ndim == 2: + stack = stack[np.newaxis, ...] + + n_frames = stack.shape[0] + patch = int(opts.get('patch_size', 64)) + # CAREamics requires the patch to fit inside the image, so clamp to the + # smaller image dimension. Do NOT re-raise a 16px floor above that: for an + # image smaller than 16px on a side, max(16, ...) would force the patch + # back above the image size and CAREamics would reject it. + patch = min(patch, stack.shape[-1], stack.shape[-2]) + + config = create_n2v_configuration( + experiment_name='image_restoration_n2v', + data_type='array', + axes='SYX' if n_frames > 1 else 'YX', + patch_size=[patch, patch], + batch_size=min(16, max(1, n_frames)), + num_epochs=int(opts.get('epochs', 20)), + use_n2v2=bool(opts.get('use_n2v2', True)), + ) + + engine = CAREamist(config) + train_source = stack if n_frames > 1 else stack[0] + engine.train(train_source=train_source) + predicted = engine.predict(source=train_source) + + predicted = np.asarray(predicted, dtype=np.float32).squeeze() + if predicted.ndim == 2: + predicted = predicted[np.newaxis, ...] + return predicted + + +def restore_cellpose3(stack, opts, device): + """Pretrained Cellpose3 restoration (denoise/deblur/upsample), applied per frame. + + Best used as segmentation preprocessing rather than for quantitative + intensity restoration (see IMAGE_RESTORATION.md). + """ + try: + from cellpose import denoise + except ImportError: + sendError( + "Cellpose3 restoration models are not installed in this worker image.", + info="Install with `pip install cellpose>=3` and rebuild the image, or choose a different Method.", + ) + return None + + model_type = opts.get('model_type', 'denoise_cyto3') + gpu = device == 'cuda' + model = denoise.DenoiseModel(model_type=model_type, gpu=gpu) + + stack = np.asarray(stack) + results = [] + for frame in stack: + restored = model.eval(frame, channels=[0, 0]) + # DenoiseModel.eval can return the array directly or a (image, ...) tuple + # depending on cellpose version; normalize to a single 2D array. + if isinstance(restored, (list, tuple)): + restored = restored[0] + restored = np.asarray(restored, dtype=np.float32).squeeze() + # Cellpose restoration output is percentile-normalized (~0-1), not in the + # source sensor range; rescale back so it isn't clipped to near-black by + # the later _clip_to_dtype cast. + restored = _rescale_to_reference(restored, frame) + results.append(restored) + + return np.stack(results, axis=0) + + +def _zs_gaussian_psf_sigma(na, wavelength_nm, pixel_size_nm): + """Approximate the diffraction-limited PSF as an isotropic Gaussian. + + Uses the Abbe resolution criterion (resolution ~= 0.21 * lambda / NA) as a + stand-in for a full Born-Wolf PSF (c.f. deconwolf's dw_bw), so the + zero-shot deconvolution stage below has a usable, if approximate, PSF + without requiring a separate PSF-generation tool. + """ + if na <= 0 or pixel_size_nm <= 0: + return 1.0 + resolution_nm = 0.21 * wavelength_nm / na + sigma_px = resolution_nm / pixel_size_nm + return max(sigma_px, 0.5) + + +def _richardson_lucy_gaussian(image, sigma_px, num_iter): + """Classic Richardson-Lucy deconvolution using a symmetric Gaussian PSF. + + A Gaussian blur is its own adjoint, so `scipy.ndimage.gaussian_filter` can + stand in directly for both the forward and backward convolution steps of + the RL update. + """ + from scipy.ndimage import gaussian_filter + + image = np.clip(image, 0, None).astype(np.float32) + estimate = image.copy() + 1e-6 + for _ in range(max(1, num_iter)): + conv = gaussian_filter(estimate, sigma_px) + 1e-6 + relative_blur = image / conv + estimate = estimate * gaussian_filter(relative_blur, sigma_px) + return estimate + + +def _zs_self_supervised_denoise(frame, iterations, torch_device): + """Train a tiny CNN to denoise a *single* frame with no external data. + + Implements a Neighbor2Neighbor-style (Huang et al., CVPR 2021) + self-supervised objective: each 2x2 pixel block is split into two + "neighbor-subsampled" half-resolution images g1, g2 that are statistically + independent noisy realizations of the same underlying signal; the network + is trained so that f(g1) ~= g2 (and vice versa isn't needed since we only + need a converged denoiser), which requires no clean reference and no + external training data -- consistent with ZS-DeconvNet's "trains on the + single input" zero-shot framing. + + NOTE: this is a good-faith, simplified reimplementation of the *zero-shot, + self-supervised* half of ZS-DeconvNet's published dual-stage approach, not + a literal port of the vendored repo's training scripts -- see + IMAGE_RESTORATION.md for why (the upstream scripts are shell-script/CLI + driven research code whose exact interface could not be verified without + executing it in this environment). + """ + import torch + import torch.nn as nn + import torch.optim as optim + + class _TinyDenoiser(nn.Module): + def __init__(self): + super().__init__() + self.net = nn.Sequential( + nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(inplace=True), + nn.Conv2d(32, 32, 3, padding=1), nn.ReLU(inplace=True), + nn.Conv2d(32, 1, 3, padding=1), + ) + + def forward(self, x): + return x + self.net(x) # residual denoising + + def neighbor_subsample(x): + _, _, h, w = x.shape + h2, w2 = h // 2, w // 2 + x = x[:, :, :h2 * 2, :w2 * 2] + blocks = x.unfold(2, 2, 2).unfold(3, 2, 2) # (B, C, h2, w2, 2, 2) + g1 = blocks[..., 0, 0] + g2 = blocks[..., 1, 1] + return g1, g2 + + x = torch.from_numpy(np.asarray(frame, dtype=np.float32)) + x = x.unsqueeze(0).unsqueeze(0) + # Explicit >0 guard (a negative max is truthy, so `... or 1.0` would not + # catch it and would invert the frame's sign). + x_max = float(x.max().item()) + scale = x_max if x_max > 0 else 1.0 + x = (x / scale).to(torch_device) + + model = _TinyDenoiser().to(torch_device) + optimizer = optim.Adam(model.parameters(), lr=1e-3) + + for _ in range(max(1, iterations)): + g1, g2 = neighbor_subsample(x) + loss = torch.mean((model(g1) - g2) ** 2) + optimizer.zero_grad() + loss.backward() + optimizer.step() + + with torch.no_grad(): + denoised = model(x) + return denoised.squeeze().detach().cpu().numpy() * scale + + +def restore_zs_deconvnet(stack, opts, device): + """Zero-shot denoise + physics-consistent deconvolution, trained per frame. + + Dual-stage pipeline inspired by ZS-DeconvNet (Qiao et al., Nat. Commun. + 2024): (1) a tiny self-supervised network denoises each frame using only + that frame (see `_zs_self_supervised_denoise`); (2) classical + Richardson-Lucy deconvolution against an approximate Gaussian PSF (built + from NA/wavelength/pixel size) enforces physics-consistency with the + optical system, standing in for the published network's learned + deconvolution stage. See IMAGE_RESTORATION.md for the honest scope of this + simplification vs. the vendored upstream repo. + + 'ZS upsampling' intentionally does NOT change the output's pixel + dimensions (unlike the published super-resolution mode): NimbusImage + annotations and pixel-scale calibration are defined against the original + image grid, so changing resolution here would break coordinate alignment + downstream. Instead it runs additional deconvolution refinement. + """ + try: + import torch + except ImportError: + sendError( + "PyTorch is required for ZS-DeconvNet restoration but is not installed.", + info="Rebuild the image with torch installed, or choose a different Method.", + ) + return None + + stack = np.asarray(stack, dtype=np.float32) + if stack.ndim == 2: + stack = stack[np.newaxis, ...] + + iterations = max(1, int(opts.get('iterations', 300))) + upsampling = bool(opts.get('upsampling', False)) + na = float(opts.get('NA', 0.75)) or 0.75 + wavelength_nm = float(opts.get('wavelength', 520)) or 520.0 + pixel_size_nm = float(opts.get('pixel_size_xy', 325)) or 325.0 + psf_sigma_px = _zs_gaussian_psf_sigma(na, wavelength_nm, pixel_size_nm) + # 'upsampling' mode runs extra RL iterations instead of changing resolution; + # see the docstring above and IMAGE_RESTORATION.md. + rl_iterations = 20 if upsampling else 10 + + torch_device = torch.device(device) + + results = [] + for frame in stack: + denoised = _zs_self_supervised_denoise(frame, iterations, torch_device) + deconvolved = _richardson_lucy_gaussian(denoised, psf_sigma_px, rl_iterations) + results.append(deconvolved) + + return np.stack(results, axis=0) + + +# FluoResFM model construction constants, transcribed verbatim from the +# vendored repo's own inference scripts (the source of truth): +# +# * unet_sd_c -> 3_0_test_it2i.py `params` block (the real text-conditioned +# FluoResFM foundation model; README: "test_it2i.py ... test +# the FluoResFM model"). +# * care/dfcan/unifmir -> 3_1_test_i2i.py `model = ...(...)` (the repo's +# non-text baseline backbones; README: "test_i2i.py ... test +# the model without inputing the text"). +# +# The chosen backbone MUST match the architecture the checkpoint was trained +# with (FLUORESFM_WEIGHTS); load_state_dict will otherwise raise, which +# compute() surfaces as a structured error. +_FLUORESFM_CONTEXT_LENGTH = 160 # BiomedCLIP context length for the "ALL" text type +_FLUORESFM_UNET_DIVISOR = 8 # unet_sd_c has 3 downsampling levels -> /8 +# 3_1_test_i2i.py: task ids sr=1, dn=2, iso=3, dcv=4. Our task select maps to: +_FLUORESFM_UNIFMIR_TASK = {'denoise': 2, 'deconvolution': 4, 'super-resolution': 1} + + +def _fluoresfm_load_state_dict(torch, on_load_checkpoint, weights_path, torch_device): + """Load a FluoResFM ``.pt`` checkpoint into a plain state_dict. + + Mirrors 3_0_test_it2i.py / 3_1_test_i2i.py: the checkpoints store a dict + with the weights under ``"model_state_dict"``; ``on_load_checkpoint`` then + strips the ``_orig_mod.`` prefix that ``torch.compile`` adds. We also + tolerate a bare state_dict or the ``"model"``/``"state_dict"`` wrappers so + a slightly different checkpoint layout still loads. + """ + try: + ckpt = torch.load(weights_path, map_location=torch_device, weights_only=True) + except Exception: + # weights_only=True can reject checkpoints pickled with extra objects + # on older/newer torch; fall back to a full load of the local file. + ckpt = torch.load(weights_path, map_location=torch_device) + + state_dict = ckpt + if isinstance(ckpt, dict): + for key in ('model_state_dict', 'model', 'state_dict'): + if key in ckpt and isinstance(ckpt[key], dict): + state_dict = ckpt[key] + break + return on_load_checkpoint(state_dict) + + +def _fluoresfm_normalize_percentile(image, p_low=0.03, p_high=0.995): + """Percentile normalization, identical to utils.data.NormalizePercentile. + + Inlined (not imported) because the repo's ``utils/data.py`` pulls a heavy, + fragile import chain (pydicom, scipy, models.PSFmodels, methods.convolution, + skimage.io) at module load; the normalization math itself is a two-line + percentile min-max, so reproducing it exactly is both faithful and robust. + """ + image = np.asarray(image, dtype=np.float32) + vmin = np.percentile(image, p_low * 100) + vmax = np.percentile(image, p_high * 100) + if vmax == 0: + vmax = np.max(image) + amp = vmax - vmin + if amp == 0: + amp = 1 + return (image - vmin) / amp + + +def _fluoresfm_pad_reflect(torch, x, divisor): + """Reflect-pad the last two dims of ``x`` up to a multiple of ``divisor``. + + Mirrors the repo's reflect-padding before the (non-patched) forward pass so + the U-Net's down/up-sampling produces a same-size output. Returns the + padded tensor and the original (H, W) so the caller can crop back. + """ + h, w = x.shape[-2], x.shape[-1] + pad_h = (divisor - h % divisor) % divisor + pad_w = (divisor - w % divisor) % divisor + if pad_h or pad_w: + # F.pad pad order is (left, right, top, bottom) + x = torch.nn.functional.pad(x, (0, pad_w, 0, pad_h), mode='reflect') + return x, h, w + + +def restore_fluoresfm(stack, opts, device): + """Pretrained FluoResFM restoration (experimental), faithful to the repo. + + Vendored from https://github.com/qiqi-lu/fluoresfm (cloned into /fluoresfm + at build time). This mirrors the repo's own inference scripts: + + * backbone ``unet_sd_c`` -> the real text-conditioned FluoResFM foundation + model (``models.unet_sd_c.UNetModel``), conditioned on a BiomedCLIP text + embedding of the prompt. This is "FluoResFM" per the README + (``test_it2i.py``). Call sequence follows 3_0_test_it2i.py: + build embedder -> embed prompt -> build UNetModel with the exact + `params` args -> load ``["model_state_dict"]`` (strip ``_orig_mod.``) -> + normalize input (percentile 0.03/0.995) -> ``model(x, None, cond)`` -> + rescale output back to the source frame's range. + * backbones ``care``/``dfcan``/``unifmir`` -> the repo's non-text baseline + models (3_1_test_i2i.py). The prompt is ignored; ``unifmir`` uses the + task index. Provided because the interface exposes them, but they are + NOT the text-conditioned foundation model. + + Weights (FLUORESFM_WEIGHTS) and, for ``unet_sd_c``, the BiomedCLIP embedder + (FLUORESFM_EMBEDDER_DIR) are supplied at runtime (see the resolver helpers). + + Experimental: restored intensities may be hallucinated by the model and + should be validated before quantitative use (segment on the restored image, + measure on the illumination-corrected raw image). + """ + weights_path = resolve_fluoresfm_weights() + if weights_path is None: + return None # resolve_fluoresfm_weights() already called sendError + + try: + import torch + except ImportError: + sendError( + "PyTorch is required for FluoResFM but is not installed.", + info="Rebuild the image with torch installed, or choose a different Method.", + ) + return None + + backbone = opts.get('backbone', 'unet_sd_c') + task = opts.get('task', 'denoise') + prompt = opts.get('prompt') or f"fluorescence microscopy image, {task}" + + fluoresfm_repo = os.environ.get('FLUORESFM_REPO_PATH', '/fluoresfm') + if fluoresfm_repo not in sys.path: + sys.path.insert(0, fluoresfm_repo) + + # Import the REAL vendored modules for the chosen backbone. utils.optim only + # depends on numpy, so on_load_checkpoint is safe to import directly. + try: + from utils.optim import on_load_checkpoint + if backbone == 'unet_sd_c': + from models.unet_sd_c import UNetModel + from models.biomedclip_embedder import BiomedCLIPTextEmbedder + elif backbone == 'care': + from models.care import CARE + elif backbone == 'dfcan': + from models.dfcan import DFCAN + elif backbone == 'unifmir': + from models.unifmir import UniModel + else: + sendError( + f"Unknown FluoResFM backbone: {backbone}", + info="Choose one of: unet_sd_c, care, dfcan, unifmir.", + ) + return None + except ImportError as e: + sendError( + "Could not import the vendored FluoResFM model code.", + info=( + f"Import failed: {e}. Verify the FluoResFM repo is vendored at " + f"'{fluoresfm_repo}' (see Dockerfile, cloned from " + "https://github.com/qiqi-lu/fluoresfm) and that its module " + "layout still matches this worker's integration " + "(models.unet_sd_c.UNetModel + models.biomedclip_embedder." + "BiomedCLIPTextEmbedder for the text-conditioned backbone; " + "models.care.CARE / models.dfcan.DFCAN / models.unifmir.UniModel " + "for the non-text baselines; utils.optim.on_load_checkpoint). " + "Required Python deps: torch, and (for unet_sd_c) open_clip_torch" + " + transformers; (for care/dfcan) torchinfo; (for unifmir) " + "torchinfo + timm. Update entrypoint.py's restore_fluoresfm() if " + "the upstream API has changed." + ), + ) + return None + + torch_device = torch.device(device) + + # --- Build the model for the chosen backbone (args from the repo scripts) - + text_embed = None + if backbone == 'unet_sd_c': + # BiomedCLIP text embedder is only needed for the text-conditioned model. + embedder_paths = resolve_fluoresfm_embedder() + if embedder_paths is None: + return None # resolve_fluoresfm_embedder() already called sendError + json_path, bin_path = embedder_paths + embedder = BiomedCLIPTextEmbedder( + path_json=json_path, + path_bin=bin_path, + context_length=_FLUORESFM_CONTEXT_LENGTH, + device=torch_device, + ) + embedder.eval() + with torch.no_grad(): + # embedder(prompt) -> conditioning of shape [1, context_length, 768] + text_embed = embedder(prompt).to(torch_device) + + # 3_0_test_it2i.py `params` block (default ALL-v2 checkpoint). + model = UNetModel( + in_channels=1, + out_channels=1, + channels=320, + n_res_blocks=1, + attention_levels=[0, 1, 2, 3], + channel_multipliers=[1, 2, 4, 4], + n_heads=8, + tf_layers=1, + d_cond=768, + pixel_shuffle=False, + scale_factor=4, + structural_prompt=False, + n_tokens=0, + ) + elif backbone == 'care': + # 3_1_test_i2i.py CARE(...) args. + model = CARE( + in_channels=1, + out_channels=1, + n_filter_base=16, + kernel_size=5, + batch_norm=False, + dropout=0.0, + residual=True, + expansion=2, + pos_out=False, + ) + elif backbone == 'dfcan': + # 3_1_test_i2i.py DFCAN(...) args; scale_factor=1 keeps output HxW == input. + model = DFCAN(in_channels=1, scale_factor=1, num_features=64, num_groups=4) + else: # unifmir + # 3_1_test_i2i.py UniModel(...) args (srscale=1 keeps output HxW == input). + model = UniModel( + in_channels=1, + out_channels=1, + tsk=0, + img_size=(64, 64), + patch_size=1, + embed_dim=180 // 2, + depths=[6, 6, 6], + num_heads=[6, 6, 6], + window_size=8, + mlp_ratio=2, + qkv_bias=True, + qk_scale=None, + drop_rate=0, + attn_drop_rate=0, + drop_path_rate=0.1, + norm_layer=torch.nn.LayerNorm, + patch_norm=True, + use_checkpoint=False, + num_feat=32, + srscale=1, + ) + + model = model.to(torch_device) + + # --- Load checkpoint (matching architecture required) -------------------- + try: + state_dict = _fluoresfm_load_state_dict( + torch, on_load_checkpoint, weights_path, torch_device) + model.load_state_dict(state_dict) + except Exception as e: + sendError( + "Failed to load the FluoResFM checkpoint into the selected backbone.", + info=( + f"{e}. The FLUORESFM_WEIGHTS checkpoint must match the chosen " + f"'FluoResFM backbone' ('{backbone}'). Verify you selected the " + "architecture the checkpoint was trained with (the vendored " + "repo ships separate checkpoints per backbone)." + ), + ) + return None + model.eval() + + # unifmir needs a task index tensor; other backbones ignore it. + unifmir_task = torch.tensor( + _FLUORESFM_UNIFMIR_TASK.get(task, 2), device=torch_device) + + # Divisor for reflect-padding so down/up-sampling yields a same-size output. + if backbone == 'unet_sd_c' or backbone == 'unifmir': + divisor = _FLUORESFM_UNET_DIVISOR + elif backbone == 'care': + divisor = 4 # 3_1_test_i2i.py pads care inputs to a multiple of 4 + else: # dfcan (scale_factor=1) accepts any size + divisor = 1 + + stack = np.asarray(stack, dtype=np.float32) + results = [] + with torch.no_grad(): + for frame in stack: + # Repo input normalization: clip negatives, then percentile min-max. + norm = _fluoresfm_normalize_percentile(np.clip(frame, 0, None)) + x = torch.from_numpy(norm).unsqueeze(0).unsqueeze(0).to(torch_device) + x, h, w = _fluoresfm_pad_reflect(torch, x, divisor) + + if backbone == 'unet_sd_c': + # forward(x, time_steps, cond); time_steps unused (None). + restored = model(x, None, text_embed) + elif backbone == 'unifmir': + restored = model(x, unifmir_task) + else: # care / dfcan + restored = model(x) + + restored = restored[..., :h, :w] # crop reflect-padding back off + restored = restored.squeeze().float().detach().cpu().numpy() + # Map the model's (roughly normalized) output range back onto the + # source frame's, so it isn't clipped to near-black downstream. + restored = _rescale_to_reference(restored, frame) + results.append(restored) + + return np.stack(results, axis=0) + + +# Method name -> per-algorithm function name. Resolved via globals() *inside* +# compute() (not bound to function objects here at module-import time) so +# that tests can patch e.g. `entrypoint.restore_n2v` with a mock and have +# compute() pick up the patched version -- a dict built once at import time +# would instead keep permanent references to the original, unpatched +# functions. +_METHOD_FUNCTION_NAMES = { + 'n2v': 'restore_n2v', + 'cellpose3': 'restore_cellpose3', + 'zs_deconvnet': 'restore_zs_deconvnet', + 'fluoresfm': 'restore_fluoresfm', +} + + +def compute(datasetId, apiUrl, token, params): + """ + params (could change): + configurationId, + datasetId, + description: tool description, + type: tool type, + id: tool id, + name: tool name, + image: docker image, + channel: annotation channel, + assignment: annotation assignment ({XY, Z, Time}), + tags: annotation tags (list of strings), + tile: tile position (TODO: roi) ({XY, Z, Time}), + connectTo: how new annotations should be connected + """ + # Lazy import: keeps large_image off the interface/preview path; only needed during compute. See todo/worker-startup-latency.md + import large_image as li + + tileClient = tiles.UPennContrastDataset( + apiUrl=apiUrl, token=token, datasetId=datasetId) + + workerInterface = params['workerInterface'] + + method = workerInterface.get('Method', 'n2v') + allChannels = workerInterface.get('Channels to restore', {}) + channels = [int(k) for k, v in allChannels.items() if v] + print(f"Selected channels to restore: {channels}, method={method}") + + if len(channels) == 0: + sendError("No channels selected for restoration", + info="Select at least one channel in 'Channels to restore'.") + return + + if 'frames' not in tileClient.tiles or not tileClient.tiles.get('frames'): + sendError("No frames found in dataset", + info="This worker requires a dataset with at least one frame.") + return + + frames = tileClient.tiles['frames'] + + restore_fn_name = _METHOD_FUNCTION_NAMES.get(method) + if restore_fn_name is None: + sendError(f"Unknown restoration method: {method}") + return + restore_fn = globals()[restore_fn_name] + + use_gpu_requested = bool(workerInterface.get('Use GPU', True)) + device = resolve_device(use_gpu_requested) + print(f"Using device: {device}") + + source_dtype = tileClient.tiles.get('dtype', None) + method_opts = _build_method_opts(method, workerInterface) + + gc = tileClient.client + + # Group frame indices by channel: retrospective/collection-based methods + # (n2v trains on the whole channel collection; zs_deconvnet trains + # per-frame but is naturally called once per channel's stack) see the + # full (N, Y, X) stack; per-frame methods (cellpose3, fluoresfm) simply + # loop internally over that same stack. + channel_frame_indices = defaultdict(list) + for i, frame in enumerate(frames): + c = frame.get('IndexC', 0) + if c in channels: + channel_frame_indices[c].append(i) + + restored_images = {} + total_channels = len(channels) + for idx, c in enumerate(channels): + frame_indices = channel_frame_indices.get(c, []) + if not frame_indices: + continue + + sendProgress(idx / max(total_channels, 1), 'Restoring', + f"Channel {c} ({method}), {len(frame_indices)} frame(s)") + + try: + images = [tileClient.getRegion(datasetId, frame=i).squeeze() for i in frame_indices] + # np.stack is inside the try so heterogeneous per-frame shapes + # (which raise ValueError) surface as a structured error, not a + # raw traceback -- like every other failure path in this worker. + stack = np.stack(images, axis=0) + result_stack = restore_fn(stack, method_opts, device) + except Exception as e: + sendError(f"Restoration failed for channel {c} using method '{method}'", + info=str(e)) + return + + if result_stack is None: + # The restore_* function already called sendError with an + # actionable message (missing dependency / missing weights); + # abort cleanly without crashing. + return + + result_stack = _clip_to_dtype(np.asarray(result_stack), source_dtype) + + # A restorer must return exactly one output frame per input frame; + # otherwise zip() below would silently drop/misalign frames (e.g. if + # CAREamics predict() ever returns a different sample count), leaving + # some frames un-restored in the output with no error. + if len(result_stack) != len(frame_indices): + sendError(f"Restoration returned {len(result_stack)} frame(s) but " + f"{len(frame_indices)} were submitted for channel {c} (method '{method}')", + info="This is an internal mismatch; please report it.") + return + + for frame_idx, restored_image in zip(frame_indices, result_stack): + restored_images[frame_idx] = restored_image + + sendProgress(0.9, 'Assembling output', 'Writing frames') + sink = li.new() + for i, frame in enumerate(frames): + # Create a parameters dictionary with only the indices that exist in frame + # The len(k) > 5 is to avoid the 'Index' key that has no postfix to it + large_image_params = {f'{k.lower()[5:]}': v for k, v in frame.items() + if k.startswith('Index') and len(k) > 5} + + if i in restored_images: + image = restored_images[i] + else: + image = tileClient.getRegion(datasetId, frame=i).squeeze() + + sink.addTile(image, 0, 0, **large_image_params) + + # Copy over the metadata + if 'channels' in tileClient.tiles: + sink.channelNames = tileClient.tiles['channels'] + sink.mm_x = tileClient.tiles['mm_x'] + sink.mm_y = tileClient.tiles['mm_y'] + sink.magnification = tileClient.tiles['magnification'] + + sendProgress(0.95, 'Writing output', 'Saving TIFF file') + sink.write('/tmp/restored.tiff') + print("Wrote to file") + + sendProgress(0.98, 'Uploading', 'Uploading to server') + item = gc.uploadFileToFolder(datasetId, '/tmp/restored.tiff') + gc.addMetadataToItem(item['itemId'], { + 'tool': 'Image Restoration', + 'method': method, + 'channels': channels, + 'gpu_requested': use_gpu_requested, + 'device_used': device, + **method_opts, + }) + print("Uploaded file") + + sendProgress(1.0, 'Complete', 'Image restoration finished') + + +if __name__ == '__main__': + # Define the command-line interface for the entry point + parser = argparse.ArgumentParser( + description='Restore images via Noise2Void/N2V2, Cellpose3, ZS-DeconvNet, or FluoResFM') + + 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:]) + + params = json.loads(args.parameters) + datasetId = args.datasetId + apiUrl = args.apiUrl + token = args.token + + match args.request: + case 'compute': + compute(datasetId, apiUrl, token, params) + case 'interface': + interface(params['image'], apiUrl, token) diff --git a/workers/annotations/image_restoration/environment.yml b/workers/annotations/image_restoration/environment.yml new file mode 100644 index 0000000..224d055 --- /dev/null +++ b/workers/annotations/image_restoration/environment.yml @@ -0,0 +1,19 @@ +name: worker +channels: +- conda-forge +- defaults +dependencies: +- python=3.11 +- pip +- numpy>=2.0 +- scipy +- scikit-image +- shapely>=2.0.6 +- libtiff +- openslide +- libvips +- tifffile +- pip: + - careamics + - cellpose>=3,<4 + - triton diff --git a/workers/annotations/image_restoration/models_config.py b/workers/annotations/image_restoration/models_config.py new file mode 100644 index 0000000..2148f6e --- /dev/null +++ b/workers/annotations/image_restoration/models_config.py @@ -0,0 +1,13 @@ +"""Single source of truth for the Cellpose3 restoration checkpoints this +worker exposes via the 'Cellpose3 model' select (see entrypoint.py) and +pre-downloads at build time (see download_models.py). +""" + +CELLPOSE3_RESTORATION_CHECKPOINTS = [ + 'denoise_cyto3', + 'deblur_cyto3', + 'upsample_cyto3', + 'denoise_nuclei', + 'deblur_nuclei', + 'oneclick_cyto3', +] diff --git a/workers/annotations/image_restoration/tests/Dockerfile_Test b/workers/annotations/image_restoration/tests/Dockerfile_Test new file mode 100644 index 0000000..817056a --- /dev/null +++ b/workers/annotations/image_restoration/tests/Dockerfile_Test @@ -0,0 +1,14 @@ +# ./workers/annotations/image_restoration/tests/Dockerfile_Test +FROM annotations/image_restoration:latest AS test + +# Install test dependencies +SHELL ["conda", "run", "-n", "worker", "/bin/bash", "-c"] +RUN pip install pytest pytest-mock + +# Copy test files +RUN mkdir -p /tests +COPY ./workers/annotations/image_restoration/tests/*.py /tests +WORKDIR /tests + +# Command to run tests +ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "worker", "python3", "-m", "pytest", "-v"] diff --git a/workers/annotations/image_restoration/tests/__init__.py b/workers/annotations/image_restoration/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/workers/annotations/image_restoration/tests/test_image_restoration.py b/workers/annotations/image_restoration/tests/test_image_restoration.py new file mode 100644 index 0000000..886ae57 --- /dev/null +++ b/workers/annotations/image_restoration/tests/test_image_restoration.py @@ -0,0 +1,531 @@ +import sys +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +# Import the worker module under test. All heavy third-party imports +# (torch, careamics, cellpose, the vendored zs_deconvnet/fluoresfm code) are +# lazy inside the per-algorithm restore_* functions, so this module imports +# cleanly without any of those packages installed. +from entrypoint import ( + compute, + interface, + resolve_device, + resolve_fluoresfm_weights, + resolve_fluoresfm_embedder, + _build_method_opts, + _clip_to_dtype, +) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture +def mock_tile_client(): + """Mock the tiles.UPennContrastDataset with a 2-channel, multi-frame dataset.""" + with patch('annotation_client.tiles.UPennContrastDataset') as mock_client: + client = mock_client.return_value + client.tiles = { + 'frames': [ + {'IndexXY': 0, 'IndexZ': 0, 'IndexT': 0, 'IndexC': 0}, + {'IndexXY': 0, 'IndexZ': 0, 'IndexT': 0, 'IndexC': 1}, + {'IndexXY': 0, 'IndexZ': 0, 'IndexT': 1, 'IndexC': 0}, + {'IndexXY': 0, 'IndexZ': 0, 'IndexT': 1, 'IndexC': 1}, + {'IndexXY': 1, 'IndexZ': 0, 'IndexT': 0, 'IndexC': 0}, + {'IndexXY': 1, 'IndexZ': 0, 'IndexT': 0, 'IndexC': 1}, + ], + 'IndexRange': { + 'IndexXY': 2, + 'IndexZ': 1, + 'IndexT': 2, + 'IndexC': 2, + }, + 'channels': ['DAPI', 'FITC'], + 'mm_x': 0.65, + 'mm_y': 0.65, + 'magnification': 20, + 'dtype': np.uint16, + } + # Channel 0 -> frame indices 0, 2, 4 ; Channel 1 -> frame indices 1, 3, 5 + client.getRegion.return_value = np.random.randint(0, 1000, (64, 64), dtype=np.uint16) + + mock_gc = MagicMock() + mock_gc.uploadFileToFolder.return_value = {'itemId': 'test_item_id'} + client.client = mock_gc + + yield client + + +@pytest.fixture +def mock_worker_preview_client(): + with patch('annotation_client.workers.UPennContrastWorkerPreviewClient') as mock_client: + yield mock_client.return_value + + +@pytest.fixture +def mock_large_image(): + with patch('large_image.new') as mock_li_new: + mock_sink = MagicMock() + mock_li_new.return_value = mock_sink + yield mock_sink + + +def _passthrough_restore(stack, opts, device): + """Benign stand-in for a real restore_* function: identity, cast to float32.""" + return np.asarray(stack, dtype=np.float32) + + +@pytest.fixture +def mock_restore_all(): + """Patch all four per-algorithm functions with an identity mock. + + Used by tests that exercise the surrounding plumbing (dispatch, channel + filtering, output/metadata, error paths, progress) without depending on + any restoration algorithm's real numerics or third-party libraries -- + mirrors histogram_matching/tests/test_histogram_matching.py's + `entrypoint.match_histograms` patch. + """ + with patch('entrypoint.restore_n2v', side_effect=_passthrough_restore) as m_n2v, \ + patch('entrypoint.restore_cellpose3', side_effect=_passthrough_restore) as m_cellpose3, \ + patch('entrypoint.restore_zs_deconvnet', side_effect=_passthrough_restore) as m_zs, \ + patch('entrypoint.restore_fluoresfm', side_effect=_passthrough_restore) as m_fluoresfm: + yield { + 'n2v': m_n2v, + 'cellpose3': m_cellpose3, + 'zs_deconvnet': m_zs, + 'fluoresfm': m_fluoresfm, + } + + +def base_params(method='n2v', channels=None, use_gpu=False, **extra_interface): + """Build a minimal params dict. `use_gpu` defaults to False so tests never + need a real torch install (resolve_device only imports torch when GPU use + is actually requested).""" + if channels is None: + channels = {'0': True, '1': False} + worker_interface = { + 'Method': method, + 'Channels to restore': channels, + 'Use GPU': use_gpu, + } + worker_interface.update(extra_interface) + return {'workerInterface': worker_interface} + + +# --------------------------------------------------------------------------- +# interface() +# --------------------------------------------------------------------------- + +def test_interface(mock_worker_preview_client): + interface('test_image', 'http://test-api', 'test-token') + + mock_worker_preview_client.setWorkerImageInterface.assert_called_once() + call_args = mock_worker_preview_client.setWorkerImageInterface.call_args + image_arg = call_args[0][0] + interface_data = call_args[0][1] + + assert image_arg == 'test_image' + + # Method select + method_field = interface_data['Method'] + assert method_field['type'] == 'select' + assert method_field['items'] == ['n2v', 'cellpose3', 'zs_deconvnet', 'fluoresfm'] + assert method_field['default'] == 'n2v' + + # Channel selection + assert interface_data['Channels to restore']['type'] == 'channelCheckboxes' + + # Use GPU + assert interface_data['Use GPU']['type'] == 'checkbox' + assert interface_data['Use GPU']['default'] is True + + # n2v params + assert interface_data['Epochs']['type'] == 'number' + assert interface_data['Epochs']['default'] == 20 + assert interface_data['Use N2V2']['type'] == 'checkbox' + assert interface_data['Use N2V2']['default'] is True + assert interface_data['Patch size']['default'] == 64 + + # cellpose3 params + cellpose_field = interface_data['Cellpose3 model'] + assert cellpose_field['type'] == 'select' + assert cellpose_field['default'] == 'denoise_cyto3' + assert 'denoise_cyto3' in cellpose_field['items'] + assert 'upsample_cyto3' in cellpose_field['items'] + + # zs_deconvnet params + assert interface_data['ZS iterations']['type'] == 'number' + assert interface_data['ZS upsampling']['type'] == 'checkbox' + assert interface_data['Numerical Aperture (NA)']['type'] == 'number' + assert interface_data['Emission Wavelength (nm)']['type'] == 'number' + assert interface_data['Pixel Size XY (nm)']['type'] == 'number' + + # fluoresfm params + backbone_field = interface_data['FluoResFM backbone'] + assert backbone_field['type'] == 'select' + assert backbone_field['items'] == ['unet_sd_c', 'care', 'dfcan', 'unifmir'] + assert backbone_field['default'] == 'unet_sd_c' + + task_field = interface_data['FluoResFM task'] + assert task_field['type'] == 'select' + assert task_field['items'] == ['denoise', 'deconvolution', 'super-resolution'] + assert task_field['default'] == 'denoise' + assert interface_data['FluoResFM text prompt']['type'] == 'text' + + +# --------------------------------------------------------------------------- +# Dispatch +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize('method', ['n2v', 'cellpose3', 'zs_deconvnet', 'fluoresfm']) +def test_dispatch_calls_only_selected_method(mock_tile_client, mock_large_image, mock_restore_all, method): + params = base_params(method=method, channels={'0': True, '1': True}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + for name, mock_fn in mock_restore_all.items(): + if name == method: + mock_fn.assert_called() + else: + mock_fn.assert_not_called() + + +def test_unknown_method_sends_error(mock_tile_client, mock_large_image, mock_restore_all, capsys): + params = base_params(method='not_a_real_method', channels={'0': True}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + captured = capsys.readouterr() + assert '"type": "error"' in captured.out + for mock_fn in mock_restore_all.values(): + mock_fn.assert_not_called() + mock_large_image.write.assert_not_called() + + +# --------------------------------------------------------------------------- +# Channel filtering +# --------------------------------------------------------------------------- + +def test_channel_filtering_only_selected_processed(mock_tile_client, mock_large_image, mock_restore_all): + params = base_params(method='n2v', channels={'0': True, '1': False}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + # Channel 0 has 3 frames in the fixture (indices 0, 2, 4) + mock_restore_all['n2v'].assert_called_once() + stack_arg = mock_restore_all['n2v'].call_args[0][0] + assert stack_arg.shape[0] == 3 + + # All 6 frames still make it into the sink (3 restored + 3 passthrough) + assert mock_large_image.addTile.call_count == 6 + + +def test_channel_filtering_multi_channel(mock_tile_client, mock_large_image, mock_restore_all): + params = base_params(method='n2v', channels={'0': True, '1': True}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + # Called once per selected channel + assert mock_restore_all['n2v'].call_count == 2 + + +# --------------------------------------------------------------------------- +# Output plumbing +# --------------------------------------------------------------------------- + +def test_output_plumbing(mock_tile_client, mock_large_image, mock_restore_all): + params = base_params(method='n2v', channels={'0': True, '1': True}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + mock_large_image.write.assert_called_once_with('/tmp/restored.tiff') + mock_tile_client.client.uploadFileToFolder.assert_called_once_with( + 'test_dataset', '/tmp/restored.tiff') + + mock_tile_client.client.addMetadataToItem.assert_called_once() + item_id, metadata = mock_tile_client.client.addMetadataToItem.call_args[0] + assert item_id == 'test_item_id' + assert metadata['tool'] == 'Image Restoration' + assert metadata['method'] == 'n2v' + assert set(metadata['channels']) == {0, 1} + assert 'device_used' in metadata + assert 'gpu_requested' in metadata + + +def test_frame_parameter_construction(mock_tile_client, mock_large_image, mock_restore_all): + params = base_params(method='n2v', channels={'0': True}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + mock_large_image.addTile.assert_called() + for call in mock_large_image.addTile.call_args_list: + kwargs = call[1] + assert any(key in ['xy', 'z', 't', 'c'] for key in kwargs.keys()) + + +# --------------------------------------------------------------------------- +# Metadata preservation +# --------------------------------------------------------------------------- + +def test_metadata_preservation(mock_tile_client, mock_large_image, mock_restore_all): + params = base_params(method='n2v', channels={'0': True}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + assert mock_large_image.channelNames == ['DAPI', 'FITC'] + assert mock_large_image.mm_x == 0.65 + assert mock_large_image.mm_y == 0.65 + assert mock_large_image.magnification == 20 + + +# --------------------------------------------------------------------------- +# Error paths +# --------------------------------------------------------------------------- + +def test_no_channels_selected_error(mock_tile_client, mock_large_image, mock_restore_all, capsys): + params = base_params(method='n2v', channels={'0': False, '1': False}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + captured = capsys.readouterr() + assert '"error": "No channels selected for restoration"' in captured.out + assert '"type": "error"' in captured.out + mock_large_image.write.assert_not_called() + for mock_fn in mock_restore_all.values(): + mock_fn.assert_not_called() + + +def test_no_frames_error(mock_tile_client, mock_large_image, mock_restore_all, capsys): + del mock_tile_client.tiles['frames'] + params = base_params(method='n2v', channels={'0': True}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + captured = capsys.readouterr() + assert '"error": "No frames found in dataset"' in captured.out + mock_large_image.write.assert_not_called() + + +def test_restore_exception_sends_error_not_crash(mock_tile_client, mock_large_image, capsys): + with patch('entrypoint.restore_n2v', side_effect=RuntimeError("boom")): + params = base_params(method='n2v', channels={'0': True}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + captured = capsys.readouterr() + assert '"type": "error"' in captured.out + assert 'boom' in captured.out + mock_large_image.write.assert_not_called() + + +def test_single_frame_dataset_handled(mock_large_image, mock_restore_all): + """A single-frame dataset should be processed normally, not crash.""" + with patch('annotation_client.tiles.UPennContrastDataset') as mock_client: + client = mock_client.return_value + client.tiles = { + 'frames': [{'IndexXY': 0, 'IndexZ': 0, 'IndexT': 0, 'IndexC': 0}], + 'channels': ['DAPI'], + 'mm_x': 1.0, 'mm_y': 1.0, 'magnification': 40, + 'dtype': np.uint8, + } + client.getRegion.return_value = np.zeros((32, 32), dtype=np.uint8) + mock_gc = MagicMock() + mock_gc.uploadFileToFolder.return_value = {'itemId': 'solo_item'} + client.client = mock_gc + + params = base_params(method='n2v', channels={'0': True}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + mock_large_image.write.assert_called_once_with('/tmp/restored.tiff') + mock_gc.addMetadataToItem.assert_called_once() + + +# --------------------------------------------------------------------------- +# Progress reporting +# --------------------------------------------------------------------------- + +def test_progress_reporting(mock_tile_client, mock_large_image, mock_restore_all, capsys): + params = base_params(method='n2v', channels={'0': True}) + compute('test_dataset', 'http://test-api', 'test-token', params) + + captured = capsys.readouterr() + assert '"progress":' in captured.out + assert 'Complete' in captured.out + + +# --------------------------------------------------------------------------- +# resolve_device (GPU -> CPU fallback) +# --------------------------------------------------------------------------- + +def test_resolve_device_cpu_when_gpu_not_requested(): + # Should never import torch when use_gpu is False. + assert resolve_device(False) == 'cpu' + + +def test_resolve_device_gpu_available(monkeypatch): + fake_torch = MagicMock() + fake_torch.cuda.is_available.return_value = True + monkeypatch.setitem(sys.modules, 'torch', fake_torch) + + assert resolve_device(True) == 'cuda' + + +def test_resolve_device_falls_back_to_cpu_with_warning(monkeypatch, capsys): + fake_torch = MagicMock() + fake_torch.cuda.is_available.return_value = False + monkeypatch.setitem(sys.modules, 'torch', fake_torch) + + result = resolve_device(True) + + assert result == 'cpu' + captured = capsys.readouterr() + assert '"type": "warning"' in captured.out + assert 'GPU requested but CUDA is not available' in captured.out + + +def test_compute_records_gpu_device_used(mock_tile_client, mock_large_image, mock_restore_all, monkeypatch): + fake_torch = MagicMock() + fake_torch.cuda.is_available.return_value = True + monkeypatch.setitem(sys.modules, 'torch', fake_torch) + + params = base_params(method='n2v', channels={'0': True}, use_gpu=True) + compute('test_dataset', 'http://test-api', 'test-token', params) + + metadata = mock_tile_client.client.addMetadataToItem.call_args[0][1] + assert metadata['gpu_requested'] is True + assert metadata['device_used'] == 'cuda' + + +# --------------------------------------------------------------------------- +# FluoResFM weight resolution / graceful degradation +# --------------------------------------------------------------------------- + +def test_resolve_fluoresfm_weights_missing_sends_error(monkeypatch, capsys): + monkeypatch.delenv('FLUORESFM_WEIGHTS', raising=False) + + result = resolve_fluoresfm_weights() + + assert result is None + captured = capsys.readouterr() + assert '"type": "error"' in captured.out + assert 'FluoResFM weights are not available' in captured.out + + +def test_resolve_fluoresfm_weights_present(monkeypatch, tmp_path): + weights_file = tmp_path / "fluoresfm.pt" + weights_file.write_bytes(b"fake-checkpoint") + monkeypatch.setenv('FLUORESFM_WEIGHTS', str(weights_file)) + + result = resolve_fluoresfm_weights() + + assert result == str(weights_file) + + +def test_resolve_fluoresfm_embedder_missing_sends_error(monkeypatch, capsys): + monkeypatch.delenv('FLUORESFM_EMBEDDER_DIR', raising=False) + + result = resolve_fluoresfm_embedder() + + assert result is None + captured = capsys.readouterr() + assert '"type": "error"' in captured.out + assert 'BiomedCLIP text embedder is not available' in captured.out + + +def test_resolve_fluoresfm_embedder_present(monkeypatch, tmp_path): + (tmp_path / 'open_clip_config.json').write_text('{}') + (tmp_path / 'open_clip_pytorch_model.bin').write_bytes(b'fake') + monkeypatch.setenv('FLUORESFM_EMBEDDER_DIR', str(tmp_path)) + + result = resolve_fluoresfm_embedder() + + assert result == ( + str(tmp_path / 'open_clip_config.json'), + str(tmp_path / 'open_clip_pytorch_model.bin'), + ) + + +def test_fluoresfm_missing_weights_aborts_compute_cleanly(mock_tile_client, mock_large_image, monkeypatch): + """Simulates restore_fluoresfm() detecting missing weights (via + resolve_fluoresfm_weights) and returning None; compute() must abort + without writing output or crashing.""" + monkeypatch.delenv('FLUORESFM_WEIGHTS', raising=False) + + with patch('entrypoint.restore_fluoresfm', return_value=None) as mock_fluoresfm: + params = base_params(method='fluoresfm', channels={'0': True}) + compute('test_dataset', 'http://test-api', 'test-token', params) + mock_fluoresfm.assert_called_once() + + mock_large_image.write.assert_not_called() + + +# --------------------------------------------------------------------------- +# Small pure-function unit tests +# --------------------------------------------------------------------------- + +def test_build_method_opts_n2v(): + opts = _build_method_opts('n2v', {'Epochs': 50, 'Use N2V2': False, 'Patch size': 32}) + assert opts == {'epochs': 50, 'use_n2v2': False, 'patch_size': 32} + + +def test_build_method_opts_cellpose3_default(): + opts = _build_method_opts('cellpose3', {}) + assert opts == {'model_type': 'denoise_cyto3'} + + +def test_build_method_opts_zs_deconvnet(): + opts = _build_method_opts('zs_deconvnet', { + 'ZS iterations': 100, + 'ZS upsampling': True, + 'Numerical Aperture (NA)': 1.2, + 'Emission Wavelength (nm)': 488, + 'Pixel Size XY (nm)': 110, + }) + assert opts['iterations'] == 100 + assert opts['upsampling'] is True + assert opts['NA'] == 1.2 + assert opts['wavelength'] == 488 + assert opts['pixel_size_xy'] == 110 + + +def test_build_method_opts_fluoresfm_default_prompt(): + opts = _build_method_opts('fluoresfm', {'FluoResFM task': 'deconvolution', 'FluoResFM text prompt': ''}) + assert opts['task'] == 'deconvolution' + assert 'deconvolution' in opts['prompt'] + # Backbone defaults to the real text-conditioned foundation model. + assert opts['backbone'] == 'unet_sd_c' + + +def test_build_method_opts_fluoresfm_custom_prompt(): + opts = _build_method_opts('fluoresfm', { + 'FluoResFM task': 'denoise', + 'FluoResFM text prompt': 'a custom prompt', + }) + assert opts['prompt'] == 'a custom prompt' + + +def test_build_method_opts_fluoresfm_backbone_selection(): + opts = _build_method_opts('fluoresfm', { + 'FluoResFM task': 'super-resolution', + 'FluoResFM backbone': 'dfcan', + }) + assert opts['backbone'] == 'dfcan' + assert opts['task'] == 'super-resolution' + + +def test_clip_to_dtype_uint16_clips_and_casts(): + image = np.array([-5.0, 70000.0, 100.0], dtype=np.float32) + clipped = _clip_to_dtype(image, np.uint16) + + assert clipped.dtype == np.uint16 + assert clipped[0] == 0 + assert clipped[1] == 65535 + assert clipped[2] == 100 + + +def test_clip_to_dtype_handles_nan_and_inf(): + image = np.array([np.nan, np.inf, -np.inf, 5.0], dtype=np.float32) + clipped = _clip_to_dtype(image, np.uint8) + + assert not np.isnan(clipped).any() + assert np.isfinite(clipped).all() + + +def test_clip_to_dtype_none_dtype_passthrough(): + image = np.array([1.5, 2.5], dtype=np.float32) + result = _clip_to_dtype(image, None) + assert result.dtype == np.float32 + np.testing.assert_allclose(result, [1.5, 2.5])