"""
Streaming mask pipeline (v2).
Replaces the multi-copy disk chain that ``preprocess_generate_masks``
has run since day one:
originals
→ renamed + split into channel folders
→ orig/ backup
→ per-channel npy
→ batch npz on disk
→ cellpose → per-field mask npy
→ concatenated into merged/
…with a two-pass streaming pipeline that keeps only what the
downstream measure module actually reads:
Pass 1 — assemble
walk originals, parse metadata regex, build one npy stack per
field with all image channels in the C axis. Emit
``filename_map.csv`` recording every original → stack mapping.
Pass 2 — segment
stream the plate in batches of N fields, hand each batch to
Cellpose, append the mask channels to the SAME stack file.
Each batch is written to a scratch NPZ on the way through and
deleted again unless ``keep_npz=True``.
Output — ``merged/`` folder holds one file per field, each shape
``(H, W, C_image + C_mask)`` in uint16, plus:
channel_order.json {"channels": [...]}
filename_map.csv original path, plate/well/field/…, stack idx
Public API::
from spacr.pipeline_v2 import (
FilenameMapper, stream_originals_to_stack,
stream_masks_from_stack, run_v2,
)
# High-level (one call):
run_v2(src_folder, channels=(0,1,2,3), model="cyto", diameter=60)
# Low-level (two passes, run each explicitly):
mapper = FilenameMapper.discover(src_folder,
metadata_type="cellvoyager")
stacks = stream_originals_to_stack(src_folder, mapper, channels=(0,1,2,3))
stream_masks_from_stack(stacks, model="cyto", diameter=60)
This module is opt-in for one release cycle. Once the follow-up commit
wires it as the default in :func:`spacr.core.preprocess_generate_masks`
the whole disk chain above collapses to ``merged/`` alone.
"""
from __future__ import annotations
import csv
import json
import logging
import re
import shutil
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple
import numpy as np
from .logging_util import Timer, timed
LOG = logging.getLogger("spacr.pipeline_v2")
@dataclass
[docs]
class FilenameRecord:
"""One entry in the filename map.
:ivar original_path: absolute path to the source image on disk.
:ivar plate: plate id parsed from the filename.
:ivar well: well id parsed from the filename.
:ivar field: field index parsed from the filename.
:ivar channel: channel index parsed from the filename.
:ivar time: time index parsed from the filename (defaults to 1).
:ivar z: z-slice index parsed from the filename (defaults to 1).
:ivar stack_field_id: the ``field`` id used in ``merged/stack_<X>.npy``.
"""
original_path: str
plate: str
well: str
field: int
channel: int
time: int = 1
z: int = 1
stack_field_id: str = ""
[docs]
class FilenameMapper:
"""Walks a folder of microscopy images, parses each filename's
metadata via a regex, and records the mapping to a per-plate CSV.
The CSV is written next to the ``merged/`` folder (at the plate
root) so users can Excel-open ``filename_map.csv`` and see the
original path of every image in the run.
:param records: parsed filename records in file-system order.
:param metadata_type: name of the metadata convention used to parse them.
:param regex: regular-expression source used for parsing.
:ivar records: list of :class:`FilenameRecord` in file-system order.
:ivar metadata_type: which regex was used (``"cellvoyager"`` /
``"yokogawa"`` / ``"custom"``).
:ivar regex: compiled regex pattern that matched.
"""
def __init__(self, records: List[FilenameRecord],
metadata_type: str, regex: str):
"""Store parsed filename records and the metadata rule that made them."""
self.records = records
self.metadata_type = metadata_type
self.regex = regex
@classmethod
[docs]
def discover(cls, src: Path,
metadata_type: str = "auto",
custom_regex: Optional[str] = None,
exts: Sequence[str] = (".tif", ".tiff", ".png",
".jpg", ".jpeg")) -> "FilenameMapper":
"""Scan ``src`` for images + parse each name with the metadata
regex. Falls back through ``cellvoyager`` → ``yokogawa`` on
``metadata_type="auto"``.
:param src: folder to scan (not recursive; we expect images at
the top level as the current spacr layout does).
:param metadata_type: ``"auto"`` / ``"cellvoyager"`` /
``"yokogawa"`` / ``"custom"``. When ``"custom"``,
``custom_regex`` must be given.
:param custom_regex: user-supplied regex; required for
``metadata_type="custom"``.
:param exts: image file extensions to include.
:returns: a populated :class:`FilenameMapper`.
:raises ValueError: when no images are found or no regex fits.
"""
src = Path(src)
files = sorted(
p for p in src.iterdir()
if p.is_file() and p.suffix.lower() in exts
and not p.name.startswith(".")
)
if not files:
raise ValueError(f"no images found in {src}")
pattern, chosen = _resolve_regex(metadata_type, files, custom_regex)
recs: List[FilenameRecord] = []
rx = re.compile(pattern)
for f in files:
m = rx.match(f.name)
if m is None:
LOG.warning("filename didn't match %s regex: %s",
chosen, f.name)
continue
g = m.groupdict()
recs.append(FilenameRecord(
original_path=str(f.resolve()),
plate=g.get("plateID") or g.get("plate") or "plate1",
well=g.get("wellID") or g.get("well") or "A01",
field=int(g.get("fieldID") or g.get("field") or 1),
channel=int(g.get("chanID") or g.get("channel") or 1),
time=int(g.get("timeID") or g.get("time") or 1),
z=int(g.get("sliceID") or g.get("z") or 1),
))
keys = {}
for r in recs:
k = (r.plate, r.well, r.field, r.time, r.z)
if k not in keys:
keys[k] = f"{len(keys):06d}"
r.stack_field_id = keys[k]
LOG.info("discovered %d images grouped into %d fields (regex: %s)",
len(recs), len(keys), chosen)
return cls(recs, chosen, pattern)
[docs]
def save_csv(self, path: Path) -> Path:
"""Write the mapping to ``path`` as a CSV that Excel opens
cleanly. One row per (original image, resulting stack slot).
:param path: destination CSV path; its parent directory is created.
"""
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
cols = ["original_path", "plate", "well", "field", "channel",
"time", "z", "stack_field_id"]
with open(path, "w", newline="") as f:
w = csv.writer(f)
w.writerow(cols)
for r in self.records:
w.writerow([getattr(r, c) for c in cols])
(path.with_suffix(".json")).write_text(json.dumps({
"metadata_type": self.metadata_type,
"regex": self.regex,
"n_records": len(self.records),
}, indent=2))
return path
@classmethod
[docs]
def load_csv(cls, path: Path) -> "FilenameMapper":
"""Rehydrate a mapper from a previously-saved CSV.
:param path: mapping CSV previously written by :meth:`save_csv`.
"""
path = Path(path)
recs: List[FilenameRecord] = []
with open(path) as f:
for row in csv.DictReader(f):
recs.append(FilenameRecord(
original_path=row["original_path"],
plate=row["plate"], well=row["well"],
field=int(row["field"]), channel=int(row["channel"]),
time=int(row["time"]), z=int(row["z"]),
stack_field_id=row["stack_field_id"],
))
meta_path = path.with_suffix(".json")
if meta_path.exists():
meta = json.loads(meta_path.read_text())
return cls(recs, meta.get("metadata_type", "?"),
meta.get("regex", ""))
return cls(recs, "?", "")
[docs]
def by_field(self) -> Dict[str, List[FilenameRecord]]:
"""Group records by ``stack_field_id`` — one entry per field, with
one record per channel inside."""
out: Dict[str, List[FilenameRecord]] = {}
for r in self.records:
out.setdefault(r.stack_field_id, []).append(r)
return out
[docs]
def field_ids(self) -> List[str]:
"""Return the sorted list of unique ``stack_field_id`` values."""
return sorted(self.by_field().keys())
_CELLVOYAGER = (
r"(?P<plateID>.*)_(?P<wellID>.*)_T(?P<timeID>.*)F(?P<fieldID>.*)"
r"L(?P<laserID>..)A(?P<AID>..)Z(?P<sliceID>.*)C(?P<chanID>.*)"
r"\.(?:tif|tiff|png|jpg|jpeg)$"
)
_YOKOGAWA = (
r"(?P<plateID>.*)_(?P<wellID>[A-Z]\d{2})_"
r"T(?P<timeID>\d{4})F(?P<fieldID>\d{3})"
r"L(?P<laserID>\d{2})A(?P<AID>\d{2})Z(?P<sliceID>\d{2})C(?P<chanID>\d{2})"
r"\.(?:tif|tiff)$"
)
def _resolve_regex(metadata_type: str, files: List[Path],
custom_regex: Optional[str]) -> Tuple[str, str]:
"""Pick a regex + return (pattern, chosen_label)."""
if metadata_type == "custom":
if not custom_regex:
raise ValueError("metadata_type='custom' needs custom_regex")
return custom_regex, "custom"
candidates: List[Tuple[str, str]] = []
if metadata_type == "cellvoyager":
candidates.append((_CELLVOYAGER, "cellvoyager"))
elif metadata_type == "yokogawa":
candidates.append((_YOKOGAWA, "yokogawa"))
else:
candidates = [(_CELLVOYAGER, "cellvoyager"),
(_YOKOGAWA, "yokogawa")]
for pattern, name in candidates:
rx = re.compile(pattern)
if all(rx.match(f.name) for f in files):
return pattern, name
best_pattern, best_name, best_hits = candidates[0][0], candidates[0][1], -1
for pattern, name in candidates:
rx = re.compile(pattern)
hits = sum(1 for f in files if rx.match(f.name))
if hits > best_hits:
best_pattern, best_name, best_hits = pattern, name, hits
LOG.warning("no regex matched every file; best fit was %s (%d/%d)",
best_name, best_hits, len(files))
return best_pattern, best_name
@dataclass
[docs]
class StackFile:
"""One field's on-disk stack: ``merged/stack_<id>.npy`` with
shape ``(H, W, C)``.
Populated by :func:`stream_originals_to_stack` before Cellpose
runs (C = image channels only). After :func:`stream_masks_from_stack`
the same file has additional mask channels appended.
:ivar field_id: stable field identifier used in the stack filename.
:ivar path: path to the on-disk NumPy stack.
:ivar shape: ``(height, width, channels)`` shape at write time.
:ivar channels: human-readable channel names in array order.
"""
field_id: str
path: Path
shape: Tuple[int, int, int]
channels: List[str]
@timed
[docs]
def stream_originals_to_stack(
src: Path,
mapper: FilenameMapper,
channels: Sequence[int] = (0, 1, 2, 3),
channel_names: Optional[Sequence[str]] = None,
dst: Optional[Path] = None,
) -> List[StackFile]:
"""Write one ``merged/stack_<field>.npy`` per field.
Reads originals directly (no rename-into-channel-folders step),
stacks the selected channels along the C axis, and writes one
npy per field. Also emits a ``channel_order.json`` sidecar
describing which C-index holds which channel.
:param src: plate folder containing the original images.
:param mapper: :class:`FilenameMapper` produced from ``src``.
:param channels: which channel numbers (as parsed from filenames)
to include, in the order they should occupy the C axis.
:param channel_names: human names for those channels (must match
``channels`` length). Default: ``["ch0", "ch1", …]``.
:param dst: override the output folder; defaults to ``<src>/merged``.
:returns: list of :class:`StackFile`, one per field written.
"""
src = Path(src)
dst = Path(dst) if dst else src / "merged"
dst.mkdir(parents=True, exist_ok=True)
if channel_names is None:
channel_names = [f"ch{c}" for c in channels]
assert len(channel_names) == len(channels), (
"channel_names must match channels length"
)
by_field = mapper.by_field()
written: List[StackFile] = []
for field_id, recs in by_field.items():
by_ch = {r.channel: r for r in recs}
ref_shape = None
ref_dtype = np.uint16
read_cache: dict = {}
for ch in channels:
rec = by_ch.get(ch)
if rec is not None:
plane = _read_plane(rec.original_path)
read_cache[ch] = plane
ref_shape = plane.shape
ref_dtype = plane.dtype
break
if ref_shape is None:
ref_shape = (256, 256)
planes: List[np.ndarray] = []
for ch in channels:
rec = by_ch.get(ch)
if rec is None:
LOG.warning("field %s missing channel %d — inserting zeros",
field_id, ch)
planes.append(np.zeros(ref_shape, dtype=ref_dtype))
continue
plane = read_cache.get(ch)
if plane is None:
plane = _read_plane(rec.original_path)
planes.append(plane)
stack = np.stack(planes, axis=-1).astype(np.uint16)
out_path = dst / f"stack_{field_id}.npy"
np.save(out_path, stack)
written.append(StackFile(
field_id=field_id, path=out_path, shape=stack.shape,
channels=list(channel_names),
))
(dst / "channel_order.json").write_text(json.dumps({
"image_channels": list(channel_names),
"mask_channels": [],
"shape_H_W_C": "final shape is (H, W, C_image + C_mask)",
}, indent=2))
mapper.save_csv(src / "filename_map.csv")
LOG.info("wrote %d field stacks under %s + filename_map.csv",
len(written), dst)
return written
def _record_cellpose_hash(model, model_name: str) -> None:
"""Best-effort — fingerprint the Cellpose checkpoint and record it
on the currently-open :class:`spacr.run_journal.Run`, if any."""
try:
ckpt_paths = []
for attr in ("pretrained_model", "cp"):
obj = getattr(model, attr, None)
if obj is None:
continue
if isinstance(obj, (list, tuple)):
ckpt_paths.extend(obj)
else:
nested = getattr(obj, "pretrained_model", None)
if nested is not None:
if isinstance(nested, (list, tuple)):
ckpt_paths.extend(nested)
else:
ckpt_paths.append(nested)
ckpt_paths = [Path(p) for p in ckpt_paths
if p and Path(p).is_file()]
if not ckpt_paths:
return
try:
from .run_journal import current_run
run = current_run()
if run is None:
return
for ckpt in ckpt_paths:
run.record_model(model_name, ckpt)
except Exception as exc:
LOG.warning("model %r was not recorded in the run journal (%s); "
"this run's manifest will not name the weights it "
"used.", model_name, exc)
except Exception as exc:
LOG.warning("could not work out which checkpoint %r is using (%s); "
"no model provenance was recorded for this run.",
model_name, exc)
def _read_plane(path: str) -> np.ndarray:
"""Read a single 2-D image plane (H, W) as uint16."""
p = Path(path)
suf = p.suffix.lower()
if suf in (".tif", ".tiff"):
import tifffile
arr = tifffile.imread(str(p))
else:
from PIL import Image
arr = np.array(Image.open(str(p)))
if arr.ndim == 3:
arr = arr[..., 0]
return arr.astype(np.uint16, copy=False)
def _as_hwc(arr: np.ndarray) -> np.ndarray:
"""A loaded field as (H, W, C), whatever singleton axes it was saved with.
A stack written as a bare plane is (H, W) on disk, and both ends of the
mask pass need a channel axis: Cellpose is called with
``channel_axis=-1`` and refuses a 2-D image outright, and the write-back
concatenates the mask onto this same array, which numpy refuses when the
two disagree on rank. Promoting once, at load, keeps one shape convention
for the whole pass instead of a squeeze that puts the mismatch off until
the batch is already segmented.
"""
arr = np.asarray(arr)
if arr.ndim == 3:
return arr
squeezed = arr.squeeze()
if squeezed.ndim == 2:
return squeezed[..., np.newaxis]
return squeezed
def _cellpose_channel_indices(
channels_for_cellpose: Sequence[int], n_channels: int
) -> Tuple[int, ...]:
"""Resolve requested channels into persisted C-axis positions once."""
if n_channels <= 0:
raise ValueError("a V2 field has no persisted intensity channels")
indices = [
int(channel) % n_channels for channel in channels_for_cellpose
]
return tuple(dict.fromkeys(indices)) or (0,)
def _backend_mask_settings(settings: Optional[Dict[str, Any]],
object_type: str, *, diameter: Optional[float],
flow_threshold: float,
cellprob_threshold: float) -> Dict[str, Any]:
"""The settings a prefixed model's masks function reads, as V1 has them.
V1 hands the run's settings to :func:`spacr.object._cellpose3_masks`
and the object's magnification diameter as its default. The same run's
settings arrive here as ``postprocess_settings``; the thresholds and
diameter this pass was called with win over them, and the magnification
default is kept under ``_default_diameter``.
:param settings: the run's settings, or None.
:param object_type: the object being segmented.
:param diameter: the diameter this pass was called with; None keeps the
object's own setting.
:param flow_threshold: this pass's flow threshold.
:param cellprob_threshold: this pass's cell probability threshold.
:returns: a new settings dict.
"""
merged = dict(settings or {})
merged[f"{object_type}_flow_threshold"] = float(flow_threshold)
merged[f"{object_type}_cellprob_threshold"] = float(cellprob_threshold)
if diameter is not None:
merged[f"{object_type}_diameter"] = diameter
default = 30.0
if merged.get("magnification") is not None:
from .utils import _get_diam
default = _get_diam(merged["magnification"], obj=object_type)
elif diameter:
default = float(diameter)
merged["_default_diameter"] = default
return merged
@timed
[docs]
def stream_masks_from_stack(
stacks: List[StackFile],
model_name: str = "cyto",
channels_for_cellpose: Sequence[int] = (0, 0),
diameter: Optional[float] = None,
batch_fields: int = 8,
mask_channel_name: str = "mask",
keep_npz: bool = False,
npz_dir: Optional[Path] = None,
cellprob_threshold: float = 0.0,
flow_threshold: float = 0.4,
min_size: int = 15,
resample: bool = True,
postprocess_settings: Optional[Dict[str, Any]] = None,
object_type: str = "cell",
illumination_session: Optional[Any] = None,
psf_session: Optional[Any] = None,
) -> List[StackFile]:
"""Batch the field stacks through Cellpose, then append the mask
channel(s) to the SAME npy files.
:param stacks: list produced by :func:`stream_originals_to_stack`.
:param model_name: Cellpose model to use (``"cyto"``, ``"nuclei"``, …).
:param channels_for_cellpose: C-axis indices selected out of each
assembled stack before it is handed to Cellpose — taken modulo the
channel count and de-duplicated in order. It is NOT forwarded as
Cellpose's ``channels=`` argument; ``[0, 0]`` therefore yields a
single plane, not a grayscale pair. The first selected channel is
the object's own channel for absolute mean-intensity filtering.
:param diameter: expected object diameter in px (None → Cellpose
auto).
:param batch_fields: how many field stacks to load into memory at
once. Larger = faster but more RAM.
:param mask_channel_name: human name to record for the appended
mask channel (default ``"mask"``).
:param keep_npz: the intermediate batch is compressed to an NPZ under
``npz_dir`` on every batch regardless; this flag only decides
whether that file and the scratch folder survive the run.
:param npz_dir: where to write the (optional) intermediate NPZ
files. Defaults to a scratch subfolder under the stack folder.
:param illumination_session: optional
:class:`spacr.illumination.SegmentationIlluminationSession`. Its
corrector sees private selected-channel copies immediately before
normalisation/Cellpose; persisted intensity planes and scratch NPZs
remain raw, and completion is recorded only after the combined stack
has been atomically replaced.
:param psf_session: optional captured PSF session. Unmixes each whole
field first when unmixing is on, then processes selected intensities
after illumination and before normalization, padding or Cellpose. Stored image channels stay raw; only the appended labels
depend on PSF processing.
:returns: the same list, with each :class:`StackFile.shape` /
``.channels`` updated to reflect the appended mask channel.
"""
if not stacks:
return stacks
from .utils import _resolve_cellpose_pretrained, _validated_intensity_bounds
intensity_settings = postprocess_settings or {}
minimum, maximum = _validated_intensity_bounds(
intensity_settings.get(f"{object_type}_min_intensity", 0),
intensity_settings.get(f"{object_type}_max_intensity", 0),
)
from .qt.mask_engine import filters_need_intensity, settings_filters
filter_by_raw_intensity = (minimum > 0 or maximum > 0 or filters_need_intensity(
settings_filters(intensity_settings, object_type)))
scratch = Path(npz_dir) if npz_dir else stacks[0].path.parent / "_scratch"
scratch.mkdir(parents=True, exist_ok=True)
try:
from .object import _prefixed_model_route
except ImportError as e:
raise RuntimeError(
"cellpose is required for v2 mask streaming: spacr.object, "
f"which routes a model setting to its backend, could not be "
f"imported ({e})"
) from e
route = _prefixed_model_route(model_name, postprocess_settings)
if route is not None:
from . import _segmentation_backends
backend, backend_masks = route
model = _segmentation_backends._load_backend(
backend, model_name=model_name, object_type=object_type)
backend_settings = _backend_mask_settings(
postprocess_settings, object_type, diameter=diameter,
flow_threshold=flow_threshold,
cellprob_threshold=cellprob_threshold)
else:
try:
from cellpose import models as cp_models # type: ignore
except Exception as e:
raise RuntimeError(
"cellpose is required for v2 mask streaming"
) from e
from .accelerator import cellpose_kwargs
pretrained = _resolve_cellpose_pretrained(
model_name, object_type=object_type)
model = cp_models.CellposeModel(
pretrained_model=pretrained,
**cellpose_kwargs(),
)
_record_cellpose_hash(model, model_name)
for batch_start in range(0, len(stacks), batch_fields):
batch = stacks[batch_start:batch_start + batch_fields]
with Timer(
f"v2.batch[{batch_start}:{batch_start + len(batch)}] "
f"load", logger="spacr.pipeline_v2",
):
loaded = [_as_hwc(np.load(s.path)) for s in batch]
npz_path = scratch / f"batch_{batch_start:04d}.npz"
np.savez_compressed(
npz_path,
**{s.field_id: arr for s, arr in zip(batch, loaded)},
)
selected_images: List[np.ndarray] = []
raw_intensity_per_field: List[np.ndarray] = []
for sf, arr in zip(batch, loaded):
indices = _cellpose_channel_indices(
channels_for_cellpose, arr.shape[-1])
if filter_by_raw_intensity:
raw_intensity_per_field.append(arr[..., indices[0]])
source = (psf_session.unmix(arr) if psf_session is not None
else arr)
selected = source[..., list(indices)]
if illumination_session is not None:
from .measure_hooks import PreprocessingContext
context = PreprocessingContext(
file_name=sf.path.name,
channels=indices,
settings=postprocess_settings or {},
)
selected = illumination_session.correct(
sf.field_id, selected, context)
if psf_session is not None:
selected = psf_session.correct(selected)
selected_images.append(selected)
if postprocess_settings is not None:
from .io import _normalize_img_batch
selected_images = [
image[..., np.newaxis] if image.ndim == 2 else image
for image in selected_images
]
max_height = max(image.shape[0] for image in selected_images)
max_width = max(image.shape[1] for image in selected_images)
selected_images = [
np.pad(
image,
(
(0, max_height - image.shape[0]),
(0, max_width - image.shape[1]),
(0, 0),
),
)
for image in selected_images
]
normalization_settings = dict(postprocess_settings)
from .object_roles import SEGMENTED_ROLES
for role in SEGMENTED_ROLES:
normalization_settings[f"{role}_channel"] = None
normalization_settings[f"{object_type}_channel"] = 0
if object_type == "cell" and selected_images[0].ndim == 3 \
and selected_images[0].shape[-1] > 1:
normalization_settings["nucleus_channel"] = 1
selected_stack = np.stack(selected_images).copy()
normalized_stack = _normalize_img_batch(
stack=selected_stack,
channels=range(selected_stack.shape[-1]),
save_dtype=np.float32,
settings=normalization_settings,
)
intensity_per_field = [
normalized_stack[index]
for index in range(normalized_stack.shape[0])
]
else:
intensity_per_field = selected_images
cellpose_images: List[np.ndarray] = []
for intensity_img in intensity_per_field:
img = np.asarray(intensity_img, dtype=np.float32)
maximum = float(img.max()) if img.size else 0.0
if maximum > 1:
img = img / maximum
cellpose_images.append(img)
with Timer(
f"v2.batch[{batch_start}:{batch_start + len(batch)}] "
f"cellpose ({len(loaded)} fields)",
logger="spacr.pipeline_v2",
):
if route is not None:
masks, _flows = backend_masks(
model, cellpose_images, backend_settings, object_type,
min_size=int(min_size),
default_diameter=backend_settings['_default_diameter'],
batch_size=max(8, len(cellpose_images)))
out = (masks,)
else:
out = model.eval(
cellpose_images,
batch_size=len(cellpose_images),
normalize=False,
channel_axis=-1,
min_size=int(min_size),
progress=True,
diameter=diameter,
flow_threshold=float(flow_threshold),
cellprob_threshold=float(cellprob_threshold),
resample=bool(resample),
)
masks = out[0]
if isinstance(masks, np.ndarray) and masks.ndim == 2:
masks = [masks]
masks_per_field = [
np.asarray(mask).astype(np.uint16) for mask in masks
]
if postprocess_settings is not None:
from .object import merge_split_filter_masks
filter_images = intensity_per_field
if filter_by_raw_intensity:
filter_images = [
np.pad(raw, (
(0, image.shape[0] - raw.shape[0]),
(0, image.shape[1] - raw.shape[1]),
))
for raw, image in zip(raw_intensity_per_field,
cellpose_images)
]
masks_per_field = list(merge_split_filter_masks(
masks=masks_per_field,
intensity_images=filter_images,
settings=postprocess_settings,
object_type=object_type,
batch_filenames=[stack.path.name for stack in batch],
))
for sf, arr, mask in zip(batch, loaded, masks_per_field):
mask = np.asarray(mask)[:arr.shape[0], :arr.shape[1]]
combined = np.concatenate(
[arr, mask[..., None]], axis=-1
).astype(np.uint16)
from .io import _save_array_atomic
_save_array_atomic(str(sf.path), combined)
sf.shape = combined.shape
sf.channels = sf.channels + [mask_channel_name]
if illumination_session is not None:
illumination_session.mark_completed(sf.field_id)
if psf_session is not None:
psf_session.mark_completed(sf.field_id)
if not keep_npz:
try:
npz_path.unlink()
except Exception:
pass
if not keep_npz:
try:
shutil.rmtree(scratch, ignore_errors=True)
except Exception:
pass
sidecar = stacks[0].path.parent / "channel_order.json"
try:
meta = json.loads(sidecar.read_text())
meta["mask_channels"] = [mask_channel_name]
sidecar.write_text(json.dumps(meta, indent=2))
except Exception as exc:
LOG.warning("channel_order.json at %s was not updated with "
"mask_channels=%r (%s); readers of this stack will "
"not know which plane holds the mask.",
sidecar, mask_channel_name, exc)
if illumination_session is not None:
illumination_session.finish(sf.field_id for sf in stacks)
if psf_session is not None:
psf_session.finish(sf.field_id for sf in stacks)
return stacks
[docs]
def run_v2(
src: Path,
channels: Sequence[int] = (0, 1, 2, 3),
channel_names: Optional[Sequence[str]] = None,
model_name: str = "cyto",
channels_for_cellpose: Sequence[int] = (0, 0),
diameter: Optional[float] = None,
batch_fields: int = 8,
metadata_type: str = "auto",
custom_regex: Optional[str] = None,
keep_npz: bool = False,
cellprob_threshold: float = 0.0,
flow_threshold: float = 0.4,
min_size: int = 15,
resample: bool = True,
postprocess_settings: Optional[Dict[str, Any]] = None,
object_type: str = "cell",
illumination_settings: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
"""Run the entire v2 pipeline against ``src``. Convenience wrapper.
Equivalent to::
mapper = FilenameMapper.discover(src, metadata_type, custom_regex)
stacks = stream_originals_to_stack(src, mapper, channels, channel_names)
stream_masks_from_stack(stacks, model_name, channels_for_cellpose,
diameter, batch_fields, keep_npz=keep_npz)
:param src: plate folder holding the originals, scanned
non-recursively. ``filename_map.csv`` lands here and output goes
to ``<src>/merged``; neither path is overridable from this wrapper.
:param channels: channel numbers as parsed from the filename
(``C01`` gives 1), not C-axis positions. The default
``(0, 1, 2, 3)`` therefore misfits stock CellVoyager/Yokogawa
names: channel 0 never exists, so plane 0 is all zeros, and
``C04`` is dropped. Pass ``(1, 2, 3, 4)`` for those layouts.
:param channel_names: names recorded in ``channel_order.json`` and on
each :class:`StackFile`. A length mismatch with ``channels`` trips
a bare ``assert``, so it goes unchecked under ``python -O``.
:param model_name: resolved by
:func:`spacr.utils._resolve_cellpose_pretrained`, not passed as
Cellpose's ``model_type``. On Cellpose 4 every legacy name
(``"cyto"``, ``"nuclei"``, …) collapses to ``cpsam``, so only a
fine-tuned checkpoint path actually changes the weights, and a
path with no file behind it raises instead of falling back.
:param channels_for_cellpose: despite the name this never reaches
Cellpose's ``channels=`` argument; it selects C-axis indices from
the assembled stack, taken modulo the channel count (7 with C=4
becomes 3) and de-duplicated in order. The default ``(0, 0)``
collapses to one plane, and an empty sequence falls back to plane 0.
:param diameter: forwarded to ``model.eval``; ``None`` leaves Cellpose
to size objects itself. Unlike ``model_name`` it is still honoured
on Cellpose 4, which rescales the image by ``30 / diameter``.
:param batch_fields: fields loaded and segmented per batch, the
memory-versus-speed dial. ``0`` raises ``ValueError``; a negative
value segments nothing at all, yet ``channel_order.json`` is still
rewritten to claim a mask channel that no stack received.
:param metadata_type: ``"cellvoyager"``, ``"yokogawa"``, ``"custom"``
or ``"auto"``. Only ``"custom"`` is special-cased; every other
unrecognised string quietly behaves as ``"auto"`` instead of
raising.
:param custom_regex: required when ``metadata_type="custom"``, else
``ValueError``. It must supply the named groups (``plateID``,
``wellID``, ``fieldID``, ``chanID``, ``timeID``, ``sliceID``); any
group it omits silently defaults, so a regex without ``chanID``
makes every image channel 1 and therefore its own single-plane
field.
:param keep_npz: the per-batch NPZ is compressed into
``merged/_scratch`` either way — this only decides whether it and
the scratch folder survive, so ``False`` does not save the write.
:param cellprob_threshold: forwarded to ``model.eval`` through
``float()``.
:param flow_threshold: forwarded to ``model.eval`` through ``float()``.
:param min_size: forwarded to ``model.eval`` through ``int()``;
``None`` raises ``TypeError`` rather than meaning "no minimum".
:param resample: coerced with ``bool()``, so any non-empty string —
``"false"`` included — is True, while ``None`` is False.
:param postprocess_settings: ``None`` hands the raw selected planes to
Cellpose and skips post-processing entirely. Any dict, ``{}``
included, switches on both :func:`spacr.io._normalize_img_batch`
and :func:`spacr.object.merge_split_filter_masks`. The four
``*_channel`` role keys are rewritten on a copy (the caller's dict
is left alone): all cleared to None, then ``object_type`` set to 0,
plus ``nucleus_channel=1`` when ``object_type`` is ``"cell"`` and
two or more planes were selected.
:param object_type: picks the weights during ``model_name``
resolution, names the role given channel 0 in normalisation, and
is passed to the mask post-processor. An unrecognised value does
not raise — it adds a dead ``<value>_channel`` key and leaves every
real role unset.
:param illumination_settings: full Mask settings mapping. When it enables
``illumination_correction``, one model is prepared from the raw
persisted stacks after pass 1 and its session corrects only the
private Cellpose inputs in pass 2; omitted/off preserves the previous
byte-level output contract.
:returns: dict with ``mapper`` (:class:`FilenameMapper`), ``stacks``
(list of :class:`StackFile`), and ``dst`` (Path to ``merged/``).
:raises ValueError: from :meth:`FilenameMapper.discover` when ``src``
holds no images, when ``metadata_type="custom"`` has no
``custom_regex``, or from ``batch_fields=0``.
"""
src = Path(src)
mapper = FilenameMapper.discover(src, metadata_type=metadata_type,
custom_regex=custom_regex)
stacks = stream_originals_to_stack(
src, mapper, channels=channels, channel_names=channel_names,
)
from .image_quality import screen_fields
rejected_quality = set(screen_fields(src, postprocess_settings or {},
[stack.path for stack in stacks], channels))
stacks = [stack for stack in stacks if stack.path.name not in rejected_quality]
if not stacks:
return {'mapper': mapper, 'stacks': [], 'dst': src / 'merged'}
illumination_session = None
if (stacks and illumination_settings and
illumination_settings.get('illumination_correction', False)):
from .illumination import prepare_segmentation_illumination
persisted_positions = _cellpose_channel_indices(
channels_for_cellpose, len(stacks[0].channels))
illumination_session = prepare_segmentation_illumination(
illumination_settings,
src=stacks[0].path.parent,
channels=persisted_positions,
pipeline_style='v2',
)
from .psf_pipeline import _prepare_segmentation_psf
psf_session = _prepare_segmentation_psf(
postprocess_settings or {}, src,
_cellpose_channel_indices(channels_for_cellpose, len(stacks[0].channels)),
pipeline_style="v2", stack_dir=stacks[0].path.parent,
load=lambda path: _as_hwc(np.load(path)))
stream_masks_from_stack(
stacks, model_name=model_name,
channels_for_cellpose=channels_for_cellpose,
diameter=diameter, batch_fields=batch_fields,
keep_npz=keep_npz,
cellprob_threshold=cellprob_threshold,
flow_threshold=flow_threshold,
min_size=min_size,
resample=resample,
postprocess_settings=postprocess_settings,
object_type=object_type,
illumination_session=illumination_session,
psf_session=psf_session,
)
return {"mapper": mapper, "stacks": stacks,
"dst": src / "merged"}