careamics 0.0.2__py3-none-any.whl → 0.0.3__py3-none-any.whl

careamics/lvae_training/get_config.py

@@ -6,7 +6,7 @@ import os
 
 import ml_collections
 
-from careamics.lvae_training.data_utils import DataType
+from careamics.lvae_training.dataset.data_utils import DataType
 from careamics.models.lvae.utils import LossType
 
 

careamics/lvae_training/train_lvae.py

@@ -1,5 +1,5 @@
 """
-This script is meant to load data, intialize the model, and provide the logic for training it.
+This script is meant to load data, initialize the model, and provide the logic for training it.
 """
 
 import glob
@@ -20,8 +20,11 @@ from torch.utils.data import DataLoader
 sys.path.append(os.path.join(os.path.dirname(__file__), "../.."))
 print(sys.path)
 
-from careamics.lvae_training.data_modules import LCMultiChDloader, MultiChDloader
-from careamics.lvae_training.data_utils import DataSplitType
+from careamics.lvae_training.dataset.data_modules import (
+    LCMultiChDloader,
+    MultiChDloader,
+)
+from careamics.lvae_training.dataset.data_utils import DataSplitType
 from careamics.lvae_training.lightning_module import LadderVAELight
 from careamics.lvae_training.train_utils import *
 

careamics/model_io/bioimage/bioimage_utils.py

@@ -25,7 +25,7 @@ def get_unzip_path(zip_path: Union[Path, str]) -> Path:
 def create_env_text(pytorch_version: str) -> str:
     """Create environment yaml content for the bioimage model.
 
-    This installs an environemnt with the specified pytorch version and the latest
+    This installs an environment with the specified pytorch version and the latest
     changes to careamics.
 
     Parameters

careamics/model_io/bioimage/model_description.py

@@ -204,7 +204,7 @@ def create_model_description(
     config : Configuration
         CAREamics configuration.
     name : str
-        Name fo the model.
+        Name of the model.
     general_description : str
         General description of the model.
     authors : List[Author]
@@ -252,7 +252,7 @@ def create_model_description(
 
     # weights description
     architecture_descr = ArchitectureFromLibraryDescr(
-        import_from="careamics.models",
+        import_from="careamics.models.unet",
         callable=f"{config.algorithm_config.model.architecture}",
         kwargs=config.algorithm_config.model.model_dump(),
     )

careamics/model_io/bmz_io.py

@@ -12,7 +12,7 @@ from torch import __version__, load, save
 
 from careamics.config import Configuration, load_configuration, save_configuration
 from careamics.config.support import SupportedArchitecture
-from careamics.lightning.lightning_module import CAREamicsModule
+from careamics.lightning.lightning_module import FCNModule, VAEModule
 
 from .bioimage import (
     create_env_text,
@@ -22,7 +22,9 @@ from .bioimage import (
 )
 
 
-def _export_state_dict(model: CAREamicsModule, path: Union[Path, str]) -> Path:
+def _export_state_dict(
+    model: Union[FCNModule, VAEModule], path: Union[Path, str]
+) -> Path:
     """
     Export the model state dictionary to a file.
 
@@ -52,7 +54,9 @@ def _export_state_dict(model: CAREamicsModule, path: Union[Path, str]) -> Path:
     return path
 
 
-def _load_state_dict(model: CAREamicsModule, path: Union[Path, str]) -> None:
+def _load_state_dict(
+    model: Union[FCNModule, VAEModule], path: Union[Path, str]
+) -> None:
     """
     Load a model from a state dictionary.
 
@@ -74,7 +78,7 @@ def _load_state_dict(model: CAREamicsModule, path: Union[Path, str]) -> None:
 
 # TODO break down in subfunctions
 def export_to_bmz(
-    model: CAREamicsModule,
+    model: Union[FCNModule, VAEModule],
     config: Configuration,
     path_to_archive: Union[Path, str],
     model_name: str,
@@ -187,7 +191,9 @@ def export_to_bmz(
         save_bioimageio_package(model_description, output_path=path_to_archive)
 
 
-def load_from_bmz(path: Union[Path, str]) -> Tuple[CAREamicsModule, Configuration]:
+def load_from_bmz(
+    path: Union[Path, str]
+) -> Tuple[Union[FCNModule, VAEModule], Configuration]:
     """Load a model from a BioImage Model Zoo archive.
 
     Parameters
@@ -225,7 +231,14 @@ def load_from_bmz(path: Union[Path, str]) -> Tuple[CAREamicsModule, Configuratio
     config = load_configuration(config_path)
 
     # create careamics lightning module
-    model = CAREamicsModule(algorithm_config=config.algorithm_config)
+    if config.algorithm_config.model.architecture == SupportedArchitecture.UNET:
+        model = FCNModule(algorithm_config=config.algorithm_config)
+    elif config.algorithm_config.model.architecture == SupportedArchitecture.LVAE:
+        model = VAEModule(algorithm_config=config.algorithm_config)
+    else:
+        raise ValueError(
+            f"Unsupported architecture {config.algorithm_config.model.architecture}"
+        )  # TODO ugly ?
 
     # load model state dictionary
     _load_state_dict(model, weights_path)

careamics/model_io/model_io_utils.py

@@ -6,12 +6,14 @@ from typing import Tuple, Union
 import torch
 
 from careamics.config import Configuration
-from careamics.lightning.lightning_module import CAREamicsModule
+from careamics.lightning.lightning_module import FCNModule, VAEModule
 from careamics.model_io.bmz_io import load_from_bmz
 from careamics.utils import check_path_exists
 
 
-def load_pretrained(path: Union[Path, str]) -> Tuple[CAREamicsModule, Configuration]:
+def load_pretrained(
+    path: Union[Path, str]
+) -> Tuple[Union[FCNModule, VAEModule], Configuration]:
     """
     Load a pretrained model from a checkpoint or a BioImage Model Zoo model.
 
@@ -44,7 +46,9 @@ def load_pretrained(path: Union[Path, str]) -> Tuple[CAREamicsModule, Configurat
         )
 
 
-def _load_checkpoint(path: Union[Path, str]) -> Tuple[CAREamicsModule, Configuration]:
+def _load_checkpoint(
+    path: Union[Path, str]
+) -> Tuple[Union[FCNModule, VAEModule], Configuration]:
     """
     Load a model from a checkpoint and return both model and configuration.
 
@@ -78,6 +82,14 @@ def _load_checkpoint(path: Union[Path, str]) -> Tuple[CAREamicsModule, Configura
             f"checkpoint: {checkpoint.keys()}"
         ) from e
 
-    model = CAREamicsModule.load_from_checkpoint(path)
+    if cfg_dict["algorithm_config"]["model"]["architecture"] == "UNet":
+        model = FCNModule.load_from_checkpoint(path)
+    elif cfg_dict["algorithm_config"]["model"]["architecture"] == "LVAE":
+        model = VAEModule.load_from_checkpoint(path)
+    else:
+        raise ValueError(
+            "Invalid model architecture: "
+            f"{cfg_dict['algorithm_config']['model']['architecture']}"
+        )
 
     return model, Configuration(**cfg_dict)

careamics/models/__init__.py

@@ -1,7 +1,5 @@
 """Models package."""
 
-__all__ = ["model_factory", "UNet"]
-
+__all__ = ["model_factory"]
 
 from .model_factory import model_factory
-from .unet import UNet as UNet

careamics/models/activation.py

@@ -23,6 +23,8 @@ def get_activation(activation: Union[SupportedActivation, str]) -> Callable:
     """
     if activation == SupportedActivation.RELU:
         return nn.ReLU()
+    elif activation == SupportedActivation.ELU:
+        return nn.ELU()
     elif activation == SupportedActivation.LEAKYRELU:
         return nn.LeakyReLU()
     elif activation == SupportedActivation.TANH:

careamics/models/lvae/__init__.py

@@ -0,0 +1,3 @@
+__all__ = ["LadderVAE"]
+
+from .lvae import LadderVAE

careamics/models/lvae/layers.py

@@ -117,7 +117,7 @@ class ResidualBlock(nn.Module):
                     bias=conv2d_bias,
                 )
                 modules.append(conv)
-                modules.append(nonlin())
+                modules.append(nonlin)
                 if batchnorm:
                     modules.append(nn.BatchNorm2d(channels))
                 if dropout is not None:
@@ -126,7 +126,7 @@ class ResidualBlock(nn.Module):
             for i in range(2):
                 if batchnorm:
                     modules.append(nn.BatchNorm2d(channels))
-                modules.append(nonlin())
+                modules.append(nonlin)
                 conv = nn.Conv2d(
                     channels,
                     channels,
@@ -142,7 +142,7 @@ class ResidualBlock(nn.Module):
             for i in range(2):
                 if batchnorm:
                     modules.append(nn.BatchNorm2d(channels))
-                modules.append(nonlin())
+                modules.append(nonlin)
                 conv = nn.Conv2d(
                     channels,
                     channels,
@@ -189,7 +189,7 @@ class GateLayer2d(nn.Module):
         assert kernel_size % 2 == 1
         pad = kernel_size // 2
         self.conv = nn.Conv2d(channels, 2 * channels, kernel_size, padding=pad)
-        self.nonlin = nonlin()
+        self.nonlin = nonlin
 
     def forward(self, x):
         x = self.conv(x)
@@ -255,7 +255,7 @@ class ResBlockWithResampling(nn.Module):
         resample: bool, optional
             Whether to perform resampling in the first convolutional layer.
             If `False`, the first convolutional layer just maps the input to a tensor with
-            `inner_channels` channels through 1x1 convolution. Deafult is `False`.
+            `inner_channels` channels through 1x1 convolution. Default is `False`.
         res_block_kernel: Union[int, Iterable[int]], optional
             The kernel size used in the convolutions of the residual block.
             It can be either a single integer or a pair of integers defining the squared kernel.
@@ -837,7 +837,7 @@ class TopDownLayer(nn.Module):
         - In inference mode, parameters of q(z_i|z_i+1) are obtained from the inference path,
         by merging outcomes of bottom-up and top-down passes. The exception is the top layer,
         in which the parameters of q(z_L|x) are set as the output of the topmost bottom-up layer.
-        - On the contrary in prediciton/generative mode, parameters of q(z_i|z_i+1) can be obtained
+        - On the contrary in predicition/generative mode, parameters of q(z_i|z_i+1) can be obtained
         once again by merging bottom-up and top-down outputs (CONDITIONAL GENERATION), or it is
         possible to directly sample from the prior p(z_i|z_i+1) (UNCONDITIONAL GENERATION).
 
@@ -899,7 +899,7 @@ class TopDownLayer(nn.Module):
             The number of downsampling steps that has to be done in this layer (typically 1).
             Default is `False`.
         nonlin: Callable, optional
-            The non-linearity function used in the block (e.g., `nn.ReLU`). Deafault is `None`.
+            The non-linearity function used in the block (e.g., `nn.ReLU`). Default is `None`.
         merge_type: Literal["linear", "residual", "residual_ungated"], optional
             The type of merge done in the layer. It can be chosen between "linear", "residual",
             and "residual_ungated". Check the `MergeLayer` class docstring for more information
@@ -931,7 +931,7 @@ class TopDownLayer(nn.Module):
             Whether to set the top prior as learnable.
             If this is set to `False`, in the top-most layer the prior will be N(0,1).
             Otherwise, we will still have a normal distribution whose parameters will be learnt.
-            Deafult is `False`.
+            Default is `False`.
         top_prior_param_shape: Iterable[int], optional
             The size of the tensor which expresses the mean and the variance
             of the prior for the top most layer. Default is `None`.
@@ -1102,7 +1102,7 @@ class TopDownLayer(nn.Module):
             The tensor defining the parameters /mu_q and /sigma_q computed during the bottom-up deterministic pass
             at the correspondent hierarchical layer.
         var_clip_max: float, optional
-            The maximum value reachable by the log-variance of the latent distribtion.
+            The maximum value reachable by the log-variance of the latent distribution.
             Values exceeding this threshold are clipped. Default is `None`.
         mask: Union[None, torch.Tensor], optional
             A tensor that is used to mask the sampled latent tensor. Default is `None`.
@@ -1171,7 +1171,7 @@ class TopDownLayer(nn.Module):
             at the correspondent hierarchical layer.
         """
         if bu_value.shape[-2:] != p_params.shape[-2:]:
-            assert self.bottomup_no_padding_mode is True
+            assert self.bottomup_no_padding_mode is True  # TODO WTF ?
             if self.topdown_no_padding_mode is False:
                 assert bu_value.shape[-1] > p_params.shape[-1]
                 bu_value = F.center_crop(bu_value, p_params.shape[-2:])
@@ -1218,7 +1218,7 @@ class TopDownLayer(nn.Module):
             A pre-defined latent tensor. If it is not `None`, than it is used as the actual latent tensor and,
             hence, sampling does not happen. Default is `None`.
         use_mode: bool, optional
-            Wheteher the latent tensor should be set as the latent distribution mode.
+            Whether the latent tensor should be set as the latent distribution mode.
             In the case of Gaussian, the mode coincides with the mean of the distribution.
             Default is `False`.
         force_constant_output: bool, optional
@@ -1230,7 +1230,7 @@ class TopDownLayer(nn.Module):
         use_uncond_mode: bool, optional
             Whether to use the uncoditional distribution p(z) to sample latents in prediction mode.
         var_clip_max: float
-            The maximum value reachable by the log-variance of the latent distribtion.
+            The maximum value reachable by the log-variance of the latent distribution.
             Values exceeding this threshold are clipped.
         """
         # Check consistency of arguments
@@ -1241,7 +1241,7 @@ class TopDownLayer(nn.Module):
         p_params = self.get_p_params(input_, n_img_prior)
 
         # Get the parameters for the latent distribution to sample from
-        if inference_mode:
+        if inference_mode:  # TODO What's this ?
             if self.is_top_layer:
                 q_params = bu_value
                 if mode_pred is False:
@@ -1466,7 +1466,7 @@ class NormalStochasticBlock2d(nn.Module):
             A pre-defined latent tensor. If it is not `None`, than it is used as the actual latent tensor and,
             hence, sampling does not happen.
         use_mode: bool
-            Wheteher the latent tensor should be set as the latent distribution mode.
+            Whether the latent tensor should be set as the latent distribution mode.
             In the case of Gaussian, the mode coincides with the mean of the distribution.
         mode_pred: bool
             Whether the model is prediction mode.
@@ -1501,7 +1501,7 @@ class NormalStochasticBlock2d(nn.Module):
         q_params: torch.Tensor
             The input tensor to be processed.
         var_clip_max: float
-            The maximum value reachable by the log-variance of the latent distribtion.
+            The maximum value reachable by the log-variance of the latent distribution.
             Values exceeding this threshold are clipped.
         """
         _, _, q = self.process_q_params(q_params, var_clip_max)
@@ -1600,7 +1600,7 @@ class NormalStochasticBlock2d(nn.Module):
         p_params: torch.Tensor
             The input tensor to be processed.
         var_clip_max: float
-            The maximum value reachable by the log-variance of the latent distribtion.
+            The maximum value reachable by the log-variance of the latent distribution.
             Values exceeding this threshold are clipped.
         """
         if self.transform_p_params:
@@ -1636,7 +1636,7 @@ class NormalStochasticBlock2d(nn.Module):
         p_params: torch.Tensor
             The input tensor to be processed.
         var_clip_max: float
-            The maximum value reachable by the log-variance of the latent distribtion.
+            The maximum value reachable by the log-variance of the latent distribution.
             Values exceeding this threshold are clipped.
         """
         q_params = self.conv_in_q(q_params)
@@ -1681,7 +1681,7 @@ class NormalStochasticBlock2d(nn.Module):
             A pre-defined latent tensor. If it is not `None`, than it is used as the actual latent
             tensor and, hence, sampling does not happen. Default is `None`.
         use_mode: bool, optional
-            Wheteher the latent tensor should be set as the latent distribution mode.
+            Whether the latent tensor should be set as the latent distribution mode.
             In the case of Gaussian, the mode coincides with the mean of the distribution.
             Default is `False`.
         force_constant_output: bool, optional
@@ -1697,7 +1697,7 @@ class NormalStochasticBlock2d(nn.Module):
             Whether to use the uncoditional distribution p(z) to sample latents in prediction mode.
             Default is `False`.
         var_clip_max: float, optional
-            The maximum value reachable by the log-variance of the latent distribtion.
+            The maximum value reachable by the log-variance of the latent distribution.
             Values exceeding this threshold are clipped. Default is `None`.
         """
         debug_qvar_max = 0
@@ -1928,7 +1928,7 @@ class NonStochasticBlock2d(nn.Module):
             A pre-defined latent tensor. If it is not `None`, than it is used as the actual latent
             tensor and, hence, sampling does not happen. Default is `None`.
         use_mode: bool, optional
-            Wheteher the latent tensor should be set as the latent distribution mode.
+            Whether the latent tensor should be set as the latent distribution mode.
             In the case of Gaussian, the mode coincides with the mean of the distribution.
             Default is `False`.
         force_constant_output: bool, optional
@@ -1944,7 +1944,7 @@ class NonStochasticBlock2d(nn.Module):
             Whether to use the uncoditional distribution p(z) to sample latents in prediction mode.
             Default is `False`.
         var_clip_max: float, optional
-            The maximum value reachable by the log-variance of the latent distribtion.
+            The maximum value reachable by the log-variance of the latent distribution.
             Values exceeding this threshold are clipped. Default is `None`.
         """
         debug_qvar_max = 0