diff --git a/scripts/run_benchmark/run_test_segger_nebius.sh b/scripts/run_benchmark/run_test_segger_nebius.sh new file mode 100644 index 000000000..04bc72e90 --- /dev/null +++ b/scripts/run_benchmark/run_test_segger_nebius.sh @@ -0,0 +1,68 @@ +#!/bin/bash + +# Test run for the segger transcript-assignment method. +# Runs on the small S3 test resources, using the DEFAULT method at every stage +# plus segger added to the transcript-assignment stage. Segger requires a GPU, +# so this runs on Nebius (the gpu labels in labels_nebius.config pin it to the +# GPU node group). + +# get the root of the directory +REPO_ROOT=$(git rev-parse --show-toplevel) + +# ensure that the command below is run from the root of the repository +cd "$REPO_ROOT" + +set -e + +resources_test_s3=s3://openproblems-data/resources_test/task_ist_preprocessing +publish_dir_s3="/mnt/data/results/runs/$(date +%Y-%m-%d_%H-%M-%S)_segger_test" + +cat > /tmp/params_settings_segger.yaml << HERE +default_methods: + - custom_segmentation + - basic_transcript_assignment + - basic_count_aggregation + - basic_qc_filter + - alpha_shapes + - normalize_by_volume + - tacco + - no_correction +segmentation_methods: + - custom_segmentation +transcript_assignment_methods: + - basic_transcript_assignment + - segger +count_aggregation_methods: + - basic_count_aggregation +qc_filtering_methods: + - basic_qc_filter +volume_calculation_methods: + - alpha_shapes +normalization_methods: + - normalize_by_volume +celltype_annotation_methods: + - tacco +expression_correction_methods: + - no_correction +HERE + +# Write the parameters to file (input_states version, NOTE: enable `-entry_name auto` for this) +cat > /tmp/params_segger.yaml << HERE +input_states: $resources_test_s3/**/state.yaml +rename_keys: 'input_sc:output_sc;input_sp:output_sp' +save_spatial_data: false +settings: '$(yq -o json /tmp/params_settings_segger.yaml | jq -c .)' +output_state: "state.yaml" +publish_dir: "$publish_dir_s3" +HERE + +tw launch https://github.com/openproblems-bio/task_ist_preprocessing.git \ + --revision build/main \ + --pull-latest \ + --main-script target/nextflow/workflows/run_benchmark/main.nf \ + --workspace 167877437119966 \ + --compute-env 5hfmdCBxMRd4nHZaJKYEQZ \ + --params-file /tmp/params_segger.yaml \ + --entry-name auto \ + --config src/base/labels_nebius.config \ + --labels task_ist_preprocessing,test,segger diff --git a/src/methods_transcript_assignment/segger/config.vsh.yaml b/src/methods_transcript_assignment/segger/config.vsh.yaml new file mode 100644 index 000000000..bb8288d8c --- /dev/null +++ b/src/methods_transcript_assignment/segger/config.vsh.yaml @@ -0,0 +1,82 @@ +__merge__: /src/api/comp_method_transcript_assignment.yaml + +name: segger +label: "Segger Transcript Assignment" +summary: "Assign transcripts to cells using the Segger method" +description: "Segger is a tool for cell segmentation in single-molecule spatial omics datasets that leverages graph neural networks (GNNs) and heterogeneous graphs. This component runs segger's end-to-end CLI (`segger segment`), using the provided segmentation as the boundary prior, and maps the resulting per-transcript assignments back onto the transcripts." +links: + documentation: "https://elihei2.github.io/segger_dev/" + repository: "https://github.com/dpeerlab/segger" +references: + doi: "10.1101/2025.03.14.643160" + +arguments: + - name: --transcripts_key + type: string + description: The key of the transcripts within the points of the spatial data + default: transcripts + - name: --coordinate_system + type: string + description: The key of the pixel space coordinate system within the spatial data + default: global + - name: --n_epochs + type: integer + required: false + description: "Number of training epochs for the segger GNN." + default: 20 + - name: --prediction_expansion_ratio + type: double + required: false + description: "Fraction of each polygon's equivalent radius used to expand it during prediction." + default: 0.5 + - name: --prediction_mode + type: string + required: false + choices: [nucleus, cell, uniform] + description: "Which polygon set drives the prediction graph." + default: nucleus + +resources: + - type: python_script + path: script.py + +engines: + # NOTE: segger runs its tiling / training / prediction end-to-end on the GPU + # (cudf, cuspatial, torch). Develop/test on a CUDA host; `viash test` + # will not run on a machine without an NVIDIA GPU. + - type: docker + image: nvcr.io/nvidia/pytorch:25.06-py3 + setup: + - type: apt + packages: [procps, git, libxcb1] + - type: python + github: + - openproblems-bio/core#subdirectory=packages/python/openproblems + - type: python + packages: + - "spatialdata>=0.7.3" + - "anndata>=0.12.0" + - "zarr>=3.0.0" + - geopandas + - shapely + - "rasterio>=1.3" + - scikit-image + # cuspatial is required by segger.geometry.conversion at import time and is + # NOT shipped in the NGC PyTorch base image (only cudf is). No version pin, + # so pip resolves it against whatever cudf the image already has installed + # (pinning would risk diverging from the image's cudf). The NVIDIA index is + # required: cuspatial's libcudf-cu12/libcuspatial-cu12 deps are not on PyPI. + - type: docker + run: | + pip install --no-cache-dir --extra-index-url=https://pypi.nvidia.com cuspatial-cu12 + # segger trunk installed last so its (unpinned) deps resolve against the + # RAPIDS/torch stack already present in the image. + - type: python + github: dpeerlab/segger + - type: native + +runners: + - type: executable + - type: nextflow + directives: + label: [ hightime, midcpu, highmem, gpu ] diff --git a/src/methods_transcript_assignment/segger/script.py b/src/methods_transcript_assignment/segger/script.py new file mode 100644 index 000000000..ff43e2662 --- /dev/null +++ b/src/methods_transcript_assignment/segger/script.py @@ -0,0 +1,246 @@ +import os +import shutil +import subprocess +from pathlib import Path + +import dask +import numpy as np +import pandas as pd +import anndata as ad +import xarray as xr +import torch +import spatialdata as sd +from spatialdata.models import PointsModel + +## VIASH START +# Note: this section is auto-generated by viash at runtime. To edit it, make changes +# in config.vsh.yaml and then run `viash config inject config.vsh.yaml`. +par = { + 'input_ist': 'resources_test/task_ist_preprocessing/mouse_brain_combined/raw_ist.zarr', + 'input_segmentation': 'resources_test/task_ist_preprocessing/mouse_brain_combined/segmentation.zarr', + 'transcripts_key': 'transcripts', + 'coordinate_system': 'global', + 'output': './temp/methods/segger/segger_assigned_transcripts.zarr', + 'n_epochs': 20, + 'prediction_expansion_ratio': 0.5, + 'prediction_mode': 'nucleus', +} +meta = { + 'name': 'segger', + 'temp_dir': './temp/methods/segger', + 'cpus': 4, +} +## VIASH END + +# segger runs its tiling / GNN training / prediction end-to-end on the GPU +# (cudf, cuspatial, torch kernels). Fail fast with a clear message rather than +# deep inside the subprocess when no CUDA device is present. +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" +print("Using device:", DEVICE, flush=True) +if DEVICE != "cuda": + raise RuntimeError( + "segger requires a CUDA GPU end-to-end and none is available. " + "Run this component on a GPU-equipped host." + ) + +# Working layout for the synthesized Xenium dataset and segger's output. +work_root = Path(meta.get("temp_dir") or "/tmp") / f"segger_{os.getpid()}" +work_root.mkdir(parents=True, exist_ok=True) +XENIUM_DIR = work_root / "xenium_input" +XENIUM_DIR.mkdir(parents=True, exist_ok=True) +SEGGER_OUT_DIR = work_root / "segger_output" + +# Load data +print('Reading input files', flush=True) +sdata = sd.read_zarr(par['input_ist']) +sdata_segm = sd.read_zarr(par['input_segmentation']) + +transcripts_coord_systems = sd.transformations.get_transformation(sdata[par["transcripts_key"]], get_all=True).keys() +assert par['coordinate_system'] in transcripts_coord_systems, f"Coordinate system '{par['coordinate_system']}' not found in input data." +segmentation_coord_systems = sd.transformations.get_transformation(sdata_segm["segmentation"], get_all=True).keys() +assert par['coordinate_system'] in segmentation_coord_systems, f"Coordinate system '{par['coordinate_system']}' not found in input data." + +########################################## +# boundaries.parquet (from segmentation) # +########################################## + +# segger consumes a Xenium-layout directory: it links the prior boundary +# polygons (cell/nucleus_boundaries.parquet) to the transcripts by cell_id. +# We derive those boundaries from the provided segmentation, bringing them +# into the transcripts' coordinate system so they align with the transcript +# x/y written below. +print('Computing boundaries from segmentation', flush=True) +boundaries = sd.to_polygons(sdata_segm['segmentation'])[['geometry']] + +print('Transforming boundaries to transcripts coordinate system', flush=True) +trans_segm_to_global = sd.transformations.get_transformation(sdata_segm["segmentation"], get_all=True)[par['coordinate_system']] +trans_global_to_transcripts = sd.transformations.get_transformation(sdata[par['transcripts_key']], get_all=True)[par['coordinate_system']].inverse() +trans = sd.transformations.Sequence([trans_segm_to_global, trans_global_to_transcripts]) +boundaries = sd.transform(boundaries, trans, par['coordinate_system']) + +# Xenium boundary schema: one row per polygon vertex (cell_id int64, vertex_x/y float64). +boundaries.index.name = "cell_id" +boundaries_df = boundaries['geometry'].get_coordinates().rename(columns={'x': 'vertex_x', 'y': 'vertex_y'}) +boundaries_df = boundaries_df.reset_index() +boundaries_df['cell_id'] = boundaries_df['cell_id'].astype(np.int64) +boundaries_df['vertex_x'] = boundaries_df['vertex_x'].astype(np.float64) +boundaries_df['vertex_y'] = boundaries_df['vertex_y'].astype(np.float64) +boundaries_df = boundaries_df[['cell_id', 'vertex_x', 'vertex_y']] +# segger expects both files to exist; we have no cell vs nucleus split, so +# the same polygon set drives whichever `--prediction-mode` selects. +boundaries_df.to_parquet(XENIUM_DIR / "nucleus_boundaries.parquet", index=False) +boundaries_df.to_parquet(XENIUM_DIR / "cell_boundaries.parquet", index=False) +del boundaries, boundaries_df + +####################### +# transcripts.parquet # +####################### + +# Basic prior assignment: label each transcript with the segmentation cell it +# falls in. Transcript x/y are transformed into the segmentation's pixel space +# only for this lookup; the parquet keeps the transcripts' native coordinates. +print('Transforming transcripts coordinates', flush=True) +# Multi-partition parquet files each start with a 0-based index, producing duplicate index +# values in the combined dask DataFrame. sd.transform() internally creates a pd.Series with +# index=transformed.index; when that dask index is computed it triggers an assign expression +# that fails on duplicate/lazy indices. Fix: materialize to pandas and rebuild as a single +# dask partition with a clean RangeIndex before transforming. +transcripts_input = sdata[par['transcripts_key']] +transcripts_reset = dask.dataframe.from_pandas(transcripts_input.compute().reset_index(drop=True), npartitions=1) +transcripts_reset.attrs.update(transcripts_input.attrs) +transcripts = sd.transform(transcripts_reset, to_coordinate_system=par['coordinate_system']) + +# In case of a translation transformation of the segmentation (e.g. crop of the data), we need to adjust the transcript coordinates +trans = sd.transformations.get_transformation(sdata_segm["segmentation"], get_all=True)[par['coordinate_system']].inverse() +transcripts = sd.transform(transcripts, trans, par['coordinate_system']) + +print('Assigning transcripts to segmentation cell ids', flush=True) +y_coords = transcripts.y.compute().to_numpy(dtype=np.int64) +x_coords = transcripts.x.compute().to_numpy(dtype=np.int64) +if isinstance(sdata_segm["segmentation"], xr.DataTree): + label_image = sdata_segm["segmentation"]["scale0"].image.to_numpy() +else: + label_image = sdata_segm["segmentation"].to_numpy() +# Clip to the label image bounds (transcripts can sit just outside after transform). +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 + +# Canonical transcripts frame (native coordinates, clean RangeIndex). Its row +# order matches prior_cell_id and the row_index segger reports back below. +tx_pd = transcripts_reset.compute() +n_tx = len(tx_pd) + +if "transcript_id" in tx_pd.columns: + transcript_id = tx_pd["transcript_id"].to_numpy().astype(np.uint64) +else: + transcript_id = np.arange(n_tx, dtype=np.uint64) + +tx_out = pd.DataFrame({ + "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 > 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)) +tx_out.to_parquet(XENIUM_DIR / "transcripts.parquet", index=False) +del transcripts, y_coords, x_coords, label_image + +################ +# Run segger # +################ + +SEGGER_OUT_DIR.mkdir(parents=True, exist_ok=True) +cmd = [ + "segger", "segment", + "-i", str(XENIUM_DIR), + "-o", str(SEGGER_OUT_DIR), + "--n-epochs", str(par["n_epochs"]), + "--prediction-expansion-ratio", str(par["prediction_expansion_ratio"]), + "--prediction-mode", par["prediction_mode"], +] +# cudf's `validate_setup` aborts `import cudf` on hosts without a usable NVIDIA +# driver; segger imports cudf eagerly during tiling, so skip that validation. +env = os.environ.copy() +env.setdefault("RAPIDS_NO_INITIALIZE", "1") +env.setdefault("CUDF_NO_INITIALIZE", "1") +env.setdefault("RMM_NO_INITIALIZE", "1") +print("Running segger:", " ".join(cmd), flush=True) +subprocess.run(cmd, check=True, env=env) + +seg_pq = SEGGER_OUT_DIR / "segger_segmentation.parquet" +if not seg_pq.exists(): + raise RuntimeError(f"Expected segger output not found: {seg_pq}") + +###################################################### +# Map segger assignments back onto the transcripts # +###################################################### + +print('Reading segger output', flush=True) +seg = pd.read_parquet(seg_pq) + +# Keep only confident assignments. Mirror segger's canonical valid_cell_id +# predicate (drop null / "-1" / "UNASSIGNED" / "NONE", case-insensitive) so +# the unassigned sentinels don't leak into the output. +_UNASSIGNED = {"-1", "UNASSIGNED", "NONE", "NAN", ""} +seg_id_str = seg["segger_cell_id"].astype(str).str.strip().str.upper() +keep_mask = ( + seg["keep"].astype(bool) + & seg["segger_cell_id"].notna() + & ~seg_id_str.isin(_UNASSIGNED) +) +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 (i.e. tx_pd +# row order). 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[row_idx] = codes.astype(np.int64) + 1 + +############################################# +# Build transcript-assignment output object # +############################################# + +transformations = sd.transformations.get_transformation(sdata[par['transcripts_key']], get_all=True) +tx_result = tx_pd.copy() +tx_result["cell_id"] = cell_id_per_tx +tx_result["transcript_id"] = transcript_id +transcripts_points = PointsModel.parse(tx_result, transformations=transformations) + +# Cell table: unique assigned cell ids (0/background excluded). +unique_cells = np.unique(cell_id_per_tx) +unique_cells = unique_cells[unique_cells != 0] +cell_id_col = pd.Series(unique_cells, name='cell_id', index=unique_cells) +assert 0 not in cell_id_col, "Found '0' in cell_id column of assignment output cell matrix" + +sdata_transcripts_only = sd.SpatialData( + points={ + "transcripts": transcripts_points + }, + tables={ + "table": ad.AnnData( + obs=pd.DataFrame(cell_id_col), + ) + } +) + +print('Write transcripts with cell ids', flush=True) +if os.path.exists(par["output"]): + shutil.rmtree(par["output"]) +sdata_transcripts_only.write(par['output']) + +# Best-effort cleanup so the worker doesn't accumulate state across runs. +try: + shutil.rmtree(work_root) +except OSError as e: + print(f"(non-fatal) cleanup of {work_root} failed: {e}", flush=True)