Add format option on the full pipeline

pipeline/convert_npy_nrrd.py

@@ -14,11 +14,29 @@
 # limitations under the License.
 
 import argparse
+from collections import OrderedDict
 from pathlib import Path
 
 import nrrd
 import numpy as np
 
+# Match voxcell requirements:
+HEADER = OrderedDict(
+    [
+        ("type", "uint32"),
+        ("dimension", None),
+        ("space dimension", 3),
+        ("sizes", None),
+        (
+            "space directions",
+            np.array([[25.0, 0.0, 0.0], [0.0, 25.0, 0.0], [0.0, 0.0, 25.0]]),
+        ),
+        ("endian", "little"),
+        ("encoding", "gzip"),
+        ("space origin", np.array([0.0, 0.0, 0.0])),
+    ]
+)
+
 
 def parse_args():
     """Parse command line arguments.
@@ -44,13 +62,26 @@ def parse_args():
         Path to the output volume.
         """,
     )
+    parser.add_argument(
+        "-h",
+        "--header",
+        type=Path,
+        help=f"""\
+        If specified, the header {HEADER} is saved.
+        """,
+    )
     return parser.parse_args()
 
 
-def main(input_path: Path, output_path: Path) -> int:
+def main(input_path: Path, output_path: Path, header: bool) -> int:
 
     array = np.load(input_path)
-    nrrd.write(str(output_path), array)
+    if header:
+        HEADER["dimension"] = len(array.shape)
+        HEADER["sizes"] = np.array(array.shape)
+        nrrd.write(str(output_path), array, header=HEADER)
+    else:
+        nrrd.write(str(output_path), array)
     return 0
 
 

pipeline/full_pipeline.py

@@ -106,6 +106,16 @@ def parse_args():
         Path of the interpolator checkpoints.
         """,
     )
+    parser.add_argument(
+        "-s",
+        "--saving-format",
+        type=str,
+        choices=("nrrd", "npy"),
+        default="npy",
+        help="""\
+        Format to save the output volumes.
+        """,
+    )
     parser.add_argument(
         "-e",
         "--expression",
@@ -135,6 +145,7 @@ def main(
     interpolator_name: str,
     interpolator_checkpoint: Path | str | None,
     output_dir: Path | str,
+    saving_format: str,
     expression: bool = False,
     force: bool = False,
 ) -> int:
@@ -226,6 +237,7 @@ def main(
                 metadata_path=aligned_results_dir / f"{experiment_id}-metadata.json",
                 interpolator_name=interpolator_name,
                 interpolator_checkpoint=interpolator_checkpoint,
+                saving_format=saving_format,
                 reference_path=nissl_path,
                 output_dir=interpolation_results_dir,
             )

pipeline/interpolate_gene.py

@@ -63,6 +63,15 @@ def parse_args():
         Path of the interpolator checkpoints.
         """,
     )
+    parser.add_argument(
+        "--saving-format",
+        type=str,
+        choices=("nrrd", "npy"),
+        default="npy",
+        help="""\
+        Format to save the output volumes.
+        """,
+    )
     parser.add_argument(
         "--reference-path",
         type=Path,
@@ -137,10 +146,12 @@ def main(
     metadata_path: Path | str,
     interpolator_name: str,
     interpolator_checkpoint: str | Path | None,
+    saving_format: str,
     reference_path: str | Path,
     output_dir: Path | str | None = None,
 ) -> int:
     """Implement main function."""
+    import nrrd
     import numpy as np
     from atlinter.data import GeneDataset
     from utils import check_and_load
@@ -194,12 +205,18 @@ def main(
 
     output_dir = Path(output_dir)
     output_dir.mkdir(parents=True, exist_ok=True)
+    output_path = str(output_dir / f"{experiment_id}-{interpolator_name}-interpolated-{image_type}")
 
-    np.save(
-        output_dir
-        / f"{experiment_id}-{interpolator_name}-interpolated-{image_type}.npy",
-        predicted_volume,
-    )
+    if saving_format == "npy":
+        np.save(
+            output_path + ".npy",
+            predicted_volume,
+        )
+    else:
+        from convert_npy_nrrd import HEADER
+        HEADER["dimension"] = len(predicted_volume.shape)
+        HEADER["sizes"] = np.array(predicted_volume.shape)
+        nrrd.write(output_path + ".nrrd", predicted_volume, header=HEADER)
 
     return 0