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

careamics/lightning/callbacks/prediction_writer_callback/write_strategy.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Optional, Protocol, Union
+from typing import Any, Protocol, Union
 
 import numpy as np
 from numpy.typing import NDArray
@@ -25,7 +25,7 @@ class WriteStrategy(Protocol):
         trainer: Trainer,
         pl_module: LightningModule,
         prediction: Any,  # TODO: change to expected type
-        batch_indices: Optional[Sequence[int]],
+        batch_indices: Sequence[int] | None,
         batch: Any,  # TODO: change to expected type
         batch_idx: int,
         dataloader_idx: int,
@@ -133,7 +133,7 @@ class CacheTiles(WriteStrategy):
         trainer: Trainer,
         pl_module: LightningModule,
         prediction: tuple[NDArray, list[TileInformation]],
-        batch_indices: Optional[Sequence[int]],
+        batch_indices: Sequence[int] | None,
         batch: tuple[NDArray, list[TileInformation]],
         batch_idx: int,
         dataloader_idx: int,
@@ -259,7 +259,7 @@ class WriteTilesZarr(WriteStrategy):
         trainer: Trainer,
         pl_module: LightningModule,
         prediction: Any,
-        batch_indices: Optional[Sequence[int]],
+        batch_indices: Sequence[int] | None,
         batch: Any,
         batch_idx: int,
         dataloader_idx: int,
@@ -346,7 +346,7 @@ class WriteImage(WriteStrategy):
         trainer: Trainer,
         pl_module: LightningModule,
         prediction: NDArray,
-        batch_indices: Optional[Sequence[int]],
+        batch_indices: Sequence[int] | None,
         batch: NDArray,
         batch_idx: int,
         dataloader_idx: int,

careamics/lightning/callbacks/prediction_writer_callback/write_strategy_factory.py

@@ -1,6 +1,6 @@
 """Module containing convenience function to create `WriteStrategy`."""
 
-from typing import Any, Optional
+from typing import Any
 
 from careamics.config.support import SupportedData
 from careamics.file_io import SupportedWriteType, WriteFunc, get_write_func
@@ -11,9 +11,9 @@ from .write_strategy import CacheTiles, WriteImage, WriteStrategy
 def create_write_strategy(
     write_type: SupportedWriteType,
     tiled: bool,
-    write_func: Optional[WriteFunc] = None,
-    write_extension: Optional[str] = None,
-    write_func_kwargs: Optional[dict[str, Any]] = None,
+    write_func: WriteFunc | None = None,
+    write_extension: str | None = None,
+    write_func_kwargs: dict[str, Any] | None = None,
 ) -> WriteStrategy:
     """
     Create a write strategy from convenient parameters.
@@ -78,8 +78,8 @@ def create_write_strategy(
 
 def _create_tiled_write_strategy(
     write_type: SupportedWriteType,
-    write_func: Optional[WriteFunc],
-    write_extension: Optional[str],
+    write_func: WriteFunc | None,
+    write_extension: str | None,
     write_func_kwargs: dict[str, Any],
 ) -> WriteStrategy:
     """
@@ -130,7 +130,7 @@ def _create_tiled_write_strategy(
 
 
 def select_write_func(
-    write_type: SupportedWriteType, write_func: Optional[WriteFunc] = None
+    write_type: SupportedWriteType, write_func: WriteFunc | None = None
 ) -> WriteFunc:
     """
     Return a function to write images.
@@ -177,7 +177,7 @@ def select_write_func(
 
 
 def select_write_extension(
-    write_type: SupportedWriteType, write_extension: Optional[str] = None
+    write_type: SupportedWriteType, write_extension: str | None = None
 ) -> str:
     """
     Return an extension to add to file paths.

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,7 +1,7 @@
 """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
@@ -90,7 +90,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:
@@ -333,18 +333,16 @@ class VAEModule(L.LightningModule):
         self.model: nn.Module = model_factory(self.algorithm_config.model)
 
         # 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(
-                config=self.algorithm_config.noise_model_likelihood,
-                noise_model=self.noise_model,
-            )
+        self.noise_model_likelihood: NoiseModelLikelihood | None = 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
         )
 
@@ -380,7 +378,7 @@ class VAEModule(L.LightningModule):
 
     def training_step(
         self, batch: tuple[Tensor, Tensor], batch_idx: Any
-    ) -> Optional[dict[str, Tensor]]:
+    ) -> dict[str, Tensor] | None:
         """Training step.
 
         Parameters
@@ -603,7 +601,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 +632,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.
 

careamics/lightning/predict_data_module.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Literal, Optional, Union
+from typing import Any, Literal, Union
 
 import numpy as np
 import pytorch_lightning as L
@@ -65,9 +65,9 @@ class PredictDataModule(L.LightningDataModule):
         self,
         pred_config: InferenceConfig,
         pred_data: Union[Path, str, NDArray],
-        read_source_func: Optional[Callable] = None,
+        read_source_func: Callable | None = None,
         extension_filter: str = "",
-        dataloader_params: Optional[dict] = None,
+        dataloader_params: dict | None = None,
     ) -> None:
         """
         Constructor.
@@ -173,7 +173,7 @@ class PredictDataModule(L.LightningDataModule):
                 self.pred_data, self.data_type, self.extension_filter
             )
 
-    def setup(self, stage: Optional[str] = None) -> None:
+    def setup(self, stage: str | None = None) -> None:
         """
         Hook called at the beginning of predict.
 
@@ -231,13 +231,13 @@ def create_predict_datamodule(
     axes: str,
     image_means: list[float],
     image_stds: list[float],
-    tile_size: Optional[tuple[int, ...]] = None,
-    tile_overlap: Optional[tuple[int, ...]] = None,
+    tile_size: tuple[int, ...] | None = None,
+    tile_overlap: tuple[int, ...] | None = None,
     batch_size: int = 1,
     tta_transforms: bool = True,
-    read_source_func: Optional[Callable] = None,
+    read_source_func: Callable | None = None,
     extension_filter: str = "",
-    dataloader_params: Optional[dict] = None,
+    dataloader_params: dict | None = None,
 ) -> PredictDataModule:
     """Create a CAREamics prediction Lightning datamodule.
 

careamics/lightning/train_data_module.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Literal, Optional, Union
+from typing import Any, Literal, Union
 
 import numpy as np
 import pytorch_lightning as L
@@ -121,10 +121,10 @@ class TrainDataModule(L.LightningDataModule):
         self,
         data_config: DataConfig,
         train_data: Union[Path, str, NDArray],
-        val_data: Optional[Union[Path, str, NDArray]] = None,
-        train_data_target: Optional[Union[Path, str, NDArray]] = None,
-        val_data_target: Optional[Union[Path, str, NDArray]] = None,
-        read_source_func: Optional[Callable] = None,
+        val_data: Union[Path, str, NDArray] | None = None,
+        train_data_target: Union[Path, str, NDArray] | None = None,
+        val_data_target: Union[Path, str, NDArray] | None = None,
+        read_source_func: Callable | None = None,
         extension_filter: str = "",
         val_percentage: float = 0.1,
         val_minimum_split: int = 5,
@@ -477,15 +477,15 @@ def create_train_datamodule(
     patch_size: list[int],
     axes: str,
     batch_size: int,
-    val_data: Optional[Union[str, Path, NDArray]] = None,
-    transforms: Optional[list[TransformModel]] = None,
-    train_target_data: Optional[Union[str, Path, NDArray]] = None,
-    val_target_data: Optional[Union[str, Path, NDArray]] = None,
-    read_source_func: Optional[Callable] = None,
+    val_data: Union[str, Path, NDArray] | None = None,
+    transforms: list[TransformModel] | None = None,
+    train_target_data: Union[str, Path, NDArray] | None = None,
+    val_target_data: Union[str, Path, NDArray] | None = None,
+    read_source_func: Callable | None = None,
     extension_filter: str = "",
     val_percentage: float = 0.1,
     val_minimum_patches: int = 5,
-    dataloader_params: Optional[dict] = None,
+    dataloader_params: dict | None = None,
     use_in_memory: bool = True,
 ) -> TrainDataModule:
     """Create a TrainDataModule.

careamics/losses/lvae/losses.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Literal, Union
 
 import numpy as np
 import torch
@@ -112,7 +112,7 @@ def get_kl_divergence_loss(
     rescaling: Literal["latent_dim", "image_dim"],
     aggregation: Literal["mean", "sum"],
     free_bits_coeff: float,
-    img_shape: Optional[tuple[int]] = None,
+    img_shape: tuple[int] | None = None,
 ) -> torch.Tensor:
     """Compute the KL divergence loss.
 
@@ -273,9 +273,9 @@ def musplit_loss(
     model_outputs: tuple[torch.Tensor, dict[str, Any]],
     targets: torch.Tensor,
     config: LVAELossConfig,
-    gaussian_likelihood: Optional[GaussianLikelihood],
-    noise_model_likelihood: Optional[NoiseModelLikelihood] = None,  # TODO: ugly
-) -> Optional[dict[str, torch.Tensor]]:
+    gaussian_likelihood: GaussianLikelihood | None,
+    noise_model_likelihood: NoiseModelLikelihood | None = None,  # TODO: ugly
+) -> dict[str, torch.Tensor] | None:
     """Loss function for muSplit.
 
     Parameters
@@ -351,9 +351,9 @@ def denoisplit_loss(
     model_outputs: tuple[torch.Tensor, dict[str, Any]],
     targets: torch.Tensor,
     config: LVAELossConfig,
-    gaussian_likelihood: Optional[GaussianLikelihood] = None,
-    noise_model_likelihood: Optional[NoiseModelLikelihood] = None,
-) -> Optional[dict[str, torch.Tensor]]:
+    gaussian_likelihood: GaussianLikelihood | None = None,
+    noise_model_likelihood: NoiseModelLikelihood | None = None,
+) -> dict[str, torch.Tensor] | None:
     """Loss function for DenoiSplit.
 
     Parameters
@@ -430,7 +430,7 @@ def denoisplit_musplit_loss(
     config: LVAELossConfig,
     gaussian_likelihood: GaussianLikelihood,
     noise_model_likelihood: NoiseModelLikelihood,
-) -> Optional[dict[str, torch.Tensor]]:
+) -> dict[str, torch.Tensor] | None:
     """Loss function for DenoiSplit.
 
     Parameters

careamics/model_io/bioimage/model_description.py

@@ -1,10 +1,10 @@
 """Module use to build BMZ model description."""
 
 from pathlib import Path
-from typing import Optional, Union
+from typing import Union
 
 import numpy as np
-from bioimageio.spec._internal.io import resolve_and_extract
+from bioimageio.spec._internal.io import extract
 from bioimageio.spec.model.v0_5 import (
     ArchitectureFromLibraryDescr,
     Author,
@@ -12,7 +12,6 @@ from bioimageio.spec.model.v0_5 import (
     AxisId,
     BatchAxis,
     ChannelAxis,
-    EnvironmentFileDescr,
     FileDescr,
     FixedZeroMeanUnitVarianceAlongAxisKwargs,
     FixedZeroMeanUnitVarianceDescr,
@@ -36,7 +35,7 @@ from ._readme_factory import readme_factory
 def _create_axes(
     array: np.ndarray,
     data_config: DataConfig,
-    channel_names: Optional[list[str]] = None,
+    channel_names: list[str] | None = None,
     is_input: bool = True,
 ) -> list[AxisBase]:
     """Create axes description.
@@ -105,7 +104,7 @@ def _create_inputs_ouputs(
     data_config: DataConfig,
     input_path: Union[Path, str],
     output_path: Union[Path, str],
-    channel_names: Optional[list[str]] = None,
+    channel_names: list[str] | None = None,
 ) -> tuple[InputTensorDescr, OutputTensorDescr]:
     """Create input and output tensor description.
 
@@ -197,7 +196,7 @@ def create_model_description(
     config_path: Union[Path, str],
     env_path: Union[Path, str],
     covers: list[Union[Path, str]],
-    channel_names: Optional[list[str]] = None,
+    channel_names: list[str] | None = None,
     model_version: str = "0.1.0",
 ) -> ModelDescr:
     """Create model description.
@@ -269,7 +268,7 @@ def create_model_description(
             source=weights_path,
             architecture=architecture_descr,
             pytorch_version=Version(torch_version),
-            dependencies=EnvironmentFileDescr(source=env_path),
+            dependencies=FileDescr(source=Path(env_path)),
         ),
     )
 
@@ -322,9 +321,11 @@ def extract_model_path(model_desc: ModelDescr) -> tuple[Path, Path]:
     """
     if model_desc.weights.pytorch_state_dict is None:
         raise ValueError("No model weights found in model description.")
-    weights_path = resolve_and_extract(
-        model_desc.weights.pytorch_state_dict.source
-    ).path
+
+    # extract the zip model and return the directory
+    model_dir = extract(model_desc.root)
+
+    weights_path = model_dir.joinpath(model_desc.weights.pytorch_state_dict.source.path)
 
     for file in model_desc.attachments:
         file_path = file.source if isinstance(file.source, Path) else file.source.path
@@ -332,7 +333,7 @@ def extract_model_path(model_desc: ModelDescr) -> tuple[Path, Path]:
             continue
         file_path = Path(file_path)
         if file_path.name == "careamics.yaml":
-            config_path = resolve_and_extract(file.source).path
+            config_path = model_dir.joinpath(file.source.path)
             break
     else:
         raise ValueError("Configuration file not found.")

careamics/model_io/bmz_io.py

@@ -2,7 +2,7 @@
 
 import tempfile
 from pathlib import Path
-from typing import Optional, Union
+from typing import Union
 
 import numpy as np
 from bioimageio.core import load_model_description, test_model
@@ -90,8 +90,8 @@ def export_to_bmz(
     authors: list[dict],
     input_array: np.ndarray,
     output_array: np.ndarray,
-    covers: Optional[list[Union[Path, str]]] = None,
-    channel_names: Optional[list[str]] = None,
+    covers: list[Union[Path, str]] | None = None,
+    channel_names: list[str] | None = None,
     model_version: str = "0.1.0",
 ) -> None:
     """Export the model to BioImage Model Zoo format.
@@ -187,7 +187,7 @@ def export_to_bmz(
 
         # test model description
         test_kwargs = (
-            model_description.config.get("bioimageio", {})
+            model_description.config.bioimageio.model_dump()
             .get("test_kwargs", {})
             .get("pytorch_state_dict", {})
         )

careamics/models/layers.py

@@ -4,7 +4,7 @@ Layer module.
 This submodule contains layers used in the CAREamics models.
 """
 
-from typing import Optional, Union
+from typing import Union
 
 import torch
 import torch.nn as nn
@@ -207,8 +207,8 @@ def get_pascal_kernel_1d(
     kernel_size: int,
     norm: bool = False,
     *,
-    device: Optional[torch.device] = None,
-    dtype: Optional[torch.dtype] = None,
+    device: torch.device | None = None,
+    dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
     """Generate Yang Hui triangle (Pascal's triangle) for a given number.
 
@@ -270,8 +270,8 @@ def _get_pascal_kernel_nd(
     norm: bool = True,
     dim: int = 2,
     *,
-    device: Optional[torch.device] = None,
-    dtype: Optional[torch.dtype] = None,
+    device: torch.device | None = None,
+    dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
     """Generate pascal filter kernel by kernel size.
 

careamics/prediction_utils/lvae_prediction.py

@@ -1,6 +1,6 @@
 """Module containing pytorch implementations for obtaining predictions from an LVAE."""
 
-from typing import Any, Optional
+from typing import Any
 
 import torch
 
@@ -18,7 +18,7 @@ def lvae_predict_single_sample(
     model: LVAE,
     likelihood_obj: LikelihoodModule,
     input: torch.Tensor,
-) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+) -> tuple[torch.Tensor, torch.Tensor | None]:
     """
     Generate a single sample prediction from an LVAE model, for a given input.
 
@@ -57,7 +57,7 @@ def lvae_predict_tiled_batch(
     model: LVAE,
     likelihood_obj: LikelihoodModule,
     input: tuple[Any],
-) -> tuple[tuple[Any], Optional[tuple[Any]]]:
+) -> tuple[tuple[Any], tuple[Any] | None]:
     # TODO: fix docstring return types, ... too many output options
     """
     Generate a single sample prediction from an LVAE model, for a given input.
@@ -98,7 +98,7 @@ def lvae_predict_mmse_tiled_batch(
     likelihood_obj: LikelihoodModule,
     input: tuple[Any],
     mmse_count: int,
-) -> tuple[tuple[Any], tuple[Any], Optional[tuple[Any]]]:
+) -> tuple[tuple[Any], tuple[Any], tuple[Any] | None]:
     # TODO: fix docstring return types, ... hard to make readable
     """
     Generate the MMSE (minimum mean squared error) prediction, for a given input.
@@ -137,7 +137,7 @@ def lvae_predict_mmse_tiled_batch(
 
     input_shape = x.shape
     output_shape = (input_shape[0], model.target_ch, *input_shape[2:])
-    log_var: Optional[torch.Tensor] = None
+    log_var: torch.Tensor | None = None
     # pre-declare empty array to fill with individual sample predictions
     sample_predictions = torch.zeros(size=(mmse_count, *output_shape))
     for mmse_idx in range(mmse_count):