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

careamics/losses/loss_factory.py

@@ -6,8 +6,9 @@ This module contains a factory function for creating loss functions.
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Callable, Union
+from typing import Union
 
 from torch import Tensor as tensor
 

careamics/lvae_training/dataset/multicrop_dset.py

@@ -1,5 +1,5 @@
 """
-Here, we have multiple folders, each containing images of a single channel. 
+Here, we have multiple folders, each containing images of a single channel.
 """
 
 from collections import defaultdict

careamics/model_io/bioimage/bioimage_utils.py

@@ -45,7 +45,7 @@ def create_env_text(pytorch_version: str, torchvision_version: str) -> str:
     env = (
         f"name: careamics\n"
         f"dependencies:\n"
-        f"  - python=3.10\n"
+        f"  - python=3.12\n"
         f"  - pip\n"
         f"  - pip:\n"
         f"    - torch=={pytorch_version}\n"

careamics/model_io/bioimage/model_description.py

@@ -148,7 +148,7 @@ def _create_inputs_ouputs(
     inv_means = []
     inv_stds = []
     if means and stds:
-        for mean, std in zip(means, stds):
+        for mean, std in zip(means, stds, strict=False):
             inv_means.append(-mean / (std + eps))
             inv_stds.append(1 / (std + eps) - eps)
 

careamics/model_io/bmz_io.py

@@ -200,7 +200,7 @@ def export_to_bmz(
 
 
 def load_from_bmz(
-    path: Union[Path, str, HttpUrl]
+    path: Union[Path, str, HttpUrl],
 ) -> tuple[Union[FCNModule, VAEModule], Configuration]:
     """Load a model from a BioImage Model Zoo archive.
 

careamics/model_io/model_io_utils.py

@@ -12,7 +12,7 @@ from careamics.utils import check_path_exists
 
 
 def load_pretrained(
-    path: Union[Path, str]
+    path: Union[Path, str],
 ) -> tuple[Union[FCNModule, VAEModule], Configuration]:
     """
     Load a pretrained model from a checkpoint or a BioImage Model Zoo model.
@@ -47,7 +47,7 @@ def load_pretrained(
 
 
 def _load_checkpoint(
-    path: Union[Path, str]
+    path: Union[Path, str],
 ) -> tuple[Union[FCNModule, VAEModule], Configuration]:
     """
     Load a model from a checkpoint and return both model and configuration.

careamics/models/activation.py

@@ -1,6 +1,7 @@
 """Activations for CAREamics models."""
 
-from typing import Callable, Union
+from collections.abc import Callable
+from typing import Union
 
 import torch.nn as nn
 

careamics/models/unet.py

@@ -205,16 +205,23 @@ class UnetDecoder(nn.Module):
         decoder_blocks: list[nn.Module] = []
         for n in range(depth):
             decoder_blocks.append(upsampling)
-            in_channels = (num_channels_init * 2 ** (depth - n)) * groups
-            out_channels = in_channels // 2
+
+            in_channels = (num_channels_init * 2 ** (depth - n - 1)) * groups
+            # final decoder block has the same number in and out features
+            out_channels = in_channels // 2 if n != depth - 1 else in_channels
+            if not (n2v2 and (n == depth - 1)):
+                in_channels = in_channels * 2  # accounting for skip connection concat
+
             decoder_blocks.append(
                 Conv_Block(
                     conv_dim,
-                    in_channels=(
-                        in_channels + in_channels // 2 if n > 0 else in_channels
-                    ),
+                    in_channels=in_channels,
                     out_channels=out_channels,
-                    intermediate_channel_multiplier=2,
+                    # TODO: Tensorflow n2v implementation has intermediate channel
+                    #   multiplication for skip_skipone=True but not skip_skipone=False
+                    #   this needs to be benchmarked.
+                    # final decoder block doesn't multiply the intermediate features
+                    intermediate_channel_multiplier=2 if n != depth - 1 else 1,
                     dropout_perc=dropout,
                     activation="ReLU",
                     use_batch_norm=use_batch_norm,
@@ -241,6 +248,7 @@ class UnetDecoder(nn.Module):
         """
         x: torch.Tensor = features[0]
         skip_connections: tuple[torch.Tensor, ...] = features[-1:0:-1]
+        depth = len(skip_connections)
 
         x = self.bottleneck(x)
 
@@ -249,10 +257,8 @@ class UnetDecoder(nn.Module):
             if isinstance(module, nn.Upsample):
                 # divide index by 2 because of upsampling layers
                 skip_connection: torch.Tensor = skip_connections[i // 2]
-                if self.n2v2:
-                    if x.shape != skip_connections[-1].shape:
-                        x = self._interleave(x, skip_connection, self.groups)
-                else:
+                # top level skip connection not added for n2v2
+                if (not self.n2v2) or (self.n2v2 and (i // 2 < depth - 1)):
                     x = self._interleave(x, skip_connection, self.groups)
         return x
 

careamics/prediction_utils/prediction_outputs.py

@@ -87,7 +87,7 @@ def combine_batches(
 
 
 def _combine_tiled_batches(
-    predictions: list[tuple[NDArray, list[TileInformation]]]
+    predictions: list[tuple[NDArray, list[TileInformation]]],
 ) -> tuple[list[NDArray], list[TileInformation]]:
     """
     Combine batches from tiled output.

careamics/prediction_utils/stitch_prediction.py

@@ -94,7 +94,7 @@ def stitch_prediction_single(
     input_shape = (1, tile_channels, *tile_infos[0].array_shape[1:])
     predicted_image = np.zeros(input_shape, dtype=np.float32)
 
-    for tile, tile_info in zip(tiles, tile_infos):
+    for tile, tile_info in zip(tiles, tile_infos, strict=False):
 
         # Compute coordinates for cropping predicted tile
         crop_slices: tuple[Union[builtins.ellipsis, slice], ...] = (

careamics/transforms/n2v_manipulate_torch.py

@@ -27,6 +27,8 @@ class N2VManipulateTorch:
         N2V manipulation configuration.
     seed : Optional[int], optional
         Random seed, by default None.
+    device : str
+        The device on which operations take place, e.g. "cuda", "cpu" or "mps".
 
     Attributes
     ----------
@@ -48,6 +50,7 @@ class N2VManipulateTorch:
         self,
         n2v_manipulate_config: N2VManipulateModel,
         seed: Optional[int] = None,
+        device: Optional[str] = None,
     ):
         """Constructor.
 
@@ -57,6 +60,8 @@ class N2VManipulateTorch:
             N2V manipulation configuration.
         seed : Optional[int], optional
             Random seed, by default None.
+        device : str
+            The device on which operations take place, e.g. "cuda", "cpu" or "mps".
         """
         self.masked_pixel_percentage = n2v_manipulate_config.masked_pixel_percentage
         self.roi_size = n2v_manipulate_config.roi_size
@@ -78,15 +83,16 @@ class N2VManipulateTorch:
 
         # PyTorch random generator
         # TODO refactor into careamics.utils.torch_utils.get_device
-        if torch.cuda.is_available():
-            device = "cuda"
-        elif torch.backends.mps.is_available() and platform.processor() in (
-            "arm",
-            "arm64",
-        ):
-            device = "mps"
-        else:
-            device = "cpu"
+        if device is None:
+            if torch.cuda.is_available():
+                device = "cuda"
+            elif torch.backends.mps.is_available() and platform.processor() in (
+                "arm",
+                "arm64",
+            ):
+                device = "mps"
+            else:
+                device = "cpu"
 
         self.rng = (
             torch.Generator(device=device).manual_seed(seed)

careamics/transforms/pixel_manipulation_torch.py

@@ -75,50 +75,23 @@ def _apply_struct_mask_torch(
     return patch
 
 
-def _odd_jitter_func_torch(step: float, rng: torch.Generator) -> torch.Tensor:
-    """
-    Randomly sample a jitter to be applied to the masking grid.
-
-    This is done to account for cases where the step size is not an integer.
-
-    Parameters
-    ----------
-    step : float
-        Step size of the grid, output of np.linspace.
-    rng : torch.Generator
-        Random number generator.
-
-    Returns
-    -------
-    torch.Tensor
-        Array of random jitter to be added to the grid.
-    """
-    step_floor = torch.floor(torch.tensor(step))
-    odd_jitter = (
-        step_floor
-        if step_floor == step
-        else torch.randint(high=2, size=(1,), generator=rng)
-    )
-    return step_floor if odd_jitter == 0 else step_floor + 1
-
-
 def _get_stratified_coords_torch(
     mask_pixel_perc: float,
     shape: tuple[int, ...],
-    rng: Optional[torch.Generator] = None,
+    rng: torch.Generator,
 ) -> torch.Tensor:
     """
     Generate coordinates of the pixels to mask.
 
-    # TODO add more details
     Randomly selects the coordinates of the pixels to mask in a stratified way, i.e.
-    the distance between masked pixels is approximately the same.
+    the distance between masked pixels is approximately the same. This is achieved by
+    defining a grid and sampling a pixel in each grid square. The grid is defined such
+    that the resulting density of masked pixels is the desired masked pixel percentage.
 
     Parameters
     ----------
     mask_pixel_perc : float
-        Actual (quasi) percentage of masked pixels across the whole image. Used in
-        calculating the distance between masked pixels across each axis.
+        Expected value for percentage of masked pixels across the whole image.
     shape : tuple[int, ...]
         Shape of the input patch.
     rng : torch.Generator or None
@@ -129,60 +102,51 @@ def _get_stratified_coords_torch(
     np.ndarray
         Array of coordinates of the masked pixels.
     """
-    if rng is None:
-        rng = torch.Generator()
-
-    # Calculate the maximum distance between masked pixels. Inversely proportional to
-    # the percentage of masked pixels.
-    mask_pixel_distance = round((100 / mask_pixel_perc) ** (1 / len(shape)))
-
-    pixel_coords = []
-    steps = []
-
-    # loop over dimensions
-    for axis_size in shape:
-        # number of pixels to mask along the axis
-        num_pixels = int(torch.ceil(torch.tensor(axis_size / mask_pixel_distance)))
-
-        # create 1D grid of coordinates for the axis
-        axis_pixel_coords = torch.linspace(
-            0,
-            axis_size - (axis_size // num_pixels),
-            num_pixels,
-            dtype=torch.int32,
-        )
-
-        # calculate the step size between coordinates
-        step = (
-            axis_pixel_coords[1] - axis_pixel_coords[0]
-            if len(axis_pixel_coords) > 1
-            else axis_size
-        )
-
-        pixel_coords.append(axis_pixel_coords)
-        steps.append(step)
-
-    # create a 2D meshgrid of coordinates
-    coordinate_grid_list = torch.meshgrid(*pixel_coords, indexing="ij")
-    coordinate_grid = torch.stack(
-        [g.flatten() for g in coordinate_grid_list], dim=-1
-    ).to(rng.device)
-
-    # add a random jitter increment so that the coordinates do not lie on the grid
-    random_increment = torch.randint(
-        high=int(_odd_jitter_func_torch(float(max(steps)), rng)),
-        size=torch.tensor(coordinate_grid.shape).to(rng.device).tolist(),
-        generator=rng,
-        device=rng.device,
+    # Implementation logic:
+    #    find a box size s.t sampling 1 pixel within the box will result in the desired
+    # pixel percentage. Make a grid of these boxes that cover the patch (the area of
+    # the grid will be greater than or equal to the area of the patch) and sample 1
+    # pixel in each box. The density of masked pixels is an intensive property therefore
+    # any subset of this area will have the desired expected masked pixel percentage.
+    # We can get our desired patch with our desired expected masked pixel percentage by
+    # simply filtering out masked pixels that lie outside of our patch bounds.
+
+    batch_size = shape[0]
+    spatial_shape = shape[1:]
+
+    n_dims = len(spatial_shape)
+    expected_area_per_pixel = 1 / (mask_pixel_perc / 100)
+
+    # keep the grid size in floats for a more accurate expected masked pixel percentage
+    grid_size = expected_area_per_pixel ** (1 / n_dims)
+    grid_dims = torch.ceil(torch.tensor(spatial_shape) / grid_size).int()
+
+    # coords on a fixed grid (top left corner)
+    coords = torch.stack(
+        torch.meshgrid(
+            torch.arange(batch_size, dtype=torch.float),
+            *[torch.arange(0, grid_dims[i].item()) * grid_size for i in range(n_dims)],
+            indexing="ij",
+        ),
+        -1,
+    ).reshape(-1, n_dims + 1)
+
+    # add random offset to get a random coord in each grid box
+    # also keep the offset in floats
+    offset = (
+        torch.rand((len(coords), n_dims), device=rng.device, generator=rng) * grid_size
     )
-    coordinate_grid += random_increment
+    coords = coords.to(rng.device)
+    coords[:, 1:] += offset
+    coords = torch.floor(coords).int()
 
-    # make sure no coordinate lie outside the range
-    return torch.clamp(
-        coordinate_grid,
-        torch.zeros_like(torch.tensor(shape)).to(device=rng.device),
-        torch.tensor([v - 1 for v in shape]).to(device=rng.device),
-    )
+    # filter pixels out of bounds
+    out_of_bounds = (
+        coords[:, 1:]
+        >= torch.tensor(spatial_shape, device=rng.device).reshape(1, n_dims)
+    ).any(1)
+    coords = coords[~out_of_bounds]
+    return coords
 
 
 def uniform_manipulate_torch(
@@ -198,7 +162,7 @@ def uniform_manipulate_torch(
 
     # TODO add more details, especially about batch
 
-    Manipulated pixels are selected unformly selected in a subpatch, away from a grid
+    Manipulated pixels are selected uniformly selected in a subpatch, away from a grid
     with an approximate uniform probability to be selected across the whole patch.
     If `struct_params` is not None, an additional structN2V mask is applied to the
     data, replacing the pixels in the mask with random values (excluding the pixel
@@ -254,18 +218,21 @@ def uniform_manipulate_torch(
     random_increment = roi_span[
         torch.randint(
             low=min(roi_span),
-            high=max(roi_span) + 1,  # TODO check this, it may exclude one value
-            size=subpatch_centers.shape,
+            high=max(roi_span) + 1,
+            # one less coord dim: we shouldn't add a random increment to the batch coord
+            size=(subpatch_centers.shape[0], subpatch_centers.shape[1] - 1),
             generator=rng,
             device=patch.device,
         )
     ]
 
     # compute the replacement pixel coordinates
-    replacement_coords = torch.clamp(
-        subpatch_centers + random_increment,
-        torch.zeros_like(torch.tensor(patch.shape)).to(device=patch.device),
-        torch.tensor([v - 1 for v in patch.shape]).to(device=patch.device),
+    replacement_coords = subpatch_centers.clone()
+    # only add random increment to the spatial dimensions, not the batch dimension
+    replacement_coords[:, 1:] = torch.clamp(
+        replacement_coords[:, 1:] + random_increment,
+        torch.zeros_like(torch.tensor(patch.shape[1:])).to(device=patch.device),
+        torch.tensor([v - 1 for v in patch.shape[1:]]).to(device=patch.device),
     )
 
     # replace the pixels in the patch
@@ -313,7 +280,7 @@ def median_manipulate_torch(
     struct_params : StructMaskParameters or None, optional
         Parameters for the structN2V mask (axis and span).
     rng : torch.default_generator or None, optional
-        Random number generato, by default None.
+        Random number generator, by default None.
 
     Returns
     -------

careamics/utils/lightning_utils.py

@@ -32,13 +32,13 @@ def read_csv_logger(experiment_name: str, log_folder: Union[str, Path]) -> dict:
         lines = f.readlines()
 
         header = lines[0].strip().split(",")
-        metrics = {value: [] for value in header}
+        metrics: dict[str, list] = {value: [] for value in header}
         print(metrics)
 
         for single_line in lines[1:]:
             values = single_line.strip().split(",")
 
-            for k, v in zip(header, values):
+            for k, v in zip(header, values, strict=False):
                 metrics[k].append(v)
 
     # train and val are not logged on the same row and can have different lengths

careamics/utils/metrics.py

@@ -4,7 +4,8 @@ Metrics submodule.
 This module contains various metrics and a metrics tracking class.
 """
 
-from typing import Callable, Optional, Union
+from collections.abc import Callable
+from typing import Optional, Union
 
 import numpy as np
 import torch

careamics/utils/torch_utils.py

@@ -5,6 +5,7 @@ These functions are used to control certain aspects and behaviours of PyTorch.
 """
 
 import inspect
+import platform
 from typing import Union
 
 import torch
@@ -16,6 +17,28 @@ from ..utils.logging import get_logger
 logger = get_logger(__name__)  # TODO are logger still needed?
 
 
+def get_device() -> str:
+    """
+    Get the device on which operations take place.
+
+    Returns
+    -------
+    str
+        The device on which operations take place, e.g. "cuda", "cpu" or "mps".
+    """
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.backends.mps.is_available() and platform.processor() in (
+        "arm",
+        "arm64",
+    ):
+        device = "mps"
+    else:
+        device = "cpu"
+
+    return device
+
+
 def filter_parameters(
     func: type,
     user_params: dict,

{careamics-0.0.12.dist-info → careamics-0.0.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: careamics
-Version: 0.0.12
+Version: 0.0.14
 Summary: Toolbox for running N2V and friends.
 Project-URL: homepage, https://careamics.github.io/
 Project-URL: repository, https://github.com/CAREamics/careamics
@@ -10,27 +10,28 @@ License-File: LICENSE
 Classifier: Development Status :: 3 - Alpha
 Classifier: License :: OSI Approved :: BSD License
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Typing :: Typed
-Requires-Python: >=3.9
+Requires-Python: >=3.10
 Requires-Dist: bioimageio-core==0.7
-Requires-Dist: matplotlib<=3.10.1
+Requires-Dist: matplotlib<=3.10.3
 Requires-Dist: numpy<2.0.0
 Requires-Dist: pillow<=11.2.1
 Requires-Dist: psutil<=7.0.0
 Requires-Dist: pydantic<=2.12,>=2.11
-Requires-Dist: pytorch-lightning<=2.5.0.post0,>=2.2
+Requires-Dist: pytorch-lightning<=2.5.2,>=2.2
 Requires-Dist: pyyaml!=6.0.0,<=6.0.2
 Requires-Dist: scikit-image<=0.25.2
-Requires-Dist: tifffile<=2025.3.13
-Requires-Dist: torch<=2.6.0,>=2.0
-Requires-Dist: torchvision<=0.21.0
-Requires-Dist: typer<=0.15.2,>=0.12.3
+Requires-Dist: tifffile<=2025.5.10
+Requires-Dist: torch<=2.7.1,>=2.0
+Requires-Dist: torchvision<=0.22.1
+Requires-Dist: typer<=0.16.0,>=0.12.3
 Requires-Dist: xarray<2025.3.0
 Requires-Dist: zarr<3.0.0
+Provides-Extra: czi
+Requires-Dist: pylibczirw<6.0.0,>=4.1.2; extra == 'czi'
 Provides-Extra: dev
 Requires-Dist: onnx; extra == 'dev'
 Requires-Dist: pre-commit; extra == 'dev'