Merge pull request #85 from ChristianHinge/dev/april_03

Reworked workflow version 0.0.5

Author: ChristianHinge

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zerodose"
-version = "0.0.4"
+version = "0.0.5"
 description = "Zerodose"
 authors = ["Christian Hinge <christian.hinge@regionh.dk>"]
 license = "MIT"

src/zerodose/__main__.py

@@ -1,12 +1,14 @@
 """Command-line interface."""
-import re
 from typing import Sequence
 from typing import Union
 
 import click
 
-from zerodose.run_abnormality import create_abnormality_maps
-from zerodose.run_synthesize import synthesize_baselines
+from zerodose.pipeline import create_abnormality_maps
+from zerodose.pipeline import normalize_to_pet
+from zerodose.pipeline import run_full
+from zerodose.pipeline import synthesize_baselines
+from zerodose.utils import _create_output_fname
 
 
 @click.group()
@@ -79,50 +81,48 @@ def main() -> None:
     Useful if the input images already are in MNI space""",
 )
 
+outputspace_option = click.option(
+    "--space",
+    "outputspace",
+    type=click.Choice(
+        [
+            "mr",
+            "pet",
+        ]
+    ),
+    default="mr",
+    help="Which space to ",
+)
+
 
 @main.command()
 @mri_option
 @mask_option
 @sbpet_output_option
 @verbose_option
 @device_option
-@no_registration_option
 def syn(
     mri_fnames: Sequence[str],
     mask_fnames: Sequence[str],
     out_fnames: Union[Sequence[str], None] = None,
     verbose: bool = True,
     device: str = "cuda:0",
-    no_registration: bool = False,
 ) -> None:
     """Synthesize baseline PET images."""
     if out_fnames is None or len(out_fnames) == 0:
         out_fnames = [
-            _create_output_fname(mri_fname, suffix="_sb") for mri_fname in mri_fnames
+            _create_output_fname(mri_fname, suffix="_sbraw") for mri_fname in mri_fnames
         ]
 
-    do_registration = not no_registration
     synthesize_baselines(
         mri_fnames,
         mask_fnames,
         out_fnames,
         verbose=verbose,
         device=device,
-        do_registration=do_registration,
     )
 
 
-def _create_output_fname(mri_fname, suffix="_sb", file_type=".nii.gz"):
-    """Create output filename from input filename."""
-    out_fname = mri_fname
-    if out_fname.endswith(".nii.gz"):
-        out_fname = re.sub(".nii.gz$", "", out_fname)
-    if out_fname.endswith(".nii"):
-        out_fname = re.sub(".nii$", "", out_fname)
-    out_fname += suffix + file_type
-    return out_fname
-
-
 pet_option = click.option(
     "-p",
     "--pet",
@@ -149,6 +149,14 @@ def _create_output_fname(mri_fname, suffix="_sb", file_type=".nii.gz"):
     multiple=True,
 )
 
+no_resample_mask_option = click.option(
+    "--no-resample-mask",
+    "no_resample_mask",
+    is_flag=True,
+    default=False,
+    help="Print verbose output.",
+)
+
 
 @pet_option
 @sbpet_option
@@ -179,3 +187,180 @@ def abn(
         verbose=verbose,
         device=device,
     )
+
+
+no_resample_sbpet_option = click.option(
+    "--no-resample-sbpet",
+    "no_resample_sbpet",
+    is_flag=True,
+    default=False,
+    help="Print verbose output.",
+)
+
+
+@pet_option
+@sbpet_option
+@mask_option
+@abn_output_option
+@verbose_option
+@device_option
+@main.command()
+def norm(
+    pet_fnames: Sequence[str],
+    sbpet_fnames: Sequence[str],
+    mask_fnames: Sequence[str],
+    out_fnames: Union[Sequence[str], None] = None,
+    verbose: bool = False,
+    device: str = "cuda:0",
+):
+    """Normalize sbPET images to PET images."""
+    if out_fnames is None or len(out_fnames) == 0:
+        out_fnames = [
+            _create_output_fname(pet_name, suffix="_sb") for pet_name in pet_fnames
+        ]
+
+    normalize_to_pet(
+        pet_fnames=pet_fnames,
+        sbpet_fnames=sbpet_fnames,
+        mask_fnames=mask_fnames,
+        out_fnames=out_fnames,
+        verbose=verbose,
+        device=device,
+    )
+
+
+no_image_option = click.option(
+    "--no-image",
+    "no_image",
+    is_flag=True,
+    default=False,
+    help="Print verbose output.",
+)
+
+mri_option_single = click.option(
+    "-i",
+    "--mri",
+    "mri_fname",
+    type=click.Path(exists=True),
+    required=True,
+)
+mask_option_single = click.option(
+    "-m",
+    "--mask",
+    "mask_fname",
+    type=click.Path(exists=True),
+    required=True,
+)
+
+pet_option_single = click.option(
+    "-p",
+    "--pet",
+    "pet_fname",
+    type=click.Path(exists=True),
+    required=False,
+)
+
+sbpet_output_option_single = click.option(
+    "-os",
+    "--out-sbpet",
+    "out_sbpet",
+    type=click.Path(),
+    required=False,
+)
+
+abn_output_option_single = click.option(
+    "-oa",
+    "--out-abn",
+    "out_abn",
+    type=click.Path(),
+    required=False,
+)
+
+img_output_option_single = click.option(
+    "-oi",
+    "--out-img",
+    "out_img",
+    type=click.Path(),
+    required=False,
+)
+
+no_abnormality_option = click.option(
+    "--no-abn",
+    "no_abnormality",
+    is_flag=True,
+    default=False,
+    help="Print verbose output.",
+)
+
+no_normalization_option = click.option(
+    "--no-norm",
+    "no_normalization",
+    is_flag=True,
+    default=False,
+    help="Print verbose output.",
+)
+
+no_image_option = click.option(
+    "--no-img",
+    "no_image",
+    is_flag=True,
+    default=False,
+    help="Print verbose output.",
+)
+
+
+@mri_option_single
+@mask_option_single
+@pet_option_single
+@sbpet_output_option_single
+@abn_output_option_single
+@img_output_option_single
+@no_registration_option
+@no_abnormality_option
+@no_normalization_option
+@no_image_option
+@verbose_option
+@device_option
+@outputspace_option
+@main.command()
+def run(
+    mri_fname,
+    mask_fname,
+    pet_fname,
+    out_sbpet,
+    out_abn,
+    out_img,
+    no_registration,
+    no_abnormality,
+    no_normalization,
+    no_image,
+    verbose,
+    device,
+    outputspace,
+):
+    """Run full pipeline."""
+    do_registration = not no_registration
+    do_abnormality = not no_abnormality
+    do_normalization = not no_normalization
+    do_image = not no_image
+
+    if out_sbpet is None:
+        out_sbpet = _create_output_fname(pet_fname, suffix="_sb")
+    if out_abn is None:
+        out_abn = _create_output_fname(pet_fname, suffix="_abn")
+
+    run_full(
+        mri_fname=mri_fname,
+        mask_fname=mask_fname,
+        out_sbpet=out_sbpet,
+        pet_fname=pet_fname,
+        out_abn=out_abn,
+        out_img=out_img,
+        do_registration=do_registration,
+        do_abnormality=do_abnormality,
+        do_normalization=do_normalization,
+        do_image=do_image,
+        verbose=verbose,
+        device=device,
+        outputspace=outputspace,
+    )

src/zerodose/dataset.py

@@ -1,25 +1,21 @@
 """Dataset class for the ZeroDose project."""
 from typing import Any
 from typing import Dict
-from typing import List
 
 import torchio as tio
 
-from zerodose import utils
-from zerodose.niftyreg_wrapper import NiftyRegistration
 from zerodose.processing import Binarize
-from zerodose.processing import PadAndCropToMNI
+from zerodose.processing import PadAndCropMNI
 from zerodose.processing import ToFloat32
 
 
 class SubjectDataset(tio.data.SubjectsDataset):
     """Dataset class for the ZeroDose project."""
 
-    def __init__(self, mri_fnames, mask_fnames, out_fnames, do_registration=True):
+    def __init__(self, mri_fnames, mask_fnames, out_fnames):
         """Initialize the dataset."""
-        transforms = self._get_augmentation_transform_val(
-            do_registration=do_registration
-        )
+        transforms = self._get_augmentation_transform_val()
+
         subjects = [
             self._make_subject_predict(mr_f, ma_f, ou_f)
             for mr_f, ma_f, ou_f in zip(  # noqa
@@ -28,7 +24,6 @@ def __init__(self, mri_fnames, mask_fnames, out_fnames, do_registration=True):
                 out_fnames,
             )
         ]
-
         super().__init__(subjects, transforms)
 
     def _make_subject_dict(self, mr_path, mask_path) -> dict:
@@ -44,23 +39,14 @@ def _make_subject_dict(self, mr_path, mask_path) -> dict:
     def _make_subject_predict(self, mr_path, mask_path, out_fname) -> tio.Subject:
         subject_dict = self._make_subject_dict(mr_path, mask_path)
         subject_dict["out_fname"] = out_fname
-
         return tio.Subject(subject_dict)
 
-    def _get_augmentation_transform_val(self, do_registration=True) -> tio.Compose:
-        augmentations: List[tio.Transform] = []
-
-        if do_registration:
-            ref = utils.get_mni_template()
-            augmentations.append(NiftyRegistration(floating_image="mr", ref=ref))
-
-        augmentations.extend(
-            [
-                Binarize(include=["mask"]),
-                tio.transforms.ZNormalization(include=["mr"], masking_method="mask"),
-                PadAndCropToMNI(include=["mr", "mask"]),
-                ToFloat32(include=["mr"]),
-            ]
-        )
+    def _get_augmentation_transform_val(self) -> tio.Compose:
+        augmentations = [
+            Binarize(include=["mask"]),
+            tio.transforms.ZNormalization(include=["mr"], masking_method="mask"),
+            PadAndCropMNI(),
+            ToFloat32(include=["mr"]),
+        ]
 
         return tio.Compose(augmentations)