Merge branch 'main' into new-indexing

.flake8

@@ -1,2 +1,3 @@
 [flake8]
 max-line-length = 79
+extend-ignore = W503

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,10 +1,10 @@
 #!/usr/bin/env python
+"""Script to merge two labelmaps with their BIDS metadata."""
 from __future__ import annotations
 
 from argparse import ArgumentParser
 from os import PathLike
 from pathlib import Path
-from typing import Optional
 
 import nibabel as nib
 import numpy as np
@@ -50,10 +50,13 @@ def split_labels(
     atlas: np.ndarray,
     metadata: pd.DataFrame,
     prefix: str = "",
-    exceptions: Optional[list[str]] = None,
+    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])
     all_labels: pd.Series[int] = metadata["index"]
     if not set(unique_vals) <= set(all_labels):
@@ -68,8 +71,8 @@ def split_labels(
         dict(
             [
                 assemble_mask(atlas, metadata, label, prefix)
-                for label in all_labels
-                if label not in exceptions
+                for label in unique_vals
+                if label not in exceptions + drops
             ]
         )
     )
@@ -145,7 +148,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,
     )
@@ -163,21 +184,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

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "labelmerge"
-version = "0.2.1"
+version = "0.3.0"
 description = "Snakebids app for merging multiple label maps."
 authors = [
     "Jason Kai <[email protected]>",