diff --git a/scripts/create_resources/combine/process_datasets_vizgen_nebius.sh b/scripts/create_resources/combine/process_datasets_vizgen_nebius.sh index c6eb4ad78..48abf64fd 100644 --- a/scripts/create_resources/combine/process_datasets_vizgen_nebius.sh +++ b/scripts/create_resources/combine/process_datasets_vizgen_nebius.sh @@ -12,7 +12,14 @@ cd "$REPO_ROOT" set -e -input_dir="s3://openproblems-data/resources/datasets" +# Read the FRESH loader outputs + SC refs from the scratch raw folder, NOT the S3 +# resources/datasets copy. The S3 vizgen loader outputs are stale (Oct 2025, uint16): +# the old loader rasterized ~700k cells into a uint16 label, so cell IDs > 65535 wrapped +# and most cells were lost -> the segmentation covered only a partial band. The rebuilt +# scratch loader outputs are uint32 / full-field. (process_dataset itself is correct; it +# faithfully cropped the already-broken S3 label, which is why the band appeared.) +input_dir="/scratch/task_ist_preprocessing/raw" +#input_dir="s3://openproblems-data/resources/datasets" # stale (uint16 labels) -- do not use #publish_dir="s3://openproblems-data/resources/task_ist_preprocessing/datasets" publish_dir='/scratch/task_ist_preprocessing/datasets' diff --git a/src/data_processors/process_dataset/script.py b/src/data_processors/process_dataset/script.py index 438d54108..31b3ce1f8 100644 --- a/src/data_processors/process_dataset/script.py +++ b/src/data_processors/process_dataset/script.py @@ -36,7 +36,21 @@ ### VIASH END -def get_crop_coords(sdata, max_n_pixels=20000*20000): #50000*50000): +def _full_res_image(img): + """Full-resolution image DataArray, whether `img` is single-scale (a plain + xarray DataArray) or multiscale (a DataTree with scale0/scale1/... nodes). + + spatialdata stores an image either way depending on the loader, and indexing + a single-scale DataArray with ["scale0"] raises KeyError (it looks for a + coordinate named 'scale0'), so pick the top-res array explicitly. + """ + from xarray import DataArray + if isinstance(img, DataArray): + return img + return img["scale0"].image + + +def get_crop_coords(sdata, max_n_pixels=20000*20000): #50000*50000): """Get the crop coordinates to subset the sdata to max_n_pixels Arguments @@ -52,7 +66,7 @@ def get_crop_coords(sdata, max_n_pixels=20000*20000): #50000*50000): The crop coordinates """ - _, h, w = sdata['image']["scale0"].image.shape + _, h, w = _full_res_image(sdata['image']).shape #h, w = sdata # Check if the image is already below the maximum number of pixels diff --git a/src/methods_transcript_assignment/proseg/script.py b/src/methods_transcript_assignment/proseg/script.py index b11049ce4..c8fcd0690 100644 --- a/src/methods_transcript_assignment/proseg/script.py +++ b/src/methods_transcript_assignment/proseg/script.py @@ -23,7 +23,13 @@ 'samples': 200, 'burnin_samples': 200, 'voxel_size': 1.0, - 'burnin_voxel_size': 2.0 + 'burnin_voxel_size': 2.0, + 'cell_compactness': 0.04, + 'nuclear_reassignment_prob': 0.2, + 'diffusion_probability': 0.2, + 'ncomponents': 10, + 'diffusion_sigma_far': 4.0, + 'diffusion_sigma_near': 1.0 } meta = { 'name': 'proseg', @@ -113,14 +119,14 @@ command_suffix = ( f'''--nthreads {n_threads} ''' -# these should not need to be changed - f'''--ncomponents 10 ''' +# model knobs (exposed as config arguments; defaults == proseg's own CLI defaults) + f'''--ncomponents {par['ncomponents']} ''' # f'''--no-diffusion ''' # should be off by default (?) - f'''--diffusion-probability 0.2 ''' - f'''--diffusion-sigma-far 4 ''' - f'''--diffusion-sigma-near 1 ''' - f'''--nuclear-reassignment-prob 0.2 ''' - f'''--cell-compactness 0.03 ''' + f'''--diffusion-probability {par['diffusion_probability']} ''' + f'''--diffusion-sigma-far {par['diffusion_sigma_far']} ''' + f'''--diffusion-sigma-near {par['diffusion_sigma_near']} ''' + f'''--nuclear-reassignment-prob {par['nuclear_reassignment_prob']} ''' + f'''--cell-compactness {par['cell_compactness']} ''' # these can be changed as arguments f'''--voxel-layers {par['voxel_layers']} ''' f'''--samples {par['samples']} ''' diff --git a/src/methods_transcript_assignment/segger/NOTES.md b/src/methods_transcript_assignment/segger/NOTES.md index 4fd957c9b..0ce302627 100644 --- a/src/methods_transcript_assignment/segger/NOTES.md +++ b/src/methods_transcript_assignment/segger/NOTES.md @@ -54,14 +54,19 @@ Links: docs https://elihei2.github.io/segger_dev/ · repo https://github.com/dpe `cell_id` (prior; background → `"UNASSIGNED"`), plus dummy `qv=40.0` and `overlaps_nucleus`. The **prior `cell_id`** per transcript comes from looking up the label image at each transcript's integer pixel coords (truncated to int64, - then clipped to image bounds). `tx_pd` (clean RangeIndex) is the **canonical - transcript frame** over ALL transcripts. **Only the IN-BOUNDS transcripts are written - to the parquet** (see the empty-`bd`-batch gotcha): transcripts whose pixel coords fall - outside the label image are **excluded from segger's input** and stay unassigned (0) in - the output. So the parquet row order == the **in-bounds subset**, and segger's reported - `row_index` indexes that subset — `seg_orig_idx = np.nonzero(in_bounds)[0]` maps it back - to the full-frame position in step 5. On datasets whose segmentation covers the whole - transcript field this is a no-op (`n_oob == 0`, parquet == full frame). + then **clamped** to image bounds — a few can round a pixel past the raster edge, same + clamp as `basic_transcript_assignment` / `baysor` / `proseg`). `tx_pd` (clean RangeIndex) + is the **canonical transcript frame** over ALL transcripts, and **ALL of them are written + to the parquet**, so the parquet row order == `tx_pd` and segger's reported `row_index` + indexes it directly (step 5). Guards against a total coordinate-frame mismatch by raising + if *every* transcript clamps out of bounds (`n_oob == n_tx`). + **History:** this step formerly EXCLUDED out-of-bounds transcripts from the parquet and + remapped `row_index` through a `seg_orig_idx` array — a workaround for when the + process_dataset crop left ~38% of transcripts outside the cropped labels. That crop bug is + fixed (`crop_points_by_global_xy` crops transcripts to the same global box as the labels) + and a full run on re-processed Xenium + MERFISH confirmed `n_oob == 0`, so the exclusion + was a no-op and was removed in favour of the plain clamp. This is INDEPENDENT of the + empty-`bd`-batch crash (see that gotcha) — that one is intra-field, not OOB-driven. 4. **Run segger** — `segger segment -i -o --n-epochs --prediction-graph-buffer-ratio --prediction-mode --node-representation-dim`, launched via a `run_segger(node_dim)` @@ -265,30 +270,36 @@ Nextflow runner labels: `[hightime, midcpu, highmem, gpuh100]`. guards right after reading the segmentation (`int(seg_arr.data.max()) == 0 → raise` with an actionable message), mirroring pciseq's guard. Note the benchmark retries an exit-1 task 3× before ignoring it, so a guarded fail still burns 3 retries — re-processing is the cure. -- **Transcripts outside the label window crash segger's encoder with an empty `bd` batch.** +- **Empty-`bd`-batch crash — a cell-free tile, NOT out-of-bounds transcripts (STILL OPEN).** segger builds a heterogeneous graph with `'tx'` (transcript) and `'bd'` (boundary/cell) node types and applies a `Positional2dEmbedder` to **both**. That embedder does - `torch.zeros((batch.max()+1, 2))` — and `batch.max()` throws `RuntimeError: max(): Expected - reduction dim to be specified for input.numel() == 0` on an **empty** tensor. A training - mini-batch drawn entirely from **tiles that have transcripts but no boundaries** has zero - `'bd'` nodes → that empty `batch` → crash. Because the Lightning DataLoader shuffles, this - fires **mid-training** (e.g. epoch 8/19), not on batch 0 — so it looks intermittent. Source - of the transcript-only tiles: the combine step crops the image/labels to ~20000² but the - transcripts span a larger field, so a large fraction (observed **~38%**, 7.6M/19.7M on - mouse-brain rep3) sit outside the label. **Fix (`script.py`, in our adapter — no segger - patch):** compute an `in_bounds` mask and write **only in-bounds transcripts** to - `transcripts.parquet`; OOB transcripts (which can't belong to any segmented cell anyway) - stay unassigned (0) in the output, and segger's `row_index` is mapped back through - `seg_orig_idx = np.nonzero(in_bounds)[0]`. This also removes the earlier "40% clamped → - garbage prior" quality issue. **Fallback if empty-`bd` batches ever persist** (e.g. a - sparse in-coverage tile with transcripts but no cells): also patch segger's - `ist_encoder.py` `Positional2dEmbedder.forward` to no-op on `batch.numel()==0` (a - from_torch-style Docker `sed` shim), but prefer the adapter filter — it stops the degenerate - tile from forming rather than tolerating it. Separately worth chasing upstream: **why ~38% - of transcripts fall outside the label at all** — `sd.query.bounding_box` in the combine step - (`process_dataset/script.py:258`) should crop transcripts and labels together, so this may - be a combine/crop inconsistency affecting every assignment method (the basic ones just mark - those transcripts background and don't crash). + `torch.zeros((batch.max()+1, 2))` (`ist_encoder.py`) — and `batch.max()` throws + `RuntimeError: max(): Expected reduction dim to be specified for input.numel() == 0` on an + **empty** tensor. A mini-batch drawn from **tiles that have transcripts but no boundaries** + has zero `'bd'` nodes → empty `batch` → crash. + - **Original hypothesis (now disproven as the cause):** the transcript-only tiles came from + the combine step cropping labels to ~20000² while transcripts spanned a larger field + (~38%, 7.6M/19.7M on mouse-brain rep3, sat outside the label). The adapter worked around it + by EXCLUDING out-of-bounds transcripts from `transcripts.parquet`. + - **What actually happened:** the crop bug was fixed at the source (`crop_points_by_global_xy` + crops transcripts to the same global box as the labels). A full validation run (Seqera + `1HlVGwXktVakGm`) then showed segger **completes on Xenium** (mouse brain rep3, human breast) + but **still crashes on MERFISH** (`2022_vizgen_human_breast_cancer_merfish_combined/rep1`) — + with `0/19222237 transcripts fall outside the label image` (**n_oob == 0**), training all 80 + batches, then dying at **validation tile ~21/28**. So the crash is **intra-field**: a tile + covering tissue that has transcripts but no *segmented cells* (MERSCOPE fields have acellular + / off-tissue gaps + a big extracellular transcript fraction; `custom_segmentation` copies the + vendor mask, so transcript-only tiles exist). Xenium coverage is dense → never trips. + - Because n_oob is 0, the OOB exclusion was a confirmed no-op and has been **removed** + (replaced by the plain edge clamp; see step 3). It never addressed this crash anyway. + - **The real fix (still to do — needs a container rebuild):** patch segger's + `Positional2dEmbedder.forward` in `ist_encoder.py` to no-op on empty input, e.g. return + `pos.new_zeros((pos.shape[0], 2 * self.dim))` when `pos.numel() == 0`, applied as a + from_torch-style Docker `sed`/python shim in `config.vsh.yaml`. This is a **segger source + patch (image rebuild)**, unlike the OOB simplification which is script-only. Verify the shim + on a live `build_main` pod (`debug-component-k8s`) before baking it in — the GPU/RAPIDS + rebuild is the expensive feedback loop. Not upstream in segger as of pinned `0233cf62` + (PR #74 is a *different* fix — Shapely polygon validity, not this). - **No CI test coverage** — GPU-only; can only be exercised on a CUDA host. - **Small panels break segger's gene-embedding PCA.** segger's `setup_anndata` does `PCA(n_components=in_channels=128)` on a genes×genes correlation matrix built from genes @@ -352,8 +363,8 @@ So 0.05 is *not* a deviation; it is segger's current default. (`transcripts_key` ### Tiers - **Tier 0 — inputs, not parameters.** The biggest levers aren't in the sweep: the segmentation - prior fed in (which segmentation method, how good its boundaries) and the fraction of - transcripts that fall inside the label window (OOB transcripts are dropped — see the empty-`bd` + prior fed in (which segmentation method, how good its boundaries) and how densely those + boundaries cover the transcript field (sparse coverage → cell-free tiles → the empty-`bd` gotcha). Both are fixed by the upstream stage here. - **Tier 1 — highest impact on *what* gets assigned:** - **`prediction_mode`** {`nucleus`,`cell`,`uniform`} — which polygon set drives the prediction diff --git a/src/methods_transcript_assignment/segger/script.py b/src/methods_transcript_assignment/segger/script.py index 6bbb28ad6..71e1eb29d 100644 --- a/src/methods_transcript_assignment/segger/script.py +++ b/src/methods_transcript_assignment/segger/script.py @@ -160,35 +160,40 @@ label_image = sdata_segm["segmentation"]["scale0"].image.to_numpy() else: label_image = sdata_segm["segmentation"].to_numpy() -# Transcripts can land outside the label image after the transform (e.g. the combine step -# crops the image/labels to ~20000x20000 but transcripts span a larger field -> here ~38% -# of them fall outside). Those OOB transcripts must NOT be fed to segger: segger tiles the -# transcript field into a heterogeneous graph with 'tx' (transcript) and 'bd' (boundary/ -# cell) node types, and a tile covering a transcripts-only region (no boundaries) produces a -# training mini-batch with ZERO 'bd' nodes. segger's positional encoder then runs -# `torch.zeros((batch.max()+1, 2))` on an empty batch -> `RuntimeError: max(): Expected -# reduction dim to be specified for input.numel() == 0`, which kills training MID-RUN (only -# when the DataLoader's shuffle happens to draw an all-OOB mini-batch, so it can survive -# several epochs first). A transcript outside the segmentation can't belong to any segmented -# cell anyway, so we EXCLUDE it from the segger input and leave it unassigned (cell_id 0) in -# the final output. Datasets whose segmentation covers the whole transcript field are -# unaffected (in_bounds all-True, n_oob == 0). -in_bounds = ( - (y_coords >= 0) & (y_coords < label_image.shape[0]) - & (x_coords >= 0) & (x_coords < label_image.shape[1]) -) -n_oob = int((~in_bounds).sum()) -y_look = np.clip(y_coords, 0, label_image.shape[0] - 1) -x_look = np.clip(x_coords, 0, label_image.shape[1] - 1) -prior_cell_id = label_image[y_look, x_look].astype(np.int64) # 0 == background -prior_cell_id[~in_bounds] = 0 # OOB -> background (excluded from segger regardless) -print(f"{n_oob}/{len(in_bounds)} transcripts fall outside the " - f"{label_image.shape[0]}x{label_image.shape[1]} label image; excluding them from the " - f"segger input (they stay unassigned in the output).", flush=True) +# Clamp transcript pixel coords to the label-image bounds before the lookup: after the +# transform a few can round a pixel past the raster edge. Matches the clamp in +# basic_transcript_assignment / baysor / proseg. +# +# HISTORY: this step used to EXCLUDE out-of-bounds transcripts from segger's input (and remap +# segger's row_index back through a `seg_orig_idx` array), from when the process_dataset crop +# left a large fraction (~38% on mouse-brain rep3) of transcripts OUTSIDE the cropped labels. +# That crop bug is fixed -- crop_points_by_global_xy now crops transcripts to the SAME global +# box as the labels -- and a full run on re-processed Xenium + MERFISH confirmed n_oob == 0, so +# the exclusion + remap were a confirmed no-op. Removed here in favour of the plain clamp; edge +# rounding is all that remains, and it's handled below. NOTE: this does NOT touch the separate +# empty-`bd`-batch crash (a cell-free INTERIOR tile), which is intra-field, not OOB-driven, and +# still needs the segger-encoder guard -- see NOTES.md. +n_oob = int(np.count_nonzero( + (y_coords < 0) | (y_coords >= label_image.shape[0]) + | (x_coords < 0) | (x_coords >= label_image.shape[1]) +)) +y_coords = np.clip(y_coords, 0, label_image.shape[0] - 1) +x_coords = np.clip(x_coords, 0, label_image.shape[1] - 1) +prior_cell_id = label_image[y_coords, x_coords].astype(np.int64) # 0 == background +print(f"Clamped {n_oob}/{len(prior_cell_id)} transcripts outside the " + f"{label_image.shape[0]}x{label_image.shape[1]} label image to its edge", flush=True) +# Every transcript clamping out of bounds means transcripts and segmentation are almost +# certainly in mismatched coordinate frames -- fail fast rather than assign everything to edge +# pixels (mirrors the intent of the old zero-in-bounds guard). +if n_oob == len(prior_cell_id): + raise ValueError( + "Every transcript falls outside the segmentation label image, so the clamp would map " + "them all to edge pixels. This usually means the transcripts and segmentation are in " + "mismatched coordinate frames (check the dataset's crop / transforms)." + ) # Canonical transcripts frame (native coordinates, clean RangeIndex). Its row order matches -# prior_cell_id / in_bounds; the IN-BOUNDS subset's order matches the parquet we write and -# hence segger's reported row_index (mapped back through seg_orig_idx below). +# prior_cell_id AND the parquet we write below, and hence segger's reported row_index. tx_pd = transcripts_reset.compute() n_tx = len(tx_pd) @@ -197,30 +202,19 @@ else: transcript_id = np.arange(n_tx, dtype=np.uint64) -# Full-frame positions of the transcripts actually handed to segger. segger's output -# row_index indexes into the (in-bounds) parquet row order; seg_orig_idx maps it back. -seg_orig_idx = np.nonzero(in_bounds)[0] -if seg_orig_idx.size == 0: - raise ValueError( - "No transcripts fall within the segmentation label image, so segger has nothing to " - "assign. This usually means the transcripts and segmentation are in mismatched " - "coordinate frames (check the dataset's crop / transforms)." - ) - -n_in = int(in_bounds.sum()) tx_out = pd.DataFrame({ - "transcript_id": transcript_id[in_bounds], - "x_location": tx_pd["x"].to_numpy().astype(np.float32)[in_bounds], - "y_location": tx_pd["y"].to_numpy().astype(np.float32)[in_bounds], - "feature_name": tx_pd["feature_name"].astype(str).to_numpy()[in_bounds], + "transcript_id": transcript_id, + "x_location": tx_pd["x"].to_numpy().astype(np.float32), + "y_location": tx_pd["y"].to_numpy().astype(np.float32), + "feature_name": tx_pd["feature_name"].astype(str).to_numpy(), # segger's Xenium loader keys transcripts to boundaries by cell_id string; # background (0) becomes the UNASSIGNED sentinel. - "cell_id": np.where(prior_cell_id[in_bounds] > 0, prior_cell_id[in_bounds].astype(str), "UNASSIGNED"), - "qv": np.full(n_in, 40.0, dtype=np.float32), - "overlaps_nucleus": (prior_cell_id[in_bounds] > 0).astype(np.int8), + "cell_id": np.where(prior_cell_id > 0, prior_cell_id.astype(str), "UNASSIGNED"), + "qv": np.full(n_tx, 40.0, dtype=np.float32), + "overlaps_nucleus": (prior_cell_id > 0).astype(np.int8), }) if "z" in tx_pd.columns: - tx_out.insert(3, "z_location", tx_pd["z"].to_numpy().astype(np.float32)[in_bounds]) + tx_out.insert(3, "z_location", tx_pd["z"].to_numpy().astype(np.float32)) tx_out.to_parquet(XENIUM_DIR / "transcripts.parquet", index=False) del transcripts, y_coords, x_coords, label_image @@ -341,15 +335,15 @@ def run_segger(node_dim): seg = seg[keep_mask] print(f"segger kept {len(seg)} assignments of {n_tx} transcripts", flush=True) -# segger reports row_index into the transcripts.parquet we wrote, which is the IN-BOUNDS -# subset (OOB transcripts were excluded above). Map that back to the full-frame position via -# seg_orig_idx, then factorize the (string) segger_cell_id into contiguous positive integers; -# unassigned / out-of-bounds transcripts stay 0. +# segger reports row_index into the transcripts.parquet we wrote. Since we now write ALL +# transcripts (the full canonical frame, no OOB exclusion), row_index indexes tx_pd directly. +# Factorize the (string) segger_cell_id into contiguous positive integers; unassigned +# transcripts stay 0. cell_id_per_tx = np.zeros(n_tx, dtype=np.int64) if len(seg): row_idx = seg["row_index"].to_numpy().astype(np.int64) codes, _ = pd.factorize(seg["segger_cell_id"].astype(str), sort=True) - cell_id_per_tx[seg_orig_idx[row_idx]] = codes.astype(np.int64) + 1 + cell_id_per_tx[row_idx] = codes.astype(np.int64) + 1 ############################################# # Build transcript-assignment output object #