diff --git a/pipeline/convert_npy_nrrd.py b/pipeline/convert_npy_nrrd.py index 4ed418c..c7f57f7 100644 --- a/pipeline/convert_npy_nrrd.py +++ b/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 diff --git a/pipeline/full_pipeline.py b/pipeline/full_pipeline.py index 17365ca..bffae7a 100644 --- a/pipeline/full_pipeline.py +++ b/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, ) diff --git a/pipeline/interpolate_gene.py b/pipeline/interpolate_gene.py index 82a0ee9..97f16fb 100644 --- a/pipeline/interpolate_gene.py +++ b/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