Merge pull request #35 from khanlab/drops

Add setting to allow labels to be dropped from input maps

labelmerge/config/snakebids.yml

@@ -80,7 +80,13 @@ parse_args:
     help: Space separated integer labels from the base labelmap to keep over overlay labels at the same voxels.
     nargs: '*'
   --overlay_exceptions:
-    help: Space separated integer labels from the overlay image to discard.
+    help: Space separated integer labels from the overlay image to be overwritten by base labels at the same voxels.
+    nargs: '*'
+  --base_drops:
+    help: Space separated integer labels from the base image to drop from the output.
+    nargs: '*'
+  --overlay_drops:
+    help: Space separated integer labels from the overlay image to drop from the output.
     nargs: '*'
 
 # Workflow specific config

labelmerge/workflow/rules/label_harmonization.smk

@@ -83,10 +83,17 @@ rule merge_labels:
         overlay_exceptions=f"--overlay_exceptions {' '.join(config['overlay_exceptions'])}"
         if config.get("overlay_exceptions")
         else "",
+        base_drops=f"--base_drops {' '.join(config['base_drops'])}"
+        if config.get("base_drops")
+        else "",
+        overlay_drops=f"--overlay_drops {' '.join(config['overlay_drops'])}"
+        if config.get("overlay_drops")
+        else "",
     resources:
         script=str(Path(workflow.basedir) / "scripts" / "labelmerge.py"),
     shell:
         "python3 {resources.script} {input.base_map} {input.base_metadata} "
         "{input.overlay_map} {input.overlay_metadata} "
         "{output.merged_map} {output.merged_metadata} "
-        "{params.base_exceptions} {params.overlay_exceptions}"
+        "{params.base_exceptions} {params.overlay_exceptions} "
+        "{params.base_drops} {params.overlay_drops}"

labelmerge/workflow/scripts/labelmerge.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python
+"""Script to merge two labelmaps with their BIDS metadata."""
 from __future__ import annotations
 
 from argparse import ArgumentParser
@@ -49,15 +50,20 @@ def split_labels(
     atlas: np.ndarray,
     metadata: pd.DataFrame,
     prefix: str = "",
-    exceptions: list[str] = [],
+    exceptions: list[int] | None = None,
+    drops: list[int] | None = None,
 ) -> list[xr.Dataset]:
+    if exceptions is None:
+        exceptions = []
+    if drops is None:
+        drops = []
     unique_vals = np.unique(atlas[atlas > 0])
     normal_ds = xr.Dataset(
         dict(
             [
                 assemble_mask(atlas, metadata, label, prefix)
                 for label in unique_vals
-                if label not in exceptions
+                if label not in exceptions + drops
             ]
         )
     )
@@ -133,7 +139,25 @@ def gen_parser() -> ArgumentParser:
         nargs="*",
         help=(
             "Space separated list of integer labels from the overlay image to "
-            "discard."
+            "be overwritten by labels from the base image."
+        ),
+        type=int,
+    )
+    parser.add_argument(
+        "--base_drops",
+        nargs="*",
+        help=(
+            "Space separated list of integer labels from the base image to "
+            "drop from the output map."
+        ),
+        type=int,
+    )
+    parser.add_argument(
+        "--overlay_drops",
+        nargs="*",
+        help=(
+            "Space separated list of integer labels from the overlay image to "
+            "drop from the output map."
         ),
         type=int,
     )
@@ -151,21 +175,28 @@ def main():
     overlay_exceptions = (
         args.overlay_exceptions if args.overlay_exceptions else []
     )
+    base_drops = args.base_drops if args.base_drops else []
+    overlay_drops = args.overlay_drops if args.overlay_drops else []
     base_datasets = split_labels(
         base_data,
         base_metadata,
         prefix="base ",
         exceptions=base_exceptions,
+        drops=base_drops,
     )
     overlay_datasets = split_labels(
         overlay_data,
         overlay_metadata,
         prefix="overlay ",
         exceptions=overlay_exceptions,
+        drops=overlay_drops,
     )
     # Note that overlay exceptions are ignored
     merged_map, merged_metadata = merge_labels(
-        base_datasets[0:1] + overlay_datasets[0:1] + base_datasets[1:]
+        overlay_datasets[1:]
+        + base_datasets[0:1]
+        + overlay_datasets[0:1]
+        + base_datasets[1:]
     )
     merged_img = nib.Nifti1Image(
         dataobj=merged_map, affine=base_affine, header=base_header