Add segger transcript assignment (CLI-based)

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

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 ]

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)