Prototype cli

antspynet/cli/mouse_brain_extraction.py

@@ -0,0 +1,84 @@
+#!/usr/bin/env python
+
+import argparse
+import sys
+
+import ants
+import antspynet
+
+
+def build_parser():
+    parser = argparse.ArgumentParser(
+        description="Run mouse brain extraction with ANTsPyNet."
+    )
+    parser.add_argument(
+        "input_image",
+        type=str,
+        help="Path to the input mouse image.",
+    )
+    parser.add_argument(
+        "output_image",
+        type=str,
+        help="Path to write the extracted probability image.",
+    )
+    parser.add_argument(
+        "--modality",
+        default="t2",
+        choices=["t2", "ex5coronal", "ex5sagittal"],
+        help='Mouse image modality passed to mouse_brain_extraction (default: "t2").',
+    )
+    parser.add_argument(
+        "--isotropic-output",
+        action="store_true",
+        help="Return the probability image in isotropic space before writing it.",
+    )
+    parser.add_argument(
+        "--axis",
+        type=int,
+        default=2,
+        choices=[0, 1, 2],
+        help="Axis index used for ex5 modalities (default: 2).",
+    )
+    parser.add_argument(
+        "--verbose",
+        action="store_true",
+        help="Print progress while running the model.",
+    )
+    return parser
+
+
+def run(input_image, output_image, modality="t2", isotropic_output=False, axis=2, verbose=False):
+    image = ants.image_read(input_image)
+    probability_image = antspynet.mouse_brain_extraction(
+        image,
+        modality=modality,
+        return_isotropic_output=isotropic_output,
+        which_axis=axis,
+        verbose=verbose,
+    )
+    ants.image_write(probability_image, output_image)
+    return output_image
+
+
+def main(argv=None):
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    try:
+        run(
+            args.input_image,
+            args.output_image,
+            modality=args.modality,
+            isotropic_output=args.isotropic_output,
+            axis=args.axis,
+            verbose=args.verbose,
+        )
+    except Exception as error:
+        print(f"Error: {error}", file=sys.stderr)
+        return 1
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

pyproject.toml

@@ -38,3 +38,6 @@ find = { include = ["antspynet","antspynet.*"] }
 [tool.setuptools]
 license-files = []
 
+[project.scripts]
+ants_mouse_brain_extraction = "antspynet.cli.mouse_brain_extraction:main"
+

tests/test_mouse_cli.py

@@ -0,0 +1,84 @@
+import os
+import tempfile
+import unittest
+from contextlib import redirect_stderr
+from io import StringIO
+from unittest.mock import patch
+
+import ants
+import numpy as np
+
+from antspynet.cli.mouse_brain_extraction import main
+
+
+class TestMouseBrainExtractionCli(unittest.TestCase):
+    def test_main_routes_to_mouse_brain_extraction(self):
+        input_image = ants.from_numpy(np.ones((4, 4, 4)))
+        output_image = ants.from_numpy(np.full((4, 4, 4), 0.25))
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            input_path = os.path.join(tmpdir, "input.nii.gz")
+            output_path = os.path.join(tmpdir, "output.nii.gz")
+
+            calls = {}
+
+            def fake_image_read(path):
+                calls["image_read"] = path
+                return input_image
+
+            def fake_mouse_brain_extraction(
+                image,
+                modality,
+                return_isotropic_output,
+                which_axis,
+                verbose,
+            ):
+                calls["mouse_brain_extraction"] = {
+                    "image": image,
+                    "modality": modality,
+                    "return_isotropic_output": return_isotropic_output,
+                    "which_axis": which_axis,
+                    "verbose": verbose,
+                }
+                return output_image
+
+            def fake_image_write(image, path):
+                calls["image_write"] = {"image": image, "path": path}
+
+            with patch("antspynet.cli.mouse_brain_extraction.ants.image_read", side_effect=fake_image_read), \
+                 patch("antspynet.cli.mouse_brain_extraction.antspynet.mouse_brain_extraction", side_effect=fake_mouse_brain_extraction), \
+                 patch("antspynet.cli.mouse_brain_extraction.ants.image_write", side_effect=fake_image_write):
+                exit_code = main([
+                    input_path,
+                    output_path,
+                    "--modality", "t2",
+                    "--isotropic-output",
+                    "--axis", "1",
+                    "--verbose",
+                ])
+
+        self.assertEqual(exit_code, 0)
+        self.assertEqual(calls["image_read"], input_path)
+        self.assertEqual(calls["image_write"]["path"], output_path)
+        self.assertIs(calls["image_write"]["image"], output_image)
+        self.assertEqual(calls["mouse_brain_extraction"]["modality"], "t2")
+        self.assertTrue(calls["mouse_brain_extraction"]["return_isotropic_output"])
+        self.assertEqual(calls["mouse_brain_extraction"]["which_axis"], 1)
+        self.assertTrue(calls["mouse_brain_extraction"]["verbose"])
+
+    def test_main_returns_nonzero_and_writes_error_to_stderr_on_failure(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            input_path = os.path.join(tmpdir, "input.nii.gz")
+            output_path = os.path.join(tmpdir, "output.nii.gz")
+
+            with patch("antspynet.cli.mouse_brain_extraction.run", side_effect=RuntimeError("boom")):
+                stderr = StringIO()
+                with redirect_stderr(stderr):
+                    exit_code = main([input_path, output_path])
+
+        self.assertEqual(exit_code, 1)
+        self.assertEqual(stderr.getvalue(), "Error: boom\n")
+
+
+if __name__ == "__main__":
+    unittest.main()