careamics 0.1.0rc2__py3-none-any.whl → 0.1.0rc4__py3-none-any.whl

careamics/lightning_prediction_loop.py

@@ -0,0 +1,116 @@
+from typing import Optional
+
+import pytorch_lightning as L
+from pytorch_lightning.loops.fetchers import _DataLoaderIterDataFetcher
+from pytorch_lightning.loops.utilities import _no_grad_context
+from pytorch_lightning.trainer import call
+from pytorch_lightning.utilities.types import _PREDICT_OUTPUT
+
+from careamics.prediction import stitch_prediction
+
+
+class CAREamicsPredictionLoop(L.loops._PredictionLoop):
+    """
+    CAREamics prediction loop.
+
+    This class extends the PyTorch Lightning `_PredictionLoop` class to include
+    the stitching of the tiles into a single prediction result.
+    """
+
+    def _on_predict_epoch_end(self) -> Optional[_PREDICT_OUTPUT]:
+        """
+        Calls `on_predict_epoch_end` hook.
+
+        Adapted from the parent method.
+
+        Returns
+        -------
+            the results for all dataloaders
+        """
+        trainer = self.trainer
+        call._call_callback_hooks(trainer, "on_predict_epoch_end")
+        call._call_lightning_module_hook(trainer, "on_predict_epoch_end")
+
+        if self.return_predictions:
+            ########################################################
+            ################ CAREamics specific code ###############
+            if len(self.predicted_array) == 1:
+                # TODO does this make sense to here? (force numpy array)
+                return self.predicted_array[0].numpy()
+            else:
+                # TODO revisit logic
+                return [element.numpy() for element in self.predicted_array]
+            ########################################################
+        return None
+
+    @_no_grad_context
+    def run(self) -> Optional[_PREDICT_OUTPUT]:
+        """
+        Runs the prediction loop.
+
+        Adapted from the parent method in order to stitch the predictions.
+
+        Returns
+        -------
+        Optional[_PREDICT_OUTPUT]
+            Prediction output
+        """
+        self.setup_data()
+        if self.skip:
+            return None
+        self.reset()
+        self.on_run_start()
+        data_fetcher = self._data_fetcher
+        assert data_fetcher is not None
+
+        self.predicted_array = []
+        self.tiles = []
+        self.stitching_data = []
+
+        while True:
+            try:
+                if isinstance(data_fetcher, _DataLoaderIterDataFetcher):
+                    dataloader_iter = next(data_fetcher)
+                    # hook's batch_idx and dataloader_idx arguments correctness cannot
+                    # be guaranteed in this setting
+                    batch = data_fetcher._batch
+                    batch_idx = data_fetcher._batch_idx
+                    dataloader_idx = data_fetcher._dataloader_idx
+                else:
+                    dataloader_iter = None
+                    batch, batch_idx, dataloader_idx = next(data_fetcher)
+                self.batch_progress.is_last_batch = data_fetcher.done
+
+                # run step hooks
+                self._predict_step(batch, batch_idx, dataloader_idx, dataloader_iter)
+
+                ########################################################
+                ################ CAREamics specific code ###############
+                is_tiled = len(self.predictions[batch_idx]) == 2
+                if is_tiled:
+                    # extract the last tile flag and the coordinates (crop and stitch)
+                    last_tile, *stitch_data = self.predictions[batch_idx][1]
+
+                    # append the tile and the coordinates to the lists
+                    self.tiles.append(self.predictions[batch_idx][0])
+                    self.stitching_data.append(stitch_data)
+
+                    # if last tile, stitch the tiles and add array to the prediction
+                    if any(last_tile):
+                        predicted_batches = stitch_prediction(
+                            self.tiles, self.stitching_data
+                        )
+                        self.predicted_array.append(predicted_batches)
+                        self.tiles.clear()
+                        self.stitching_data.clear()
+                else:
+                    # simply add the prediction to the list
+                    self.predicted_array.append(self.predictions[batch_idx])
+                ########################################################
+            except StopIteration:
+                break
+            finally:
+                self._restarting = False
+        return self.on_run_end()
+
+    # TODO predictions aren't stacked, list returned

careamics/losses/__init__.py

@@ -1,4 +1,7 @@
 """Losses module."""
 
 
-from .loss_factory import create_loss_function as create_loss_function
+from .loss_factory import loss_factory
+
+# from .noise_model_factory import noise_model_factory as noise_model_factory
+# from .noise_models import GaussianMixtureNoiseModel, HistogramNoiseModel

careamics/losses/loss_factory.py

@@ -3,22 +3,21 @@ Loss factory module.
 
 This module contains a factory function for creating loss functions.
 """
-from typing import Callable
+from typing import Callable, Union
 
-from careamics.config import Configuration
-from careamics.config.algorithm import Loss
+from ..config.support import SupportedLoss
+from .losses import mae_loss, mse_loss, n2v_loss
 
-from .losses import n2v_loss
 
-
-def create_loss_function(config: Configuration) -> Callable:
-    """
-    Create loss function based on Configuration.
+# TODO add tests
+# TODO add custom?
+def loss_factory(loss: Union[SupportedLoss, str]) -> Callable:
+    """Return loss function.
 
     Parameters
     ----------
-    config : Configuration
-        Configuration.
+    loss: SupportedLoss
+        Requested loss.
 
     Returns
     -------
@@ -30,9 +29,20 @@ def create_loss_function(config: Configuration) -> Callable:
     NotImplementedError
         If the loss is unknown.
     """
-    loss_type = config.algorithm.loss
-
-    if loss_type == Loss.N2V:
+    if loss == SupportedLoss.N2V:
         return n2v_loss
+
+    # elif loss_type == SupportedLoss.PN2V:
+    #     return pn2v_loss
+
+    elif loss == SupportedLoss.MAE:
+        return mae_loss
+
+    elif loss == SupportedLoss.MSE:
+        return mse_loss
+
+    # elif loss_type == SupportedLoss.DICE:
+    #     return dice_loss
+
     else:
-        raise NotImplementedError(f"Loss {loss_type} is not yet supported.")
+        raise NotImplementedError(f"Loss {loss} is not yet supported.")

careamics/losses/losses.py

@@ -3,14 +3,34 @@ Loss submodule.
 
 This submodule contains the various losses used in CAREamics.
 """
+
 import torch
 
+# TODO if we are only using the DiceLoss, can we just implement it?
+# from segmentation_models_pytorch.losses import DiceLoss
+from torch.nn import L1Loss, MSELoss
+
+
+def mse_loss(samples: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+    """
+    Mean squared error loss.
+
+    Returns
+    -------
+    torch.Tensor
+        Loss value.
+    """
+    loss = MSELoss()
+    return loss(samples, labels)
+
 
 def n2v_loss(
-    samples: torch.Tensor, labels: torch.Tensor, masks: torch.Tensor, device: str
+    manipulated_patches: torch.Tensor,
+    original_patches: torch.Tensor,
+    masks: torch.Tensor,
 ) -> torch.Tensor:
     """
-    N2V Loss function (see Eq.7 in Krull et al).
+    N2V Loss function described in A Krull et al 2018.
 
     Parameters
     ----------
@@ -20,15 +40,55 @@ def n2v_loss(
         Noisy patches.
     masks : torch.Tensor
         Array containing masked pixel locations.
-    device : str
-        Device to use.
 
     Returns
     -------
     torch.Tensor
         Loss value.
     """
-    errors = (labels - samples) ** 2
+    errors = (original_patches - manipulated_patches) ** 2
     # Average over pixels and batch
     loss = torch.sum(errors * masks) / torch.sum(masks)
     return loss
+
+
+def mae_loss(samples: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+    """
+    N2N Loss function described in to J Lehtinen et al 2018.
+
+    Parameters
+    ----------
+    samples : torch.Tensor
+        Raw patches.
+    labels : torch.Tensor
+        Different subset of noisy patches.
+
+    Returns
+    -------
+    torch.Tensor
+        Loss value.
+    """
+    loss = L1Loss()
+    return loss(samples, labels)
+
+
+# def pn2v_loss(
+#     samples: torch.Tensor,
+#     labels: torch.Tensor,
+#     masks: torch.Tensor,
+#     noise_model: HistogramNoiseModel,
+# ) -> torch.Tensor:
+#     """Probabilistic N2V loss function described in A Krull et al., CVF (2019)."""
+#     likelihoods = noise_model.likelihood(labels, samples)
+#     likelihoods_avg = torch.log(torch.mean(likelihoods, dim=0, keepdim=True)[0, ...])
+
+#     # Average over pixels and batch
+#     loss = -torch.sum(likelihoods_avg * masks) / torch.sum(masks)
+#     return loss
+
+
+# def dice_loss(
+#     samples: torch.Tensor, labels: torch.Tensor, mode: str = "multiclass"
+# ) -> torch.Tensor:
+#     """Dice loss function."""
+#     return DiceLoss(mode=mode)(samples, labels.long())

careamics/losses/noise_model_factory.py

@@ -0,0 +1,40 @@
+from typing import Type, Union
+
+from ..config.noise_models import NoiseModel, NoiseModelType
+from .noise_models import GaussianMixtureNoiseModel, HistogramNoiseModel
+
+
+def noise_model_factory(
+    noise_config: NoiseModel,
+) -> Type[Union[HistogramNoiseModel, GaussianMixtureNoiseModel, None]]:
+    """Create loss model based on Configuration.
+
+    Parameters
+    ----------
+    config : Configuration
+        Configuration.
+
+    Returns
+    -------
+    Noise model
+
+    Raises
+    ------
+    NotImplementedError
+        If the noise model is unknown.
+    """
+    noise_model_type = noise_config.model_type if noise_config else None
+
+    if noise_model_type == NoiseModelType.HIST:
+        return HistogramNoiseModel
+
+    elif noise_model_type == NoiseModelType.GMM:
+        return GaussianMixtureNoiseModel
+
+    elif noise_model_type is None:
+        return None
+
+    else:
+        raise NotImplementedError(
+            f"Noise model {noise_model_type} is not yet supported."
+        )