|
| 1 | +#!/usr/bin/env python |
| 2 | +''' |
| 3 | +Module : compute_cellpose_norm |
| 4 | +Description : Compute per-channel global intensity bounds (low, high) for a |
| 5 | + SpatialData .zarr image, to be handed to Cellpose as its |
| 6 | + `normalize={"lowhigh": ...}` argument. |
| 7 | +
|
| 8 | + Cellpose normalises every patch to that patch's own 1-99 |
| 9 | + percentile. Identical cells therefore land on different scales in |
| 10 | + a dense versus a sparse patch, which is the source of the residual |
| 11 | + tiling grid once sopa's CLAHE/gaussian stages are disabled. |
| 12 | + Computing one (low, high) pair per channel from the WHOLE image |
| 13 | + and applying it to every patch removes that patch dependence. |
| 14 | +
|
| 15 | + Bounds are emitted in the SAME channel order they are passed on |
| 16 | + the command line, which must match the order in which the caller |
| 17 | + passes `--channels` to `sopa segmentation cellpose` (membrane then |
| 18 | + nucleus for whole-cell, nucleus alone for nuclear-only). The |
| 19 | + output is a JSON/Python-literal list of [low, high] pairs, ready to |
| 20 | + splice into `--method-kwargs`. |
| 21 | +Copyright : (c) WEHI SODA Hub, 2025 |
| 22 | +License : MIT |
| 23 | +Maintainer : Marek Cmero (@mcmero) |
| 24 | +Portability : POSIX |
| 25 | +''' |
| 26 | +import json |
| 27 | +from pathlib import Path |
| 28 | +from typing import Annotated, List, Optional |
| 29 | + |
| 30 | +import numpy as np |
| 31 | +import typer |
| 32 | + |
| 33 | +app = typer.Typer() |
| 34 | + |
| 35 | + |
| 36 | +def _get_image(zarr_path: str): |
| 37 | + """Return the single image element of a SpatialData store as a DataArray or |
| 38 | + DataTree (multiscale). Uses sopa's helper when available, else falls back to |
| 39 | + the sole image in the store.""" |
| 40 | + import spatialdata as sd |
| 41 | + |
| 42 | + sdata = sd.read_zarr(zarr_path) |
| 43 | + try: |
| 44 | + from sopa.utils import get_spatial_image |
| 45 | + |
| 46 | + return get_spatial_image(sdata) |
| 47 | + except Exception: |
| 48 | + keys = list(sdata.images.keys()) |
| 49 | + if len(keys) != 1: |
| 50 | + raise ValueError( |
| 51 | + f"Expected exactly one image in {zarr_path}, found {keys}" |
| 52 | + ) |
| 53 | + return sdata.images[keys[0]] |
| 54 | + |
| 55 | + |
| 56 | +def _pick_level(image, max_pixels: int): |
| 57 | + """From a (possibly multiscale) image, return a single DataArray to compute |
| 58 | + percentiles on. For a multiscale pyramid, pick the finest level whose Y*X is |
| 59 | + at or below `max_pixels` (falling back to the coarsest if even that is |
| 60 | + larger). p1/p99 are robust to this downsampling, and it bounds memory on |
| 61 | + large slides.""" |
| 62 | + # Multiscale images are DataTree-like: they expose scale groups. |
| 63 | + scales = None |
| 64 | + try: |
| 65 | + scales = [k for k in image.keys() if str(k).startswith("scale")] |
| 66 | + except Exception: |
| 67 | + scales = None |
| 68 | + |
| 69 | + if scales: |
| 70 | + # scale0 is full resolution; scales are ordered coarser as the index grows. |
| 71 | + def _yx(scale): |
| 72 | + arr = image[scale]["image"] |
| 73 | + return int(arr.sizes["y"]) * int(arr.sizes["x"]) |
| 74 | + |
| 75 | + scales_sorted = sorted(scales, key=lambda s: _yx(s), reverse=True) |
| 76 | + chosen = scales_sorted[-1] # coarsest as a safe default |
| 77 | + for s in scales_sorted: # finest first |
| 78 | + if _yx(s) <= max_pixels: |
| 79 | + chosen = s |
| 80 | + break |
| 81 | + return image[chosen]["image"] |
| 82 | + |
| 83 | + # Single-scale DataArray. |
| 84 | + return image |
| 85 | + |
| 86 | + |
| 87 | +def _percentile_bounds(values: np.ndarray, plow: float, phigh: float, |
| 88 | + max_pixels: int) -> List[float]: |
| 89 | + """Compute [low, high] percentile bounds from a 2D intensity array. |
| 90 | +
|
| 91 | + Pure numpy so it is unit-testable without xarray/spatialdata. Stride- |
| 92 | + subsamples a single-scale array far above the cap to bound memory, and |
| 93 | + guards a flat channel: Cellpose divides by (high - low), so equal bounds |
| 94 | + would give inf/NaN.""" |
| 95 | + values = np.asarray(values) |
| 96 | + n = values.size |
| 97 | + if n > max_pixels: |
| 98 | + stride = int(np.ceil(np.sqrt(n / max_pixels))) |
| 99 | + values = values[::stride, ::stride] |
| 100 | + |
| 101 | + lo, hi = np.percentile(values.ravel(), [plow, phigh]) |
| 102 | + lo, hi = float(lo), float(hi) |
| 103 | + if hi <= lo: |
| 104 | + hi = lo + 1.0 |
| 105 | + return [lo, hi] |
| 106 | + |
| 107 | + |
| 108 | +def _bounds_for_channel(arr, channel: str, plow: float, phigh: float, |
| 109 | + max_pixels: int) -> List[float]: |
| 110 | + """Compute [low, high] percentile bounds for one named channel.""" |
| 111 | + a = arr.sel(c=channel).data # dask or numpy, shape (Y, X) |
| 112 | + return _percentile_bounds(np.asarray(a), plow, phigh, max_pixels) |
| 113 | + |
| 114 | + |
| 115 | +@app.command() |
| 116 | +def main( |
| 117 | + zarr: Annotated[str, typer.Argument(help="Path to the SpatialData .zarr store")], |
| 118 | + channel: Annotated[ |
| 119 | + List[str], |
| 120 | + typer.Option( |
| 121 | + "--channel", |
| 122 | + help="Channel name to compute bounds for. Repeat in the SAME order " |
| 123 | + "the channels are passed to `sopa segmentation cellpose` " |
| 124 | + "(membrane then nucleus for whole-cell).", |
| 125 | + ), |
| 126 | + ], |
| 127 | + percentiles: Annotated[ |
| 128 | + str, |
| 129 | + typer.Option(help="Low,high percentiles for the bounds."), |
| 130 | + ] = "1,99", |
| 131 | + max_pixels: Annotated[ |
| 132 | + int, |
| 133 | + typer.Option( |
| 134 | + help="Compute on the finest multiscale level with at most this many " |
| 135 | + "pixels (single-scale images above it are stride-subsampled).", |
| 136 | + ), |
| 137 | + ] = 8_000_000, |
| 138 | + as_method_kwargs: Annotated[ |
| 139 | + bool, |
| 140 | + typer.Option( |
| 141 | + "--as-method-kwargs", |
| 142 | + help="Emit the full Cellpose method-kwargs dict " |
| 143 | + '`{"normalize": {"lowhigh": ...}}` instead of the bare list of ' |
| 144 | + "pairs, ready to pass to `sopa segmentation cellpose --method-kwargs`.", |
| 145 | + ), |
| 146 | + ] = False, |
| 147 | + output: Annotated[ |
| 148 | + Optional[str], |
| 149 | + typer.Option("--output", "-o", help="Write JSON here instead of stdout."), |
| 150 | + ] = None, |
| 151 | +): |
| 152 | + """Emit per-channel [low, high] bounds as a JSON list of pairs, ordered to |
| 153 | + match the given `--channel` order.""" |
| 154 | + plow, phigh = (float(x) for x in percentiles.split(",")) |
| 155 | + |
| 156 | + image = _get_image(zarr) |
| 157 | + arr = _pick_level(image, max_pixels) |
| 158 | + |
| 159 | + available = {str(c) for c in arr.coords["c"].values} |
| 160 | + missing = [c for c in channel if c not in available] |
| 161 | + if missing: |
| 162 | + raise ValueError( |
| 163 | + f"Channel(s) {missing} not found in image. Available: {sorted(available)}" |
| 164 | + ) |
| 165 | + |
| 166 | + lowhigh = [ |
| 167 | + _bounds_for_channel(arr, c, plow, phigh, max_pixels) for c in channel |
| 168 | + ] |
| 169 | + |
| 170 | + payload = {"normalize": {"lowhigh": lowhigh}} if as_method_kwargs else lowhigh |
| 171 | + text = json.dumps(payload) |
| 172 | + if output: |
| 173 | + Path(output).write_text(text) |
| 174 | + else: |
| 175 | + typer.echo(text) |
| 176 | + |
| 177 | + |
| 178 | +if __name__ == "__main__": |
| 179 | + app() |
0 commit comments