careamics 0.0.14__py3-none-any.whl → 0.0.16__py3-none-any.whl

careamics/lightning/dataset_ng/data_module.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Optional, Union, overload
+from typing import Any, Union, overload
 
 import numpy as np
 import pytorch_lightning as L
@@ -124,14 +124,14 @@ class CareamicsDataModule(L.LightningDataModule):
         self,
         data_config: NGDataConfig,
         *,
-        train_data: Optional[InputType] = None,
-        train_data_target: Optional[InputType] = None,
-        val_data: Optional[InputType] = None,
-        val_data_target: Optional[InputType] = None,
-        pred_data: Optional[InputType] = None,
-        pred_data_target: Optional[InputType] = None,
+        train_data: InputType | None = None,
+        train_data_target: InputType | None = None,
+        val_data: InputType | None = None,
+        val_data_target: InputType | None = None,
+        pred_data: InputType | None = None,
+        pred_data_target: InputType | None = None,
         extension_filter: str = "",
-        val_percentage: Optional[float] = None,
+        val_percentage: float | None = None,
         val_minimum_split: int = 5,
         use_in_memory: bool = True,
     ) -> None: ...
@@ -142,16 +142,16 @@ class CareamicsDataModule(L.LightningDataModule):
         self,
         data_config: NGDataConfig,
         *,
-        train_data: Optional[InputType] = None,
-        train_data_target: Optional[InputType] = None,
-        val_data: Optional[InputType] = None,
-        val_data_target: Optional[InputType] = None,
-        pred_data: Optional[InputType] = None,
-        pred_data_target: Optional[InputType] = None,
+        train_data: InputType | None = None,
+        train_data_target: InputType | None = None,
+        val_data: InputType | None = None,
+        val_data_target: InputType | None = None,
+        pred_data: InputType | None = None,
+        pred_data_target: InputType | None = None,
         read_source_func: Callable,
-        read_kwargs: Optional[dict[str, Any]] = None,
+        read_kwargs: dict[str, Any] | None = None,
         extension_filter: str = "",
-        val_percentage: Optional[float] = None,
+        val_percentage: float | None = None,
         val_minimum_split: int = 5,
         use_in_memory: bool = True,
     ) -> None: ...
@@ -161,16 +161,16 @@ class CareamicsDataModule(L.LightningDataModule):
         self,
         data_config: NGDataConfig,
         *,
-        train_data: Optional[Any] = None,
-        train_data_target: Optional[Any] = None,
-        val_data: Optional[Any] = None,
-        val_data_target: Optional[Any] = None,
-        pred_data: Optional[Any] = None,
-        pred_data_target: Optional[Any] = None,
+        train_data: Any | None = None,
+        train_data_target: Any | None = None,
+        val_data: Any | None = None,
+        val_data_target: Any | None = None,
+        pred_data: Any | None = None,
+        pred_data_target: Any | None = None,
         image_stack_loader: ImageStackLoader,
-        image_stack_loader_kwargs: Optional[dict[str, Any]] = None,
+        image_stack_loader_kwargs: dict[str, Any] | None = None,
         extension_filter: str = "",
-        val_percentage: Optional[float] = None,
+        val_percentage: float | None = None,
         val_minimum_split: int = 5,
         use_in_memory: bool = True,
     ) -> None: ...
@@ -179,18 +179,18 @@ class CareamicsDataModule(L.LightningDataModule):
         self,
         data_config: NGDataConfig,
         *,
-        train_data: Optional[Any] = None,
-        train_data_target: Optional[Any] = None,
-        val_data: Optional[Any] = None,
-        val_data_target: Optional[Any] = None,
-        pred_data: Optional[Any] = None,
-        pred_data_target: Optional[Any] = None,
-        read_source_func: Optional[Callable] = None,
-        read_kwargs: Optional[dict[str, Any]] = None,
-        image_stack_loader: Optional[ImageStackLoader] = None,
-        image_stack_loader_kwargs: Optional[dict[str, Any]] = None,
+        train_data: Any | None = None,
+        train_data_target: Any | None = None,
+        val_data: Any | None = None,
+        val_data_target: Any | None = None,
+        pred_data: Any | None = None,
+        pred_data_target: Any | None = None,
+        read_source_func: Callable | None = None,
+        read_kwargs: dict[str, Any] | None = None,
+        image_stack_loader: ImageStackLoader | None = None,
+        image_stack_loader_kwargs: dict[str, Any] | None = None,
         extension_filter: str = "",
-        val_percentage: Optional[float] = None,
+        val_percentage: float | None = None,
         val_minimum_split: int = 5,
         use_in_memory: bool = True,
     ) -> None:
@@ -280,7 +280,7 @@ class CareamicsDataModule(L.LightningDataModule):
     def _validate_input_target_type_consistency(
         self,
         input_data: InputType,
-        target_data: Optional[InputType],
+        target_data: InputType | None,
     ) -> None:
         """Validate if the input and target data types are consistent.
 
@@ -314,7 +314,7 @@ class CareamicsDataModule(L.LightningDataModule):
         self,
         input_data,
         target_data=None,
-    ) -> tuple[list[Path], Optional[list[Path]]]:
+    ) -> tuple[list[Path], list[Path] | None]:
         """List files from input and target directories.
 
         Parameters
@@ -347,7 +347,7 @@ class CareamicsDataModule(L.LightningDataModule):
         self,
         input_data,
         target_data=None,
-    ) -> tuple[list[Path], Optional[list[Path]]]:
+    ) -> tuple[list[Path], list[Path] | None]:
         """Create a list of file paths from the input and target data.
 
         Parameters
@@ -379,7 +379,7 @@ class CareamicsDataModule(L.LightningDataModule):
     def _validate_array_input(
         self,
         input_data: InputType,
-        target_data: Optional[InputType],
+        target_data: InputType | None,
     ) -> tuple[Any, Any]:
         """Validate if the input data is a numpy array.
 
@@ -408,8 +408,8 @@ class CareamicsDataModule(L.LightningDataModule):
             )
 
     def _validate_path_input(
-        self, input_data: InputType, target_data: Optional[InputType]
-    ) -> tuple[list[Path], Optional[list[Path]]]:
+        self, input_data: InputType, target_data: InputType | None
+    ) -> tuple[list[Path], list[Path] | None]:
         """Validate if the input data is a path or a list of paths.
 
         Parameters
@@ -488,8 +488,8 @@ class CareamicsDataModule(L.LightningDataModule):
 
     def _initialize_data_pair(
         self,
-        input_data: Optional[InputType],
-        target_data: Optional[InputType],
+        input_data: InputType | None,
+        target_data: InputType | None,
     ) -> tuple[Any, Any]:
         """
         Initialize a pair of input and target data.

careamics/lightning/lightning_module.py

@@ -1,11 +1,11 @@
 """CAREamics Lightning module."""
 
 from collections.abc import Callable
-from typing import Any, Literal, Optional, Union
+from typing import Any, Literal, Union
 
 import numpy as np
 import pytorch_lightning as L
-from torch import Tensor, nn
+import torch
 
 from careamics.config import (
     N2VAlgorithm,
@@ -71,7 +71,9 @@ class FCNModule(L.LightningModule):
         Learning rate scheduler name.
     """
 
-    def __init__(self, algorithm_config: Union[UNetBasedAlgorithm, dict]) -> None:
+    def __init__(
+        self, algorithm_config: Union[UNetBasedAlgorithm, VAEBasedAlgorithm, dict]
+    ) -> None:
         """Lightning module for CAREamics.
 
         This class encapsulates the a PyTorch model along with the training, validation,
@@ -90,7 +92,7 @@ class FCNModule(L.LightningModule):
         # create preprocessing, model and loss function
         if isinstance(algorithm_config, N2VAlgorithm):
             self.use_n2v = True
-            self.n2v_preprocess: Optional[N2VManipulateTorch] = N2VManipulateTorch(
+            self.n2v_preprocess: N2VManipulateTorch | None = N2VManipulateTorch(
                 n2v_manipulate_config=algorithm_config.n2v_config
             )
         else:
@@ -98,7 +100,7 @@ class FCNModule(L.LightningModule):
             self.n2v_preprocess = None
 
         self.algorithm = algorithm_config.algorithm
-        self.model: nn.Module = model_factory(algorithm_config.model)
+        self.model: torch.nn.Module = model_factory(algorithm_config.model)
         self.loss_func = loss_factory(algorithm_config.loss)
 
         # save optimizer and lr_scheduler names and parameters
@@ -122,12 +124,12 @@ class FCNModule(L.LightningModule):
         """
         return self.model(x)
 
-    def training_step(self, batch: Tensor, batch_idx: Any) -> Any:
+    def training_step(self, batch: torch.Tensor, batch_idx: Any) -> Any:
         """Training step.
 
         Parameters
         ----------
-        batch : torch.Tensor
+        batch : torch.torch.Tensor
             Input batch.
         batch_idx : Any
             Batch index.
@@ -154,12 +156,12 @@ class FCNModule(L.LightningModule):
         self.log("learning_rate", current_lr, on_step=False, on_epoch=True, logger=True)
         return loss
 
-    def validation_step(self, batch: Tensor, batch_idx: Any) -> None:
+    def validation_step(self, batch: torch.Tensor, batch_idx: Any) -> None:
         """Validation step.
 
         Parameters
         ----------
-        batch : torch.Tensor
+        batch : torch.torch.Tensor
             Input batch.
         batch_idx : Any
             Batch index.
@@ -184,12 +186,12 @@ class FCNModule(L.LightningModule):
             logger=True,
         )
 
-    def predict_step(self, batch: Tensor, batch_idx: Any) -> Any:
+    def predict_step(self, batch: torch.Tensor, batch_idx: Any) -> Any:
         """Prediction step.
 
         Parameters
         ----------
-        batch : torch.Tensor
+        batch : torch.torch.torch.Tensor
             Input batch.
         batch_idx : Any
             Batch index.
@@ -330,21 +332,23 @@ class VAEModule(L.LightningModule):
         # self.save_hyperparameters(self.algorithm_config.model_dump())
 
         # create model
-        self.model: nn.Module = model_factory(self.algorithm_config.model)
+        self.model: torch.nn.Module = model_factory(self.algorithm_config.model)
 
+        # supervised_mode
+        self.supervised_mode = self.algorithm_config.is_supervised
         # create loss function
-        self.noise_model: Optional[NoiseModel] = noise_model_factory(
+        self.noise_model: NoiseModel | None = noise_model_factory(
             self.algorithm_config.noise_model
         )
 
-        self.noise_model_likelihood: Optional[NoiseModelLikelihood] = (
-            likelihood_factory(
+        self.noise_model_likelihood: NoiseModelLikelihood | None = None
+        if self.algorithm_config.noise_model_likelihood is not None:
+            self.noise_model_likelihood = likelihood_factory(
                 config=self.algorithm_config.noise_model_likelihood,
                 noise_model=self.noise_model,
             )
-        )
 
-        self.gaussian_likelihood: Optional[GaussianLikelihood] = likelihood_factory(
+        self.gaussian_likelihood: GaussianLikelihood | None = likelihood_factory(
             self.algorithm_config.gaussian_likelihood
         )
 
@@ -362,30 +366,43 @@ class VAEModule(L.LightningModule):
             RunningPSNR() for _ in range(self.algorithm_config.model.output_channels)
         ]
 
-    def forward(self, x: Tensor) -> tuple[Tensor, dict[str, Any]]:
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, dict[str, Any]]:
         """Forward pass.
 
         Parameters
         ----------
-        x : Tensor
+        x : torch.Tensor
             Input tensor of shape (B, (1 + n_LC), [Z], Y, X), where n_LC is the
             number of lateral inputs.
 
         Returns
         -------
-        tuple[Tensor, dict[str, Any]]
+        tuple[torch.Tensor, dict[str, Any]]
             A tuple with the output tensor and additional data from the top-down pass.
         """
         return self.model(x)  # TODO Different model can have more than one output
 
+    def set_data_stats(self, data_mean, data_std):
+        """Set data mean and std for the noise model likelihood.
+
+        Parameters
+        ----------
+        data_mean : float
+            Mean of the data.
+        data_std : float
+            Standard deviation of the data.
+        """
+        if self.noise_model_likelihood is not None:
+            self.noise_model_likelihood.set_data_stats(data_mean, data_std)
+
     def training_step(
-        self, batch: tuple[Tensor, Tensor], batch_idx: Any
-    ) -> Optional[dict[str, Tensor]]:
+        self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: Any
+    ) -> dict[str, torch.Tensor] | None:
         """Training step.
 
         Parameters
         ----------
-        batch : tuple[Tensor, Tensor]
+        batch : tuple[torch.Tensor, torch.Tensor]
             Input batch. It is a tuple with the input tensor and the target tensor.
             The input tensor has shape (B, (1 + n_LC), [Z], Y, X), where n_LC is the
             number of lateral inputs. The target tensor has shape (B, C, [Z], Y, X),
@@ -399,15 +416,29 @@ class VAEModule(L.LightningModule):
         Any
             Loss value.
         """
-        x, target = batch
+        x, *target = batch
 
         # Forward pass
         out = self.model(x)
+        if not self.supervised_mode:
+            target = x
+        else:
+            target = target[
+                0
+            ]  # hacky way to unpack. #TODO maybe should be fixed on the dataset level
 
         # Update loss parameters
         self.loss_parameters.kl_params.current_epoch = self.current_epoch
 
         # Compute loss
+        if self.noise_model_likelihood is not None:
+            if (
+                self.noise_model_likelihood.data_mean is None
+                or self.noise_model_likelihood.data_std is None
+            ):
+                raise RuntimeError(
+                    "NoiseModelLikelihood: data_mean and data_std must be set before training."
+                )
         loss = self.loss_func(
             model_outputs=out,
             targets=target,
@@ -419,15 +450,26 @@ class VAEModule(L.LightningModule):
         # Logging
         # TODO: implement a separate logging method?
         self.log_dict(loss, on_step=True, on_epoch=True)
-        # self.log("lr", self, on_epoch=True)
+
+        try:
+            optimizer = self.optimizers()
+            current_lr = optimizer.param_groups[0]["lr"]
+            self.log(
+                "learning_rate", current_lr, on_step=False, on_epoch=True, logger=True
+            )
+        except RuntimeError:
+            # This happens when the module is not attached to a trainer, e.g., in tests
+            pass
         return loss
 
-    def validation_step(self, batch: tuple[Tensor, Tensor], batch_idx: Any) -> None:
+    def validation_step(
+        self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: Any
+    ) -> None:
         """Validation step.
 
         Parameters
         ----------
-        batch : tuple[Tensor, Tensor]
+        batch : tuple[torch.Tensor, torch.Tensor]
             Input batch. It is a tuple with the input tensor and the target tensor.
             The input tensor has shape (B, (1 + n_LC), [Z], Y, X), where n_LC is the
             number of lateral inputs. The target tensor has shape (B, C, [Z], Y, X),
@@ -436,11 +478,16 @@ class VAEModule(L.LightningModule):
         batch_idx : Any
             Batch index.
         """
-        x, target = batch
+        x, *target = batch
 
         # Forward pass
         out = self.model(x)
-
+        if not self.supervised_mode:
+            target = x
+        else:
+            target = target[
+                0
+            ]  # hacky way to unpack. #TODO maybe should be fixed on the datasel level
         # Compute loss
         loss = self.loss_func(
             model_outputs=out,
@@ -466,12 +513,12 @@ class VAEModule(L.LightningModule):
         else:
             self.log("val_psnr", 0.0, on_epoch=True, prog_bar=True)
 
-    def predict_step(self, batch: Tensor, batch_idx: Any) -> Any:
+    def predict_step(self, batch: torch.Tensor, batch_idx: Any) -> Any:
         """Prediction step.
 
         Parameters
         ----------
-        batch : Tensor
+        batch : torch.Tensor
             Input batch.
         batch_idx : Any
             Batch index.
@@ -481,36 +528,86 @@ class VAEModule(L.LightningModule):
         Any
             Model output.
         """
-        if self._trainer.datamodule.tiled:
+        if self.algorithm_config.algorithm == "microsplit":
             x, *aux = batch
-        else:
-            x = batch
-            aux = []
+            # Reset model for inference with spatial dimensions only (H, W)
+            self.model.reset_for_inference(x.shape[-2:])
 
-        # apply test-time augmentation if available
-        # TODO: probably wont work with batch size > 1
-        if self._trainer.datamodule.prediction_config.tta_transforms:
-            tta = ImageRestorationTTA()
-            augmented_batch = tta.forward(x)  # list of augmented tensors
-            augmented_output = []
-            for augmented in augmented_batch:
-                augmented_pred = self.model(augmented)
-                augmented_output.append(augmented_pred)
-            output = tta.backward(augmented_output)
-        else:
-            output = self.model(x)
+            rec_img_list = []
+            for _ in range(self.algorithm_config.mmse_count):
+                # get model output
+                rec, _ = self.model(x)
 
-        # Denormalize the output
-        denorm = Denormalize(
-            image_means=self._trainer.datamodule.predict_dataset.image_means,
-            image_stds=self._trainer.datamodule.predict_dataset.image_stds,
-        )
-        denormalized_output = denorm(patch=output.cpu().numpy())
+                # get reconstructed img
+                if self.model.predict_logvar is None:
+                    rec_img = rec
+                    logvar = torch.tensor([-1])
+                else:
+                    rec_img, logvar = torch.chunk(rec, chunks=2, dim=1)
+                rec_img_list.append(rec_img.cpu().unsqueeze(0))  # add MMSE dim
+
+            # aggregate results
+            samples = torch.cat(rec_img_list, dim=0)
+            mmse_imgs = torch.mean(samples, dim=0)  # avg over MMSE dim
+            std_imgs = torch.std(samples, dim=0)  # std over MMSE dim
+
+            tile_prediction = mmse_imgs.cpu().numpy()
+            tile_std = std_imgs.cpu().numpy()
+
+            return tile_prediction, tile_std
 
-        if len(aux) > 0:  # aux can be tiling information
-            return denormalized_output, *aux
         else:
-            return denormalized_output
+            # Regular prediction logic
+            if self._trainer.datamodule.tiled:
+                # TODO tile_size should match model input size
+                x, *aux = batch
+                x = (
+                    x[0] if isinstance(x, list | tuple) else x
+                )  # TODO ugly, so far i don't know why x might be a list
+                self.model.reset_for_inference(x.shape)  # TODO should it be here ?
+            else:
+                x = batch[0] if isinstance(batch, list | tuple) else batch
+                aux = []
+                self.model.reset_for_inference(x.shape)
+
+            mmse_list = []
+            for _ in range(self.algorithm_config.mmse_count):
+                # apply test-time augmentation if available
+                if self._trainer.datamodule.prediction_config.tta_transforms:
+                    tta = ImageRestorationTTA()
+                    augmented_batch = tta.forward(x)  # list of augmented tensors
+                    augmented_output = []
+                    for augmented in augmented_batch:
+                        augmented_pred = self.model(augmented)
+                        augmented_output.append(augmented_pred)
+                    output = tta.backward(augmented_output)
+                else:
+                    output = self.model(x)
+
+                # taking the 1st element of the output, 2nd is std if
+                # predict_logvar=="pixelwise"
+                output = (
+                    output[0]
+                    if self.model.predict_logvar is None
+                    else output[0][:, 0:1, ...]
+                )
+                mmse_list.append(output)
+
+            mmse = torch.stack(mmse_list).mean(0)
+            std = torch.stack(mmse_list).std(0)  # TODO why?
+            # TODO better way to unpack if pred logvar
+            # Denormalize the output
+            denorm = Denormalize(
+                image_means=self._trainer.datamodule.predict_dataset.image_means,
+                image_stds=self._trainer.datamodule.predict_dataset.image_stds,
+            )
+
+            denormalized_output = denorm(patch=mmse.cpu().numpy())
+
+            if len(aux) > 0:  # aux can be tiling information
+                return denormalized_output, std, *aux
+            else:
+                return denormalized_output, std
 
     def configure_optimizers(self) -> Any:
         """Configure optimizers and learning rate schedulers.
@@ -539,19 +636,19 @@ class VAEModule(L.LightningModule):
     # should we refactor LadderVAE so that it already outputs
     # tuple(`mean`, `logvar`, `td_data`)?
     def get_reconstructed_tensor(
-        self, model_outputs: tuple[Tensor, dict[str, Any]]
-    ) -> Tensor:
+        self, model_outputs: tuple[torch.Tensor, dict[str, Any]]
+    ) -> torch.Tensor:
         """Get the reconstructed tensor from the LVAE model outputs.
 
         Parameters
         ----------
-        model_outputs : tuple[Tensor, dict[str, Any]]
+        model_outputs : tuple[torch.Tensor, dict[str, Any]]
             Model outputs. It is a tuple with a tensor representing the predicted mean
             and (optionally) logvar, and the top-down data dictionary.
 
         Returns
         -------
-        Tensor
+        torch.Tensor
             Reconstructed tensor, i.e., the predicted mean.
         """
         predictions, _ = model_outputs
@@ -562,18 +659,18 @@ class VAEModule(L.LightningModule):
 
     def compute_val_psnr(
         self,
-        model_output: tuple[Tensor, dict[str, Any]],
-        target: Tensor,
+        model_output: tuple[torch.Tensor, dict[str, Any]],
+        target: torch.Tensor,
         psnr_func: Callable = scale_invariant_psnr,
     ) -> list[float]:
         """Compute the PSNR for the current validation batch.
 
         Parameters
         ----------
-        model_output : tuple[Tensor, dict[str, Any]]
+        model_output : tuple[torch.Tensor, dict[str, Any]]
             Model output, a tuple with the predicted mean and (optionally) logvar,
             and the top-down data dictionary.
-        target : Tensor
+        target : torch.Tensor
             Target tensor.
         psnr_func : Callable, optional
             PSNR function to use, by default `scale_invariant_psnr`.
@@ -583,6 +680,7 @@ class VAEModule(L.LightningModule):
         list[float]
             PSNR for each channel in the current batch.
         """
+        # TODO check this! Related to is_supervised which is also wacky
         out_channels = target.shape[1]
 
         # get the reconstructed image
@@ -603,7 +701,7 @@ class VAEModule(L.LightningModule):
             for i in range(out_channels)
         ]
 
-    def reduce_running_psnr(self) -> Optional[float]:
+    def reduce_running_psnr(self) -> float | None:
         """Reduce the running PSNR statistics and reset the running PSNR.
 
         Returns
@@ -634,11 +732,11 @@ def create_careamics_module(
     use_n2v2: bool = False,
     struct_n2v_axis: Literal["horizontal", "vertical", "none"] = "none",
     struct_n2v_span: int = 5,
-    model_parameters: Optional[dict] = None,
+    model_parameters: dict | None = None,
     optimizer: Union[SupportedOptimizer, str] = "Adam",
-    optimizer_parameters: Optional[dict] = None,
+    optimizer_parameters: dict | None = None,
     lr_scheduler: Union[SupportedScheduler, str] = "ReduceLROnPlateau",
-    lr_scheduler_parameters: Optional[dict] = None,
+    lr_scheduler_parameters: dict | None = None,
 ) -> Union[FCNModule, VAEModule]:
     """Create a CAREamics Lightning module.