Add InstanSeg nucleus segmentation for H&E whole-slide images

.gitignore

@@ -32,3 +32,14 @@ xenium_example/
 /graphify-out/
 /slurm/
 /slurm-logs/
+
+# ...except the segmentation pipeline, which is part of the package's public
+# surface rather than scratch: the CLI wrappers are documented entry points and
+# the SLURM scripts are the supported way to run them on a cluster.
+!/scripts/instanseg_segment.py
+!/scripts/geojson_to_spatialdata.py
+!/slurm/
+/slurm/*
+!/slurm/segment_node_worker.sh
+!/slurm/segment_slurm.sh
+!/slurm/SEGMENTATION_PLAN.md

pyproject.toml

@@ -44,6 +44,19 @@ dependencies = [
   "tifffile",
 ]
 optional-dependencies.czi = [ "bioio", "bioio-czi" ]
+# Nucleus segmentation on H&E whole-slide images. Kept optional: it pulls
+# torch (multi-GB CUDA wheels), which most spatialrefinery uses do not need.
+optional-dependencies.segmentation = [
+  # rasterio + geojson are what InstanSeg's save_geojson=True actually imports.
+  # They are NOT taken via `instanseg-torch[io]`: that extra also carries
+  # `zarr>=2.0.0,<3`, which silently downgrades zarr/numcodecs/tiffslide and
+  # breaks spatialdata. Its stated reason ("tiffslide doesn't support zarr v3
+  # yet") is stale as of tiffslide 4.0 (Bayer-Group/tiffslide#97).
+  "geojson>=3",
+  "instanseg-torch>=0.1.1",
+  "rasterio>=1.3",
+  "tiffslide>=4",
+]
 # https://docs.pypi.org/project_metadata/#project-urls
 urls.Documentation = "https://spatialrefinery.readthedocs.io/"
 urls.Homepage = "https://github.com/peng-lab/spatialrefinery"
@@ -153,11 +166,15 @@ module = [
   "bioio.*",
   "dask_image.*",
   "geopandas.*",
+  # From the optional `segmentation` extra, so absent in the typecheck env.
+  "instanseg.*",
   "pyarrow.*",
   "scipy.*",
   "shapely.*",
   "spatialdata.*",
   "spatialdata_io.*",
+  "tiffslide.*",
+  "torch.*",
 ]
 ignore_missing_imports = true
 

scripts/geojson_to_spatialdata.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+r"""Convert a nucleus-segmentation GeoJSON plus its slide into a SpatialData zarr.
+
+Thin CLI wrapper around
+`spatialrefinery.segmentation.to_spatialdata.geojson_to_spatialdata`.
+Pairs with `instanseg_segment.py`, consuming the `cells.geojson` it writes.
+
+Usage
+-----
+    python geojson_to_spatialdata.py \\
+        --geojson-path results/slide.svs/cells.geojson \\
+        --zarr-outdir zarrs/ --wsi-path slide.svs --template-adata template.h5ad
+"""
+
+import argparse
+import logging
+import sys
+from pathlib import Path
+
+from spatialrefinery.segmentation.to_spatialdata import geojson_to_spatialdata
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Build the CLI parser."""
+    parser = argparse.ArgumentParser(
+        description="SpatialData conversion for a single segmented WSI.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("--geojson-path", required=True)
+    parser.add_argument("--zarr-outdir", required=True)
+    parser.add_argument("--wsi-path", required=True)
+    parser.add_argument("--template-adata", required=True)
+    parser.add_argument("--no-zip", action="store_true", help="Skip the .zarr.zip archive")
+    parser.add_argument("--no-skip-existing", action="store_true", help="Rebuild even if the zarr exists")
+    return parser
+
+
+def main() -> None:
+    """Convert one slide's segmentation into a SpatialData zarr."""
+    args = build_parser().parse_args()
+
+    geojson_path = Path(args.geojson_path)
+    wsi_path = Path(args.wsi_path)
+    template_adata = Path(args.template_adata)
+
+    for label, path in (("GeoJSON", geojson_path), ("WSI", wsi_path), ("Template AnnData", template_adata)):
+        if not path.exists():
+            logger.error("%s file not found: %s", label, path)
+            sys.exit(1)
+
+    # Named with the full filename, matching the segmentation stage's layout.
+    zarr_path = Path(args.zarr_outdir) / f"{wsi_path.name}.zarr"
+    if zarr_path.exists() and not args.no_skip_existing:
+        logger.info("Skipping %s: %s already exists", wsi_path.name, zarr_path)
+        sys.exit(0)
+
+    Path(args.zarr_outdir).mkdir(parents=True, exist_ok=True)
+
+    try:
+        geojson_to_spatialdata(
+            geojson_path=geojson_path,
+            zarr_path=zarr_path,
+            image_path=wsi_path,
+            template_adata_path=template_adata,
+            write_zip=not args.no_zip,
+        )
+    except Exception:
+        logger.exception("Conversion failed for %s", wsi_path.name)
+        sys.exit(1)
+
+    print(f"ZARR_PATH={zarr_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/instanseg_segment.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+"""Segment nuclei in one H&E whole-slide image with InstanSeg.
+
+Thin CLI wrapper around `spatialrefinery.segmentation.instanseg.segment_wsi`.
+
+Prints `GEOJSON_PATH=<path>` on success -- the SLURM worker greps for it to
+hand the result to the conversion stage.
+
+Usage
+-----
+    python instanseg_segment.py --wsi-path slide.svs --outdir results/
+    python instanseg_segment.py --wsi-path slide.ome.tif --outdir results/ --wsi-mpp 0.27
+"""
+
+import argparse
+import logging
+import sys
+from pathlib import Path
+
+from spatialrefinery.segmentation.instanseg import (
+    DEFAULT_DETECTION_SIZE,
+    DEFAULT_MODEL,
+    DEFAULT_OVERLAP,
+    DEFAULT_TILE_SIZE,
+    segment_wsi,
+)
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Build the CLI parser."""
+    parser = argparse.ArgumentParser(
+        description="InstanSeg nucleus segmentation for a single WSI.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("--wsi-path", required=True)
+    parser.add_argument("--outdir", required=True)
+    parser.add_argument("--gpu-id", type=int, default=0, help="CUDA index; -1 forces CPU")
+    parser.add_argument("--model", default=DEFAULT_MODEL)
+    parser.add_argument("--wsi-mpp", type=float, default=None, help="Microns per pixel; read from metadata if omitted")
+    parser.add_argument("--tile-size", type=int, default=DEFAULT_TILE_SIZE)
+    parser.add_argument("--overlap", type=int, default=DEFAULT_OVERLAP)
+    parser.add_argument("--detection-size", type=int, default=DEFAULT_DETECTION_SIZE)
+    parser.add_argument(
+        "--clahe",
+        type=float,
+        default=None,
+        metavar="CLIP",
+        help=(
+            "Run CLAHE at this clip limit over each tile before inference (e.g. 2.0). "
+            "Off by default; helps on weakly haematoxylin-stained slides where pale "
+            "nuclei are missed."
+        ),
+    )
+    parser.add_argument(
+        "--seed-threshold",
+        type=float,
+        default=None,
+        help="Override the model seed threshold (default 0.7). Lower detects fainter nuclei.",
+    )
+    parser.add_argument(
+        "--no-otsu",
+        action="store_true",
+        help="Segment every tile instead of only those inside the tissue mask",
+    )
+    parser.add_argument("--no-skip-existing", action="store_true", help="Re-segment even if cells.geojson exists")
+    return parser
+
+
+def main() -> None:
+    """Run segmentation for one slide and print its GeoJSON path."""
+    args = build_parser().parse_args()
+
+    wsi_path = Path(args.wsi_path)
+    if not wsi_path.exists():
+        logger.error("WSI file not found: %s", wsi_path)
+        sys.exit(1)
+
+    try:
+        geojson_path = segment_wsi(
+            wsi_path,
+            args.outdir,
+            pixel_size=args.wsi_mpp,
+            gpu_id=None if args.gpu_id < 0 else args.gpu_id,
+            model_type=args.model,
+            tile_size=args.tile_size,
+            overlap=args.overlap,
+            detection_size=args.detection_size,
+            use_otsu_threshold=not args.no_otsu,
+            clahe_clip=args.clahe,
+            seed_threshold=args.seed_threshold,
+            skip_existing=not args.no_skip_existing,
+        )
+    except Exception:
+        logger.exception("Segmentation failed for %s", wsi_path.name)
+        sys.exit(1)
+
+    # Contract with slurm/segment_node_worker.sh -- keep this line last.
+    print(f"GEOJSON_PATH={geojson_path}")
+
+
+if __name__ == "__main__":
+    main()