careamics 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

careamics/models/layers.py

@@ -4,7 +4,7 @@ Layer module.
 This submodule contains layers used in the CAREamics models.
 """
 
-from typing import List, Optional, Tuple, Union
+from typing import Optional, Union
 
 import torch
 import torch.nn as nn
@@ -157,22 +157,22 @@ class Conv_Block(nn.Module):
 
 
 def _unpack_kernel_size(
-    kernel_size: Union[Tuple[int, ...], int], dim: int
-) -> Tuple[int, ...]:
+    kernel_size: Union[tuple[int, ...], int], dim: int
+) -> tuple[int, ...]:
     """Unpack kernel_size to a tuple of ints.
 
     Inspired by Kornia implementation. TODO: link
 
     Parameters
     ----------
-    kernel_size : Union[Tuple[int, ...], int]
+    kernel_size : Union[tuple[int, ...], int]
         Kernel size.
     dim : int
         Number of dimensions.
 
     Returns
     -------
-    Tuple[int, ...]
+    tuple[int, ...]
         Kernel size tuple.
     """
     if isinstance(kernel_size, int):
@@ -183,20 +183,20 @@ def _unpack_kernel_size(
 
 
 def _compute_zero_padding(
-    kernel_size: Union[Tuple[int, ...], int], dim: int
-) -> Tuple[int, ...]:
+    kernel_size: Union[tuple[int, ...], int], dim: int
+) -> tuple[int, ...]:
     """Utility function that computes zero padding tuple.
 
     Parameters
     ----------
-    kernel_size : Union[Tuple[int, ...], int]
+    kernel_size : Union[tuple[int, ...], int]
         Kernel size.
     dim : int
         Number of dimensions.
 
     Returns
     -------
-    Tuple[int, ...]
+    tuple[int, ...]
         Zero padding tuple.
     """
     kernel_dims = _unpack_kernel_size(kernel_size, dim)
@@ -245,8 +245,8 @@ def get_pascal_kernel_1d(
     >>> get_pascal_kernel_1d(6)
     tensor([ 1.,  5., 10., 10.,  5.,  1.])
     """
-    pre: List[float] = []
-    cur: List[float] = []
+    pre: list[float] = []
+    cur: list[float] = []
     for i in range(kernel_size):
         cur = [1.0] * (i + 1)
 
@@ -266,7 +266,7 @@ def get_pascal_kernel_1d(
 
 
 def _get_pascal_kernel_nd(
-    kernel_size: Union[Tuple[int, int], int],
+    kernel_size: Union[tuple[int, int], int],
     norm: bool = True,
     dim: int = 2,
     *,
@@ -282,7 +282,7 @@ def _get_pascal_kernel_nd(
 
     Parameters
     ----------
-    kernel_size : Union[Tuple[int, int], int]
+    kernel_size : Union[tuple[int, int], int]
         Kernel size for the pascal kernel.
     norm : bool
         Normalize the kernel, by default True.
@@ -420,7 +420,7 @@ class MaxBlurPool(nn.Module):
     ----------
     dim : int
         Toggles between 2D and 3D.
-    kernel_size : Union[Tuple[int, int], int]
+    kernel_size : Union[tuple[int, int], int]
         Kernel size for max pooling.
     stride : int
         Stride for pooling.
@@ -433,7 +433,7 @@ class MaxBlurPool(nn.Module):
     def __init__(
         self,
         dim: int,
-        kernel_size: Union[Tuple[int, int], int],
+        kernel_size: Union[tuple[int, int], int],
         stride: int = 2,
         max_pool_size: int = 2,
         ceil_mode: bool = False,
@@ -444,7 +444,7 @@ class MaxBlurPool(nn.Module):
         ----------
         dim : int
             Dimension of the convolution.
-        kernel_size : Union[Tuple[int, int], int]
+        kernel_size : Union[tuple[int, int], int]
             Kernel size for max pooling.
         stride : int, optional
             Stride, by default 2.

careamics/models/lvae/likelihoods.py

@@ -324,6 +324,8 @@ class NoiseModelLikelihood(LikelihoodModule):
         if self.data_mean.device != correct_device_tensor.device:
             self.data_mean = self.data_mean.to(correct_device_tensor.device)
             self.data_std = self.data_std.to(correct_device_tensor.device)
+        if correct_device_tensor.device != self.noiseModel.device:
+            self.noiseModel.to_device(correct_device_tensor.device)
 
     def get_mean_lv(self, x: torch.Tensor) -> tuple[torch.Tensor, None]:
         return x, None

careamics/models/lvae/lvae.py

@@ -6,14 +6,12 @@ and Artefact Removal, Prakash et al."
 """
 
 from collections.abc import Iterable
-from typing import Union
+from typing import Optional, Union
 
 import numpy as np
 import torch
 import torch.nn as nn
 
-from careamics.config.architectures import register_model
-
 from ..activation import get_activation
 from .layers import (
     BottomUpDeterministicResBlock,
@@ -25,7 +23,6 @@ from .layers import (
 from .utils import Interpolate, ModelType, crop_img_tensor
 
 
-@register_model("LVAE")
 class LadderVAE(nn.Module):
     """
     Constructor.
@@ -837,3 +834,15 @@ class LadderVAE(nn.Module):
             # TODO check if model_3D_depth is needed ?
             top_layer_shape = (n_imgs, mu_logvar, self._model_3D_depth, h, w)
         return top_layer_shape
+
+    def reset_for_inference(self, tile_size: Optional[tuple[int, int]] = None):
+        """Should be called if we want to predict for a different input/output size."""
+        self.mode_pred = True
+        if tile_size is None:
+            tile_size = self.image_size
+        self.image_size = tile_size
+        for i in range(self.n_layers):
+            self.bottom_up_layers[i].output_expected_shape = (
+                ts // 2 ** (i + 1) for ts in tile_size
+            )
+            self.top_down_layers[i].latent_shape = tile_size