Fixes

src/data_processors/process_dataset/script.py

@@ -172,6 +172,61 @@ def crop_shapes_by_global_xy(shapes, x0, x1, y0, y1):
     set_transformation(new, trans, set_all=True)
     return new
 
+def rasterize_boundaries_to_labels(sdata, shapes_key="cell_boundaries", image_key="image"):
+    """Rasterize polygon cell boundaries into a ``cell_labels`` labels element.
+
+    Some loaders (e.g. allen_brain_cell_atlas_merfish) provide the vendor
+    segmentation only as polygon shapes and never rasterize it into a label
+    image. The ``custom_segmentation`` method hard-requires ``labels["cell_labels"]``,
+    so synthesize it here from the boundaries when it is absent — done post-crop
+    (see call site) so only the retained region is rasterized.
+
+    Mirrors the vizgen_merscope loader: ``sd.rasterize`` labels regions 1..N
+    positionally and caps a single pass at 65535 regions, so rasterize in chunks
+    and offset each chunk's labels past that; then promote to a multiscale pyramid
+    so downstream ``["scale0"]`` indexing on the copied segmentation works.
+    """
+    import dask.array as da
+    from spatialdata.models import Labels2DModel
+
+    # sd.rasterize(return_regions_as_labels=True) numbers regions 1..n positionally
+    # into a uint16 array, so a single pass encodes at most 65535 distinct cells.
+    UINT16_MAX = 65535
+    img_extent = sd.get_extent(sdata[image_key])
+    n_cells = len(sdata[shapes_key])
+    n_iter = n_cells // UINT16_MAX + bool(n_cells % UINT16_MAX)
+
+    rasterize_args = {
+        "min_coordinate": [int(img_extent["x"][0]), int(img_extent["y"][0])],
+        "max_coordinate": [int(img_extent["x"][1]), int(img_extent["y"][1])],
+        "target_coordinate_system": "global",
+        "target_unit_to_pixels": 1,
+        "return_regions_as_labels": True,
+    }
+
+    if n_iter <= 1:
+        labels_image = sd.rasterize(sdata[shapes_key], ["x", "y"], **rasterize_args)
+    else:
+        combined = None
+        template = None
+        for i in range(n_iter):
+            start = i * UINT16_MAX
+            end = min((i + 1) * UINT16_MAX, n_cells)
+            chunk = sd.rasterize(sdata[shapes_key].iloc[start:end], ["x", "y"], **rasterize_args)
+            chunk_np = np.asarray(chunk.data)
+            if combined is None:
+                combined = chunk_np.astype("uint32")
+                template = chunk
+            else:
+                mask = chunk_np > 0
+                combined[mask] = chunk_np[mask].astype("uint32") + start
+        labels_image = template.copy(data=da.from_array(combined, chunks=template.data.chunksize))
+
+    # rasterize tags the labels with a shape->category map that the Labels model
+    # does not expect; drop it before parsing (matches the vizgen loader).
+    labels_image.attrs.pop("label_index_to_category", None)
+    return Labels2DModel.parse(labels_image, scale_factors=[2, 2, 2, 2])
+
 def rechunk_sdata(sdata, CHUNK_SIZE=1024):
     """Rechunk the sdata to the given chunk size
 
@@ -383,6 +438,15 @@ def subsample_adata_group_balanced(adata, group_key, n_samples, seed=0):
 else:
     sdata_output = sdata
 
+# Synthesize cell_labels from polygon boundaries for loaders that only provide
+# shapes (e.g. allen_brain_cell_atlas_merfish). Done here — post-crop — so we
+# rasterize only the retained region (rasterizing whole-brain labels would OOM)
+# and avoid re-running the expensive stitching loader. custom_segmentation
+# requires labels["cell_labels"]; loaders that already provide it are untouched.
+if "cell_labels" not in sdata_output.labels and "cell_boundaries" in sdata_output.shapes:
+    print("No cell_labels found; rasterizing cell_boundaries -> cell_labels", flush=True)
+    sdata_output["cell_labels"] = rasterize_boundaries_to_labels(sdata_output)
+
 # Rechunk to uniform chunks before writing (NOTE: rechunking currently needed,
 # https://github.com/scverse/spatialdata/issues/929). Run unconditionally so
 # that uncropped datasets (e.g. 10x Atera, whose store has rectilinear chunk

src/methods_segmentation/stardist/NOTES.md

@@ -48,31 +48,41 @@ detector**: the script feeds it `image[0]` only (see Tier 0 below).
    (prob=0.479071, nms=0.3 for that model) as the fallback thresholds.
 4. **Percentile normalizer** (`:64-77`) — a csbdeep `Normalizer` subclass that min-max
    scales the image to its **1st / 99.8th percentiles** (`normalize_mi_ma`), the
-   recommended StarDist preprocessing but with fixed percentile bounds. `block_size`
-   and `context` are derived from the image width so a large panel is processed in
-   tiles; **`min_overlap` is derived from object size, not block size** (see step 6 and
-   the min_overlap gotcha below).
+   recommended StarDist preprocessing but with fixed percentile bounds. The **image
+   size then selects the segmentation path** (single-pass vs tiled — see step 6 and the
+   min_overlap gotcha below).
 5. **Build eval-params** (`:85-89`) — collects the newly exposed tunables
    (`prob_thresh, nms_thresh, scale`) from `par`, **dropping any that are `None`**.
    A dropped key ⇒ `predict_instances` uses the model's own optimized value (so the
    no-args call is byte-for-byte the pre-tuning behaviour). Mirrors the cellposev4
    eval-params pattern.
-6. **Segment** (`:92-131`) — `model.predict_instances_big(image[0], axes='YX',
-   block_size=…, min_overlap=…, context=…, normalizer=…, **eval_params)`.
-   `predict_instances_big` splits the image into `block_size` blocks, calls
-   `predict_instances` on each (forwarding `**eval_params` unchanged — it only
-   overrides `axes/overlap_label/return_labels/return_predict`), and reassembles the
-   labels into global coordinates. `image[0]` = first channel → a single 2D plane.
-   **The stitching invariant is that every predicted object is smaller than
-   `min_overlap`** (an object bigger than the overlap can span a block seam and can't be
-   uniquely assigned → `RuntimeError: ...violates the assumption of being smaller than
-   'min_overlap'`). So `min_overlap` is set from an **object-size** bound
-   (`max_object_diameter`, default **192 px**), *not* from `block_size` (the old
-   `block_size // 5.5` shrank it to 64 px on small panels while real blobs reached
-   ~110 px → crash). `block_size` is then grown if needed to satisfy
-   `min_overlap + 2*context < block_size`, and the call is wrapped in a **retry that
-   doubles `min_overlap` on that specific error** so a rare oversized blob self-heals
-   instead of failing the run.
+6. **Segment** (`:91-141`) — **two paths chosen by image size** (`image[0]` = first
+   channel → a single 2D plane):
+   - **Fits (largest side ≤ `BIG_PX`=4096) → `model.predict_instances(...)`.** One pass,
+     **no block stitching**, so there is *no* `min_overlap`/block-geometry constraint and
+     objects of any size are fine. `n_tiles` only sub-tiles the **forward pass** to bound
+     GPU memory (that tiling has its own automatic context and no overlap requirement).
+     This is the path all current benchmark panels take.
+   - **Large whole-slide (largest side > 4096) → `model.predict_instances_big(...)`.**
+     Tiles into `block_size` blocks and stitches, under **two** constraints: (1) a
+     *stitching* invariant that every object be smaller than `min_overlap` (else
+     `RuntimeError: ...violates the assumption of being smaller than 'min_overlap'`), and
+     (2) a *block-geometry* one — per-axis stride is `size − (min_overlap + 2·context)`
+     and stardist's `Block.cover` **asserts** consecutive write-regions overlap by ≥
+     `min_overlap`, which requires `block_size` **comfortably** larger than
+     `min_overlap + 2·context`. Here `block_size = BIG_PX = 4096 ≫ min_overlap(192) +
+     2·context`, so the geometry holds. `min_overlap` is **object-size-based**
+     (`max_object_diameter`, default 192 px, measured in original px — `predict_instances`
+     undoes `scale`), and the call is wrapped in a **retry that doubles `min_overlap`**
+     on that `RuntimeError` so a rare oversized blob self-heals.
+
+   Why the split: the old code *always* used `predict_instances_big` with
+   `block_size = image.shape[1] // 3` (~336 px on a ~1000 px panel), which (a) forced
+   tiling even on tiny images and (b) left only a ~16 px margin over `min_overlap +
+   2·context` → `Block.cover` `AssertionError`. A too-small `min_overlap` (`block_size //
+   5.5` = 64 px) had earlier caused the *stitching* `RuntimeError` on a 110 px blob. Both
+   bug classes only exist on the tiled path; single-pass `predict_instances` sidesteps
+   them entirely.
 7. **Post-process** (`:100-104`) — `convert_to_lower_dtype` downcasts the label array
    to the smallest uint that holds `max label`; wrap as an `xarray.DataArray`,
    `Labels2DModel.parse` with the copied transform, store as
@@ -129,10 +139,11 @@ wrong for a different `--model` whose optimized thresholds differ. **They need
 `viash ns build` + a container rebuild to take effect** (see `check-component`).
 
 `max_object_diameter` is a **geometry/robustness knob, not a quality knob** — it only
-sizes `predict_instances_big`'s `min_overlap`; it does not change which pixels get
-segmented. It is **not part of the quality sweep**; leave it at the 192 px default
-unless you hit the min_overlap `RuntimeError` (the script also auto-doubles it), or your
-nuclei are unusually large.
+sizes `min_overlap` on the **tiled (>4096 px) path**; on the single-pass path it is
+unused, and it never changes which pixels get segmented. It is **not part of the quality
+sweep**; leave it at the 192 px default unless a large whole-slide image hits the
+min_overlap `RuntimeError` (the script also auto-doubles it) or its nuclei are unusually
+large.
 
 Not exposed:
 - `--n_tiles` — a pure GPU-memory tiling knob. `predict_instances_big` **already** tiles
@@ -219,16 +230,23 @@ That is exactly the sweep encoded in `scripts/run_benchmark/stardist_params.yaml
   `viash ns build` + a container rebuild; a stale image silently ignores them. The sweep
   has **not yet been run end-to-end** with the new args — validated only by
   `viash config view` + a script `ast.parse`.
-- **`min_overlap` must exceed the largest object.** `predict_instances_big`'s block
-  stitching asserts every object is smaller than `min_overlap`; a bigger object throws
-  `RuntimeError: ...violates the assumption of being smaller than 'min_overlap'`. The old
-  code tied it to `block_size // 5.5`, so on small/narrow panels it fell to 64 px while
-  real blobs reached ~110 px → crash. Now `min_overlap` is **object-size-based**
-  (`max_object_diameter`, default 192 px), `block_size` is grown to keep
-  `min_overlap + 2*context < block_size`, and a **retry doubles `min_overlap`** on that
-  error. `scale` is forwarded per-block via `**kwargs` but objects are measured in
-  original pixels (predict_instances undoes `scale`), so `min_overlap` is in original px
-  regardless of `scale`; still sanity-check masks at block seams for extreme `scale`.
+- **The `predict_instances_big` tiling has TWO independent failure modes — which is why
+  small images now bypass it entirely** (single-pass `predict_instances`, see step 6):
+  1. *Stitching* — `RuntimeError: ...violates the assumption of being smaller than
+     'min_overlap'` when an object is bigger than `min_overlap`. The old
+     `min_overlap = block_size // 5.5` fell to 64 px on small panels while blobs reached
+     ~110 px.
+  2. *Block geometry* — `AssertionError` in `Block.cover` (per-axis
+     `stride = size − (min_overlap + 2·context)`; consecutive write-regions must overlap
+     by ≥ `min_overlap`). Fires when `block_size` is only *marginally* above
+     `min_overlap + 2·context`. Over-correcting fix #1 to `min_overlap=192` with
+     `block_size = image//3 ≈ 336` left a 16 px margin → this assertion tripped.
+
+  On the surviving tiled path both are avoided by construction: `block_size = 4096 ≫
+  min_overlap(192) + 2·context`, and the retry doubles `min_overlap` for a rare huge
+  object. `scale` is forwarded per-block but objects are measured in original pixels
+  (`predict_instances` undoes `scale`), so `min_overlap` is scale-independent; still
+  sanity-check masks at block seams for extreme `scale`.
 - **Only channel 0 is segmented** (`image[0]`). Fine for single-channel iST morphology;
   StarDist2D has no multi-channel mode anyway.
 - **Whole image loaded into RAM** (`:53`) — full-res plane; big panels are why the label

src/methods_segmentation/stardist/script.py

@@ -77,7 +77,7 @@ def do_after(self):
 mi, ma = np.percentile(image, [1,99.8])
 normalizer = MyNormalizer(mi, ma)
 
-# Tunable knobs forwarded through predict_instances_big -> predict_instances.
+# Tunable knobs forwarded to predict_instances / predict_instances_big.
 # A value left as None (i.e. omitted from par) means "use the model's own optimized
 # value": thresholds.json for prob_thresh/nms_thresh, no rescaling for scale. This
 # keeps the default (no-args) call identical to the pre-tuning behaviour.
@@ -86,46 +86,67 @@ def do_after(self):
     for k in ("prob_thresh", "nms_thresh", "scale")
     if par.get(k) is not None
 }
-print(f"predict_instances_big overrides: {eval_params}", flush=True)
-
-# predict_instances_big tiles the image and stitches the per-block predictions. Its
-# stitching invariant is that EVERY predicted object is smaller than `min_overlap`;
-# an object spanning a block seam that is larger than the overlap can't be assigned to
-# a single block, which raises "Found object of shape (...), which violates the
-# assumption of being smaller than 'min_overlap'". So `min_overlap` must be
-# OBJECT-SIZE-based, not block-size-based — the old `block_size // 5.5` shrank the
-# overlap below real nuclei/blobs on small panels (min_overlap fell to 64 px while
-# objects reached ~110 px). Objects are measured in ORIGINAL image pixels
-# (predict_instances undoes `scale` internally), so this bound is in original px and is
-# independent of `scale`. `context` is only the receptive-field margin discarded around
-# each block, so deriving it from the image size is fine.
-block_size = min(image.shape[1] // 3, 4096)
-context = int(min(block_size // 5.5, 128))
-min_overlap = int(par.get("max_object_diameter") or 192)  # px; must exceed largest object
-# predict_instances_big asserts: min_overlap + 2*context < block_size.
-block_size = max(block_size, min_overlap + 2 * context + 1)
-
-# Self-heal: if a rare oversized blob (merged nuclei / debris) still exceeds
-# `min_overlap`, double it (and grow block_size to keep the geometry constraint) and
-# retry, rather than failing the whole segmentation.
-while True:
-    try:
-        labels, _ = model.predict_instances_big(
-            image[0, :, :], axes='YX', block_size=block_size,
-            min_overlap=min_overlap, context=context,
-            normalizer=normalizer, **eval_params,  # n_tiles left to block_size
-        )
-        break
-    except RuntimeError as e:
-        if "min_overlap" not in str(e) or min_overlap >= 2048:
-            raise
-        min_overlap *= 2
-        block_size = max(block_size, min_overlap + 2 * context + 1)
-        print(
-            "predict_instances_big: an object exceeded min_overlap; retrying with "
-            f"min_overlap={min_overlap}, block_size={block_size}",
-            flush=True,
-        )
+print(f"stardist overrides: {eval_params}", flush=True)
+
+# Segmentation strategy — two paths, chosen by image size:
+#
+#   * predict_instances_big TILES the image and stitches the per-block predictions
+#     under TWO strict constraints. (1) a *stitching* invariant that every object be
+#     smaller than `min_overlap`; and (2) a *block-geometry* one, since the per-axis
+#     stride is `size - (min_overlap + 2*context)` and stardist's `Block.cover`
+#     asserts consecutive blocks' write-regions overlap by >= min_overlap — which
+#     needs `block_size` to be *comfortably* larger than `min_overlap + 2*context`,
+#     not just larger (the old `block_size = image.shape[1] // 3` made ~336 px blocks
+#     on a ~1000 px panel, leaving a 16 px margin, so Block.cover's assertion failed).
+#   * predict_instances processes the whole image in ONE pass: no block stitching, so
+#     NEITHER constraint exists and objects of any size are fine. `n_tiles` only
+#     sub-tiles the forward pass to bound GPU memory (that tiling has its own automatic
+#     context and no min_overlap requirement).
+#
+# So: if the image fits (largest side <= BIG_PX) use the constraint-free single-pass
+# `predict_instances`; only genuinely large whole-slide images take the tiled path,
+# where block_size=BIG_PX >> min_overlap+2*context keeps Block.cover's geometry valid.
+# `min_overlap` (the tiled path only) is OBJECT-SIZE-based: it must exceed the largest
+# object, measured in ORIGINAL image pixels (predict_instances undoes `scale`), so it is
+# independent of `scale`.
+def _n_tiles(px):
+    # Forward-pass tiling to bound GPU memory (~2048 px/tile); no stitching constraint.
+    n = max(1, int(px) // 2048)
+    return (n, n)
+
+BIG_PX = 4096
+max_dim = int(max(image.shape[1], image.shape[2]))
+
+if max_dim <= BIG_PX:
+    labels, _ = model.predict_instances(
+        image[0, :, :], axes='YX', normalizer=normalizer,
+        n_tiles=_n_tiles(max_dim), **eval_params,
+    )
+else:
+    # Large image -> tile + stitch. Self-heal: if a rare oversized blob (merged nuclei /
+    # debris) exceeds `min_overlap`, double it (and grow block_size to keep the geometry
+    # constraint) and retry, rather than failing the whole segmentation.
+    block_size = BIG_PX
+    min_overlap = int(par.get("max_object_diameter") or 192)  # px; must exceed largest object
+    context = int(min(min_overlap, 128))
+    while True:
+        try:
+            labels, _ = model.predict_instances_big(
+                image[0, :, :], axes='YX', block_size=block_size,
+                min_overlap=min_overlap, context=context, n_tiles=_n_tiles(block_size),
+                normalizer=normalizer, **eval_params,
+            )
+            break
+        except RuntimeError as e:
+            if "min_overlap" not in str(e) or min_overlap >= 2048:
+                raise
+            min_overlap *= 2
+            block_size = max(block_size, min_overlap + 2 * context + 256)
+            print(
+                "predict_instances_big: an object exceeded min_overlap; retrying with "
+                f"min_overlap={min_overlap}, block_size={block_size}",
+                flush=True,
+            )