Merge pull request #89 from ChristianHinge/dev/visualizations

dev/visualizations

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zerodose"
-version = "0.0.5"
+version = "0.0.6"
 description = "Zerodose"
 authors = ["Christian Hinge <[email protected]>"]
 license = "MIT"
@@ -21,6 +21,11 @@ click = ">=8.0.1"
 torch = ">=1.0.0"
 torchio = "^0.18.0"
 niftyreg = "^1.5.70rc1"
+matplotlib = {version = "^3.7.1", optional = true}
+opencv-python = {version = "^4.7.0.72", optional = true}
+
+[tool.poetry.extras]
+plotting = ["matplotlib","opencv-python"]
 
 [tool.poetry.dev-dependencies]
 Pygments = ">=2.10.0"

src/zerodose/__main__.py

@@ -6,7 +6,7 @@
 
 from zerodose.pipeline import create_abnormality_maps
 from zerodose.pipeline import normalize_to_pet
-from zerodose.pipeline import run_full
+from zerodose.pipeline import run_with_registration
 from zerodose.pipeline import synthesize_baselines
 from zerodose.utils import _create_output_fname
 
@@ -308,58 +308,61 @@ def norm(
     help="Print verbose output.",
 )
 
+no_reg_pet_to_mr_option = click.option(
+    "--no-pet-rigid",
+    "no_reg_pet_to_mr",
+    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_reg_pet_to_mr_option
 @no_image_option
 @verbose_option
 @device_option
 @outputspace_option
 @main.command()
-def run(
+def pipeline(
     mri_fname,
     mask_fname,
     pet_fname,
     out_sbpet,
     out_abn,
     out_img,
-    no_registration,
-    no_abnormality,
-    no_normalization,
+    no_reg_pet_to_mr,
     no_image,
     verbose,
     device,
     outputspace,
 ):
     """Run full pipeline."""
-    do_registration = not no_registration
-    do_abnormality = not no_abnormality
-    do_normalization = not no_normalization
+    reg_pet_to_mr = not no_reg_pet_to_mr
     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")
+    if out_img is None and do_image:
+        out_img = _create_output_fname(
+            pet_fname, suffix="_abn_figure", file_type=".png"
+        )
 
-    run_full(
+    run_with_registration(
         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,
+        reg_pet_to_mri=reg_pet_to_mr,
         verbose=verbose,
         device=device,
         outputspace=outputspace,

src/zerodose/dataset.py

@@ -12,33 +12,28 @@
 class SubjectDataset(tio.data.SubjectsDataset):
     """Dataset class for the ZeroDose project."""
 
-    def __init__(self, mri_fnames, mask_fnames, out_fnames):
+    def __init__(self, mri_fnames, mask_fnames, pet_fnames=None):
         """Initialize the dataset."""
         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
+            self._make_subject_predict(mr_f, ma_f)
+            for mr_f, ma_f in zip(  # noqa
                 mri_fnames,
                 mask_fnames,
-                out_fnames,
             )
         ]
-        super().__init__(subjects, transforms)
-
-    def _make_subject_dict(self, mr_path, mask_path) -> dict:
-        subject_dict: Dict[Any, Any] = {}
-        mri = mr_path
-        mask = mask_path
 
-        subject_dict["mr"] = tio.ScalarImage(mri)
-        subject_dict["mask"] = tio.LabelMap(mask)
+        if pet_fnames is not None:
+            for sub, pet_fname in zip(subjects, pet_fnames):  # noqa
+                sub.add_image(tio.ScalarImage(pet_fname), "pet")
 
-        return subject_dict
+        super().__init__(subjects, transforms)
 
-    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
+    def _make_subject_predict(self, mr_path, mask_path) -> tio.Subject:
+        subject_dict: Dict[Any, Any] = {}
+        subject_dict["mr"] = tio.ScalarImage(mr_path)
+        subject_dict["mask"] = tio.LabelMap(mask_path)
         return tio.Subject(subject_dict)
 
     def _get_augmentation_transform_val(self) -> tio.Compose:

src/zerodose/inference.py

@@ -61,3 +61,15 @@ def synthesize(mr, model, device="cuda:0"):
     )
 
     return sbpet_tensor
+
+
+def get_inference_params():
+    """Returns the default parameters for inference stitching."""
+    stride = 2
+    params = {
+        "patch_size": (32, 192, 192),
+        "patch_overlap": (32 - stride, 192 - stride, 192 - stride),
+        "batch_size": 1,
+        "sd_weight": 5,
+    }
+    return params

src/zerodose/models.py

@@ -27,6 +27,11 @@ def forward(self, mrs: torch.Tensor) -> torch.Tensor:
         """Forward pass through the model."""
         return self.generator(mrs)
 
+    @staticmethod
+    def get_default():
+        """Get the default model."""
+        return utils.get_model()
+
 
 class AbnormalityMap(nn.Module):
     """Abnormality map generator."""
@@ -38,6 +43,11 @@ def __init__(self, sigma_smooth=3) -> None:
             channels=1, kernel_size=5 * sigma_smooth, sigma=sigma_smooth, dim=3
         )
 
+    @staticmethod
+    def get_default():
+        """Get the default model."""
+        return AbnormalityMap(sigma_smooth=3)
+
     def forward(
         self, pet: torch.Tensor, sbpet: torch.Tensor, mask: torch.Tensor
     ) -> torch.Tensor:
@@ -88,3 +98,8 @@ def _scale_sbpet(self, pet, sbpet, mask):
             sbpet[~mask] = pet[~mask]
 
         return sbpet
+
+    @staticmethod
+    def get_default():
+        """Get the default model."""
+        return QuantileNormalization(quantile=0.97, sigma_normalization=3)

src/zerodose/pipeline/__init__.py

@@ -1,6 +1,7 @@
 """Pipeline for ZeroDose project."""
 from .abnormality import create_abnormality_maps
-from .full import run_full
+from .full import run_with_registration
+from .main import run
 from .niftyreg import from_mni
 from .niftyreg import to_mni
 from .normalization import normalize_to_pet
@@ -9,7 +10,8 @@
 
 __all__ = [
     "create_abnormality_maps",
-    "run_full",
+    "run",
+    "run_with_registration",
     "from_mni",
     "to_mni",
     "normalize_to_pet",