Implement inferencing

Author: jonaskordt

.gitignore

@@ -153,3 +153,6 @@ lightning_logs
 
 # Weight&Biases Logs
 wandb/
+
+# Checkpoints
+checkpoints

README.md

@@ -88,7 +88,9 @@ Additionally, the following optional arguments can be supplied in the same confi
     "epochs": 50,
     "gpus": 1,
     "loss": "dice",
-    "optimizer": "adam"
+    "optimizer": "adam",
+    "prediction_count": None,
+    "prediction_dir": "./predictions"
 }
 ```
 

brats_example_config.json

@@ -8,5 +8,7 @@
   "optimizer": "adam",
   "loss": "dice",
   "data_dir": "/dhc/groups/mpws2021cl1/Data",
-  "gpus": 1
-}
+  "gpus": 1,
+  "prediction_count": 5,
+  "prediction_dir": "/dhc/groups/mpws2021cl1/Predictions"
+}

src/datasets/__init__.py

@@ -2,3 +2,4 @@
 from datasets.data_module import ActiveLearningDataModule
 from datasets.brats_data_module import BraTSDataModule
 from datasets.pascal_voc_data_module import PascalVOCDataModule
+from datasets.brats_dataset import BraTSDataset

src/datasets/brats_data_module.py

@@ -1,5 +1,6 @@
 """ Module containing the data module for brats data """
 import os
+import random
 from typing import Any, List, Optional, Union, Tuple
 from torch.utils.data import Dataset
 
@@ -21,7 +22,9 @@ class BraTSDataModule(ActiveLearningDataModule):
     # pylint: disable=unused-argument,no-self-use
     @staticmethod
     def discover_paths(
-        dir_path: str, modality: str = "flair"
+        dir_path: str,
+        modality: str = "flair",
+        random_samples: Union[int, None] = None,
     ) -> Tuple[List[str], List[str]]:
         """
         Discover the .nii.gz file paths with a given modality
@@ -40,6 +43,9 @@ def discover_paths(
             if not case.startswith(".") and os.path.isdir(os.path.join(dir_path, case))
         ]
 
+        if random_samples is not None and random_samples < len(cases):
+            cases = random.sample(cases, random_samples)
+
         image_paths = [
             os.path.join(dir_path, case, f"{os.path.basename(case)}_{modality}.nii.gz")
             for case in cases

src/datasets/brats_dataset.py

@@ -1,5 +1,5 @@
 """ Module to load and batch brats dataset """
-from typing import Any, Callable, List, Optional, Tuple
+from typing import Any, Callable, List, Literal, Optional, Tuple
 import math
 import nibabel as nib
 import numpy as np
@@ -73,6 +73,7 @@ def __init__(
         clip_mask: bool = True,
         transform: Optional[Callable[[Any], torch.Tensor]] = None,
         target_transform: Optional[Callable[[Any], torch.Tensor]] = None,
+        dimensionality: Literal["2d", "3d"] = "2d",
     ):
 
         self.image_paths = image_paths
@@ -98,16 +99,22 @@ def __init__(
         self.transform = transform
         self.target_transform = target_transform
 
+        self.dimensionality = dimensionality
+
     def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]:
-        image_index = math.floor(index / BraTSDataset.IMAGE_DIMENSIONS[0])
-        slice_index = index - image_index * BraTSDataset.IMAGE_DIMENSIONS[0]
-        if image_index != self._current_image_index:
-            self._current_image_index = image_index
-            self._current_image = self.images[self._current_image_index]
-            self._current_mask = self.masks[self._current_image_index]
-
-        x = torch.from_numpy(self._current_image[slice_index, :, :])
-        y = torch.from_numpy(self._current_mask[slice_index, :, :])
+        if self.dimensionality == "2d":
+            image_index = math.floor(index / BraTSDataset.IMAGE_DIMENSIONS[0])
+            slice_index = index - image_index * BraTSDataset.IMAGE_DIMENSIONS[0]
+            if image_index != self._current_image_index:
+                self._current_image_index = image_index
+                self._current_image = self.images[self._current_image_index]
+                self._current_mask = self.masks[self._current_image_index]
+
+            x = torch.from_numpy(self._current_image[slice_index, :, :])
+            y = torch.from_numpy(self._current_mask[slice_index, :, :])
+        else:
+            x = torch.from_numpy(self.images[index])
+            y = torch.from_numpy(self.masks[index])
 
         if self.transform:
             x = self.transform(x)

src/inferencing.py

@@ -0,0 +1,73 @@
+""" Module containing inferencing logic """
+import os
+import uuid
+import torch
+import numpy as np
+import nibabel as nib
+from models import PytorchModel
+from datasets import BraTSDataModule, BraTSDataset
+
+
+class Inferencer:
+    """
+    The inferencer to use a given model for inferencing.
+    Args:
+        model: A model object to be used for inferencing.
+        dataset: Name of the dataset. E.g. 'brats'
+        data_dir: Main directory with the dataset. E.g. './data'
+        prediction_dir: Main directory with the predictions. E.g. './predictions'
+        prediction_count: The amount of predictions to be generated.
+    """
+
+    def __init__(
+        self,
+        model: PytorchModel,
+        dataset: str,
+        data_dir: str,
+        prediction_dir: str,
+        prediction_count: int,
+    ) -> None:
+        self.model = model
+        self.dataset = dataset
+        self.data_dir = data_dir
+        self.prediction_dir = prediction_dir
+        self.prediction_count = prediction_count
+
+    def inference(self) -> None:
+        """Run the inferencing."""
+        if not self.dataset == "brats":
+            print(f"Inferencing is not implemented for the {self.dataset} dataset.")
+            return
+
+        output_folder_name = f"model-{str(uuid.uuid4())}"
+        output_dir = os.path.join(self.prediction_dir, output_folder_name)
+        os.mkdir(output_dir)
+
+        image_paths, annotation_paths = BraTSDataModule.discover_paths(
+            self.data_dir,
+            random_samples=self.prediction_count,
+        )
+        data = BraTSDataset(
+            image_paths=image_paths,
+            annotation_paths=annotation_paths,
+            dimensionality="3d",
+        )
+
+        for i in range(self.prediction_count):
+            x, _ = data.__getitem__(i)
+
+            x = torch.swapaxes(x, 0, 1)
+            pred = self.model(x)
+            pred = torch.swapaxes(pred, 0, 1)
+
+            seg = pred.detach().cpu().numpy()[0]
+            seg = (seg >= 0.5) * 255
+            seg = np.moveaxis(seg, 0, 2)
+            seg = np.rot90(seg, 2, (0, 1))
+            seg = seg.astype("float64")
+
+            img = nib.Nifti1Image(seg, np.eye(4))
+            file_name = os.path.basename(annotation_paths[i]).replace("seg", "pred")
+            path = os.path.join(output_dir, file_name)
+            nib.save(img, path)
+            print(path)

src/main.py

@@ -1,8 +1,10 @@
 """ Main module to execute active learning pipeline from CLI """
 import json
 import os.path
+from typing import Union
 import fire
 from active_learning import ActiveLearningPipeline
+from inferencing import Inferencer
 from models import PytorchFCNResnet50, PytorchUNet
 from datasets import BraTSDataModule, PascalVOCDataModule
 from query_strategies import QueryStrategy
@@ -20,6 +22,8 @@ def run_active_learning_pipeline(
     gpus: int = 1,
     loss: str = "dice",
     optimizer: str = "adam",
+    prediction_count: Union[int, None] = None,
+    prediction_dir: str = "./predictions",
 ) -> None:
     """
     Main function to execute an active learning pipeline run, or start an active learning simulation.
@@ -60,6 +64,14 @@ def run_active_learning_pipeline(
     pipeline = ActiveLearningPipeline(data_module, model, strategy, epochs, gpus)
     pipeline.run()
 
+    if prediction_count is None:
+        return
+
+    inferencer = Inferencer(
+        model, dataset, os.path.join(data_dir, "val"), prediction_dir, prediction_count
+    )
+    inferencer.inference()
+
 
 def run_active_learning_pipeline_from_config(config_file_name: str) -> None:
     """