Add fluorescence illumination-correction and image-restoration workers#151
Open
arjunrajlab wants to merge 8 commits into
Open
Add fluorescence illumination-correction and image-restoration workers#151arjunrajlab wants to merge 8 commits into
arjunrajlab wants to merge 8 commits into
Conversation
…kers Detailed spec for two new Image Processing workers: illumination_correction (BaSiC, CIDRE-style, CellProfiler-style, flat/dark-field, destripe, plus an EVEN-style QC report) and image_restoration (Noise2Void/N2V2, Cellpose3, ZS-DeconvNet, FluoResFM). Covers interfaces, per-algorithm dispatch, dependencies, Dockerfiles, mock-based tests, and docs. Worker implementations follow in subsequent commits. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011brXFBv4hGK5hyheaGSMC7
Two new Image Processing workers, each with an algorithm-selector dropdown. illumination_correction (CPU): BaSiC (via BaSiCPy), a CIDRE-style retrospective gain/offset estimator, a CellProfiler-style illumination function (regular + background modes), classical flat/dark-field reference correction, and wavelet-FFT destriping (via pystripe, a classical stand-in for the unpackaged DL method SSCOR). Optional lightweight QC metrics report (a stand-in for EVEN's ML evaluation framework). Retrospective methods fit per-channel across the full XY/Z/Time collection. Gain surfaces are floored to bound amplification. image_restoration (GPU + CPU fallback): reference-free / self-supervised / pretrained methods only, since paired clean ground truth is usually unavailable. Noise2Void/N2V2 (via CAREamics, self-supervised), Cellpose3 restoration (pretrained), a zero-shot denoise+deconvolution pipeline inspired by ZS-DeconvNet, and FluoResFM (pretrained text-prompted foundation model, experimental; weights runtime-optional via FLUORESFM_WEIGHTS). Supervised CARE/3D-RCAN intentionally excluded. Both workers: lazy heavy imports (fast interface path, mock-testable), dtype preservation with NaN/inf guards, graceful sendError on missing deps/weights, unselected channels pass through unchanged. Registered in docker-compose.yml (build + test services) and REGISTRY.md. Mock-based test suites (17 + 31 tests) run natively without the heavy algorithm libraries installed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011brXFBv4hGK5hyheaGSMC7
Adds a new 'sscor' method to the Illumination Correction worker, alongside the existing classical pystripe 'destripe' method (which stays as the fast, no-weights option). SSCOR (github.com/lxxcontinue/SSCOR, Nat. Commun. 2023) is a CycleGAN-family codebase with no importable API, so it's vendored at a pinned commit and driven via subprocess to its restore.py CLI (mirroring how deconwolf shells out to its binary). This exposes SSCOR's inference stage with a user-supplied trained generator checkpoint: weights are not baked into the image (distributed via Google Drive, and a checkpoint requires an image-specific offline self-training stage), so they are resolved at runtime from the SSCOR_WEIGHTS env var with a graceful sendError if unset. GPU used when available, CPU fallback with a warning. SSCOR operates in 8-bit, so each frame is min/max-scaled to uint8 and rescaled back (a lossy round-trip inherent to SSCOR, documented). Missing-weights handling is an upfront check in compute() before the channel loop, matching the flatfield reference-check pattern and the correct_*() (stack, diagnostics) tuple contract. Dockerfile clones the pinned repo and installs its pix2pix-pipeline deps (torchvision/dominate/visdom/opencv/matplotlib/wandb; torch reused from basicpy). Tests (19 total) cover sscor dispatch with weights and the no-weights error path, running natively via lazy imports. REGISTRY.md and the design spec updated. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011brXFBv4hGK5hyheaGSMC7
…checkpoint name Adds a 'self-train' SSCOR mode (via a new 'SSCOR mode' select) that runs SSCOR's faithful per-image pipeline with no pre-supplied checkpoint: for each frame it samples striped/stripe-free patches from the image itself (sample_stripe.py for horizontal/vertical, sample_stripe_2.py for grid, controlled by new stripe direction/count/grid-direction/epochs params), trains a fresh CycleGAN (train.py, --save_epoch_freq == n_epochs so latest_net_G_A.pth is written once), then restores (restore.py) against that model. Training runs per frame on a GPU and is slow, so a one-time sendWarning is emitted; no SSCOR_WEIGHTS is required in this mode. The existing 'pretrained' mode (env SSCOR_WEIGHTS) remains the default and its upfront weights check is now gated on mode. Also fixes a checkpoint-naming bug in the pretrained path: SSCOR's model sets model_names=['G_A'] at inference, so load_networks loads latest_net_G_A.pth. The pretrained branch now stages the user checkpoint under that name (was latest_net_G.pth, which would have failed to load). Tests updated for the new interface fields and a self-train-needs-no-weights case (20 passing, native/mock-based via lazy imports). Docs and the design spec updated to describe both SSCOR modes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011brXFBv4hGK5hyheaGSMC7
…facts These .pyc files were inadvertently staged from a local pytest run; no other worker tracks compiled bytecode. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011brXFBv4hGK5hyheaGSMC7
Prevents __pycache__/*.pyc and .pytest_cache from being accidentally committed when running worker tests locally. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011brXFBv4hGK5hyheaGSMC7
…toration Correctness: - SSCOR: pass --load_size/--crop_size == patch_size to both restore.py invocations. restore.py defaults them to 256 and writes the network output into a patch_size-sized slice, so any non-default "SSCOR patch size" crashed with a shape mismatch (pretrained and self-train restore steps). - illumination compute(): wrap the per-method correction call in try/except so library/subprocess failures (BaSiC on a degenerate stack, a SSCOR subprocess error) send a structured error instead of a raw traceback (spec principle 5). - image_restoration compute(): move np.stack inside the try/except (heterogeneous per-frame shapes now error gracefully) and reject a restorer returning a different frame count than submitted (prevents silent frame drop/misalignment, e.g. from CAREamics predict()). - n2v: fix inverted patch-size clamp (max(16, min(patch,H,W)) could force the patch above a sub-16px image); clamp to the image dimension instead. - Cellpose3 / FluoResFM: rescale each restored frame's dynamic range back onto the source frame's before dtype cast. These pretrained models emit ~0-1 normalized output that _clip_to_dtype would otherwise clip to near-black. - ZS/FluoResFM: replace `float(max) or 1.0` normalization guards with an explicit >0 check (a negative max is truthy and would invert the frame). Build: - illumination Dockerfile: install BaSiCPy with --no-deps plus its runtime deps by hand. BaSiCPy 2.0.0 hard-pins scipy<1.13, which conflicts with the base env's numpy>=2.0 (needs scipy>=1.13) and would break the conda scikit-image/ large_image C extensions. Install CPU-only torch/torchvision explicitly to avoid pulling the multi-GB CUDA runtime into this CPU worker, and move the entrypoint COPY after the dependency layers for build-cache friendliness. Docs: - REGISTRY.md: correct the summary Total (56 -> 58) after adding two workers. All tests pass (illumination 20, restoration 31). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011brXFBv4hGK5hyheaGSMC7
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b2fcd3477c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Comment on lines
+596
to
+597
| from methods.fluoresfm import build_model, load_checkpoint | ||
| from packages.text_embedding import embed_text |
There was a problem hiding this comment.
When Method='fluoresfm' is selected and weights are provided, this import path still aborts before inference: the Dockerfile pins qiqi-lu/fluoresfm to v1.0.1, whose methods/ directory contains back_projector.py, convolution.py, and deconvolution.py, but no methods/fluoresfm.py, and packages/ does not contain text_embedding.py (see https://github.com/qiqi-lu/fluoresfm/tree/v1.0.1/methods). As a result every FluoResFM job will hit the ImportError path and report the method unavailable; this should either call the actual upstream inference entry points or the method should not be exposed as runnable.
Useful? React with 👍 / 👎.
| # 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" triton tifffile |
There was a problem hiding this comment.
This unbounded cellpose>=3 install now resolves to Cellpose 4.x in fresh builds, but the Cellpose 4 release notes state that cellpose.denoise is not available (https://github.com/MouseLand/cellpose/releases). In that environment restore_cellpose3() will fail at from cellpose import denoise, and the advertised Cellpose3 restoration method will always abort at runtime; pin the production and dev installs to a Cellpose 3 restoration release such as 3.1.1.2 or otherwise install a version that still ships cellpose.denoise.
Useful? React with 👍 / 👎.
… API)
Codex flagged two P2 issues on the image_restoration worker; both confirmed
against the upstream repos and fixed.
1. Cellpose resolved to 4.x, which removed `cellpose.denoise`, so
restore_cellpose3 always failed at import. Pin `cellpose>=3,<4` in both
Dockerfiles and environment.yml so the restoration models remain available.
2. restore_fluoresfm imported modules that do not exist in qiqi-lu/fluoresfm
(`methods.fluoresfm`, `packages.text_embedding`), so the method always hit
the ImportError path and was non-runnable. Rewrote it against the repo's
real inference API, transcribed from 3_0_test_it2i.py / 3_1_test_i2i.py:
- Default backbone `unet_sd_c` is the true text-conditioned FluoResFM
(models.unet_sd_c.UNetModel conditioned on a BiomedCLIP text embedding
from models.biomedclip_embedder.BiomedCLIPTextEmbedder), with the repo's
percentile normalization and reflect-pad-to-/8.
- Also wires the three real non-text baselines care/dfcan/unifmir from
3_1_test_i2i.py (prompt ignored; unifmir uses a task index).
- New `FluoResFM backbone` select param (must match the checkpoint).
- Checkpoint via FLUORESFM_WEIGHTS and the BiomedCLIP embedder via a new
FLUORESFM_EMBEDDER_DIR, both runtime-supplied with accurate sendError
gating (no fabricated APIs, no giant weights baked into the image).
- Dockerfiles add the embedder deps (open_clip_torch/transformers/
huggingface_hub/torchinfo/timm) and document weight/embedder sources.
Honest simplifications documented in IMAGE_RESTORATION.md: whole-frame forward
(no sliding-window stitcher), free-text prompts are out-of-distribution vs. the
repo's structured prompts, output grid kept identical to input. FluoResFM
remains experimental and is not built/GPU-run in CI.
Tests: image_restoration 34, illumination_correction 20 (all passing, mock-based).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011brXFBv4hGK5hyheaGSMC7
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds two new Image Processing workers, each a single worker with a
Methoddropdown selecting among algorithms. Both follow the existing image-processing pattern (histogram_matching/deconwolf): read all frames, process selected channels, upload a new TIFF to Girder, and attach provenance metadata. Unselected channels pass through unchanged.illumination_correction— corrects illumination/shading/vignetting/stripingImplements the full illumination shortlist:
basiccidrecellprofilerregularacross-batch /backgroundper-image; reimplemented, no CellProfiler dependency).flatfield(raw − dark) / (flat − dark)with reference frames supplied by XY coordinate.destripepystripe— fast, no weights needed.sscorpretrained(runtime checkpoint viaSSCOR_WEIGHTS) andself-train(samples patches from each frame, trains a fresh CycleGAN, then restores — needs no checkpoint).Plus an optional lightweight QC metrics report (a stand-in for EVEN, which is an ML evaluation framework rather than a corrector).
image_restoration— denoise / deblur / restore (GPU, CPU fallback)Deliberately biased toward reference-free / self-supervised / pretrained methods, since paired clean ground truth is usually unavailable:
n2vcellpose3zs_deconvnetfluoresfmFLUORESFM_WEIGHTS).Supervised methods (CARE/CSBDeep, 3D-RCAN) and time-lapse-specific denoisers (DeepCAD-RT/FAST/TeD) were intentionally excluded — noted as possible future additions.
Design notes
histogram_matching's test approach.sendErrorwith setup instructions rather than a crash. GPU methods fall back to CPU with a warning.FLUOR_CORRECTION_WORKERS_SPEC.md.Testing
illumination_correction: 20 tests ·image_restoration: 31 tests — all passing natively (mock-based).restore.pynow receives--load_size/--crop_size == patch_size(default 256 would crash for any other patch size).scipy<1.13vs the base's numpy 2.x — installed with--no-deps+ hand-listed runtime deps, plus CPU-only torch to avoid multi-GB CUDA bloat.compute()s; frame-count guard against silent drops; n2v patch-size clamp fix; Cellpose3/FluoResFM output rescaled to source range (avoids near-black output); explicit>0normalization guards.Known limitations (documented in the worker docs)
self-traintrains a fresh model per frame — faithful to the method but slow; GPU strongly recommended.flatfield, a reference frame on the corrected channel is itself corrected in the output; retrospective methods need ≥3 frames per channel to be reliable (a warning is emitted below that).🤖 Generated with Claude Code
https://claude.ai/code/session_011brXFBv4hGK5hyheaGSMC7
Generated by Claude Code