Handle Sample Area Case (#163)

tansey-lab · Nov 28, 2024 · 66c6080 · 66c6080
1 parent d4b70fb
commit 66c6080
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 5 deletions.
diff --git a/nextflow/modules/nf-core/segment/main.nf b/nextflow/modules/nf-core/segment/main.nf
@@ -9,9 +9,9 @@ process SEGMENT {
     tuple val(meta), path(xenium_dir), path(dataset), path(predictions), path(cell_typing_results), val(tile_idx)
 
     output:
-    tuple val(meta), path("${prefix}/segmentation_tiles/segmentation_tile_*.h5")                    , emit: segmentation, optional: true
-    tuple val(meta), path("${prefix}/shape_tiles/shapes_tile_*.parquet")                     , emit: shapefile, optional: true
-    tuple val(meta), path("${prefix}/anndata_tiles/anndata_tile_*.h5ad")                       , emit: anndata, optional: true
+    tuple val(meta), path("${prefix}/segmentation_tiles/segmentation_tile_*.h5") , emit: segmentation, optional: true
+    tuple val(meta), path("${prefix}/shape_tiles/shapes_tile_*.parquet")         , emit: shapefile, optional: true
+    tuple val(meta), path("${prefix}/anndata_tiles/anndata_tile_*.h5ad")         , emit: anndata, optional: true
     path  "versions.yml"                                                         , emit: versions
 
 
@@ -21,6 +21,7 @@ process SEGMENT {
     script:
     prefix = task.ext.prefix ?: "${meta.id}"
     def args = task.ext.args ?: ""
+    def sample_area_arg = params.sample_area == null ? "" : "--sample-area ${params.sample_area}"
     """
     mkdir -p "${prefix}/segmentation_tiles"
     mkdir -p "${prefix}/shape_tiles"
@@ -38,6 +39,7 @@ process SEGMENT {
         --tile-height ${params.segmentation_tile_height} \
         --tile-width ${params.segmentation_tile_width} \
         --overlap-percentage ${params.overlap_percentage} \
+        ${sample_area_arg} \
         ${args}
 
     cat <<-END_VERSIONS > versions.yml

diff --git a/src/nuc2seg/cli/segment.py b/src/nuc2seg/cli/segment.py
@@ -20,6 +20,7 @@
     load_and_filter_transcripts_as_points,
 )
 from nuc2seg.utils import get_tile_bounds
+from nuc2seg.utils import create_shapely_rectangle
 from shapely import box
 
 logger = logging.getLogger(__name__)
@@ -126,11 +127,25 @@ def get_parser():
         type=float,
         default=0.25,
     )
+    parser.add_argument(
+        "--sample-area",
+        default=None,
+        type=str,
+        help='Crop the dataset to this rectangle, provided in in "x1,y1,x2,y2" format.',
+    )
     return parser
 
 
 def main():
     args = get_parser().parse_args()
+
+    if args.sample_area:
+        sample_area = create_shapely_rectangle(
+            *[float(x) for x in args.sample_area.split(",")]
+        )
+    else:
+        sample_area = None
+
     with h5py.File(args.dataset, "r") as f:
         base_width = f["labels"].shape[0]
         base_height = f["labels"].shape[1]
@@ -139,6 +154,10 @@ def main():
         dataset = Nuc2SegDataset.load_h5(args.dataset)
         transcripts = load_and_filter_transcripts_as_points(args.transcripts)
         predictions = ModelPredictions.load_h5(args.predictions)
+        if sample_area:
+            slide_bbox = sample_area.bounds
+        else:
+            slide_bbox = dataset.bbox
     else:
         dataset = Nuc2SegDataset.load_h5(
             args.dataset,
@@ -159,9 +178,20 @@ def main():
             )
         )
 
+        if sample_area:
+            slide_bbox = (
+                tile_bbox.bounds[0] + sample_area.bounds[0],
+                tile_bbox.bounds[1] + sample_area.bounds[1],
+                tile_bbox.bounds[2] + sample_area.bounds[0],
+                tile_bbox.bounds[3] + sample_area.bounds[1],
+            )
+        else:
+            slide_bbox = tile_bbox.bounds
+
         transcripts = load_and_filter_transcripts_as_points(
-            args.transcripts, sample_area=tile_bbox
+            args.transcripts, sample_area=box(*slide_bbox)
         )
+
         predictions = ModelPredictions.load_h5(
             args.predictions,
             tile_width=args.tile_width,
@@ -205,7 +235,7 @@ def main():
         logger.warning("No cells found in segmentation, exiting")
         return
 
-    gdf["geometry"] = gdf.translate(*dataset.bbox[:2])
+    gdf["geometry"] = gdf.translate(*slide_bbox[:2])
 
     logger.info("Creating anndata")
     ad = convert_transcripts_to_anndata(