careamics 0.0.11__py3-none-any.whl → 0.0.13__py3-none-any.whl

careamics/careamist.py

@@ -1,7 +1,8 @@
 """A class to train, predict and export models in CAREamics."""
 
+from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Callable, Literal, Optional, Union, overload
+from typing import Any, Literal, Optional, Union, overload
 
 import numpy as np
 from numpy.typing import NDArray
@@ -52,6 +53,9 @@ class CAREamist:
         by default None.
     callbacks : list of Callback, optional
         List of callbacks to use during training and prediction, by default None.
+    enable_progress_bar : bool
+        Whether a progress bar will be displayed during training, validation and
+        prediction.
 
     Attributes
     ----------
@@ -77,6 +81,7 @@ class CAREamist:
         source: Union[Path, str],
         work_dir: Optional[Union[Path, str]] = None,
         callbacks: Optional[list[Callback]] = None,
+        enable_progress_bar: bool = True,
     ) -> None: ...
 
     @overload
@@ -85,6 +90,7 @@ class CAREamist:
         source: Configuration,
         work_dir: Optional[Union[Path, str]] = None,
         callbacks: Optional[list[Callback]] = None,
+        enable_progress_bar: bool = True,
     ) -> None: ...
 
     def __init__(
@@ -92,6 +98,7 @@ class CAREamist:
         source: Union[Path, str, Configuration],
         work_dir: Optional[Union[Path, str]] = None,
         callbacks: Optional[list[Callback]] = None,
+        enable_progress_bar: bool = True,
     ) -> None:
         """
         Initialize CAREamist with a configuration object or a path.
@@ -112,6 +119,9 @@ class CAREamist:
             by default None.
         callbacks : list of Callback, optional
             List of callbacks to use during training and prediction, by default None.
+        enable_progress_bar : bool
+            Whether a progress bar will be displayed during training, validation and
+            prediction.
 
         Raises
         ------
@@ -169,7 +179,7 @@ class CAREamist:
                 self.model, self.cfg = load_pretrained(source)
 
         # define the checkpoint saving callback
-        self._define_callbacks(callbacks)
+        self._define_callbacks(callbacks, enable_progress_bar)
 
         # instantiate logger
         csv_logger = CSVLogger(
@@ -202,7 +212,7 @@ class CAREamist:
             precision=self.cfg.training_config.precision,
             max_steps=self.cfg.training_config.max_steps,
             check_val_every_n_epoch=self.cfg.training_config.check_val_every_n_epoch,
-            enable_progress_bar=self.cfg.training_config.enable_progress_bar,
+            enable_progress_bar=enable_progress_bar,
             accumulate_grad_batches=self.cfg.training_config.accumulate_grad_batches,
             gradient_clip_val=self.cfg.training_config.gradient_clip_val,
             gradient_clip_algorithm=self.cfg.training_config.gradient_clip_algorithm,
@@ -215,13 +225,19 @@ class CAREamist:
         self.train_datamodule: Optional[TrainDataModule] = None
         self.pred_datamodule: Optional[PredictDataModule] = None
 
-    def _define_callbacks(self, callbacks: Optional[list[Callback]] = None) -> None:
+    def _define_callbacks(
+        self, callbacks: Optional[list[Callback]], enable_progress_bar: bool
+    ) -> None:
         """Define the callbacks for the training loop.
 
         Parameters
         ----------
         callbacks : list of Callback, optional
             List of callbacks to use during training and prediction, by default None.
+        enable_progress_bar : bool
+            Whether a progress bar will be displayed during training, validation and
+            prediction. It controls whether a `ProgressBarCallback` is added to the
+            callback list.
         """
         self.callbacks = [] if callbacks is None else callbacks
 
@@ -251,9 +267,10 @@ class CAREamist:
                     filename=self.cfg.experiment_name,
                     **self.cfg.training_config.checkpoint_callback.model_dump(),
                 ),
-                ProgressBarCallback(),
             ]
         )
+        if enable_progress_bar:
+            self.callbacks.append(ProgressBarCallback())
 
         # early stopping callback
         if self.cfg.training_config.early_stopping_callback is not None:
@@ -811,7 +828,7 @@ class CAREamist:
             source_path = source.pred_data
             source_data_type = source.data_type
             extension_filter = source.extension_filter
-        elif isinstance(source, (str, Path)):
+        elif isinstance(source, str | Path):
             source_path = source
             source_data_type = data_type or self.cfg.data_config.data_type
             extension_filter = SupportedData.get_extension_pattern(
@@ -824,7 +841,7 @@ class CAREamist:
             raise ValueError(
                 "Predicting to disk is not supported for input type 'array'."
             )
-        assert isinstance(source_path, (Path, str))  # because data_type != "array"
+        assert isinstance(source_path, str | Path)  # because data_type != "array"
         source_path = Path(source_path)
 
         file_paths = list_files(source_path, source_data_type, extension_filter)

careamics/cli/utils.py

@@ -4,7 +4,7 @@ from typing import Optional
 
 
 def handle_2D_3D_callback(
-    value: Optional[tuple[int, int, int]]
+    value: Optional[tuple[int, int, int]],
 ) -> Optional[tuple[int, ...]]:
     """
     Callback for options that require 2D or 3D inputs.

careamics/config/algorithms/n2v_algorithm_model.py

@@ -1,4 +1,4 @@
-""""N2V Algorithm configuration."""
+"""N2V Algorithm configuration."""
 
 from typing import Annotated, Literal
 

careamics/config/architectures/unet_model.py

@@ -63,6 +63,9 @@ class UNetModel(ArchitectureModel):
     """Whether information is processed independently in each channel, used to train
     channels independently."""
 
+    use_batch_norm: bool = Field(default=True, validate_default=True)
+    """Whether to use batch normalization in the model."""
+
     @field_validator("num_channels_init")
     @classmethod
     def validate_num_channels_init(cls, num_channels_init: int) -> int:

careamics/config/callback_model.py

@@ -22,52 +22,42 @@ class CheckpointModel(BaseModel):
     https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html#modelcheckpoint
     """
 
-    model_config = ConfigDict(
-        validate_assignment=True,
-    )
+    model_config = ConfigDict(validate_assignment=True, validate_default=True)
 
-    monitor: Literal["val_loss"] = Field(default="val_loss", validate_default=True)
-    """Quantity to monitor."""
+    monitor: Literal["val_loss"] = Field(default="val_loss")
+    """Quantity to monitor, currently only `val_loss`."""
 
-    verbose: bool = Field(default=False, validate_default=True)
+    verbose: bool = Field(default=False)
     """Verbosity mode."""
 
-    save_weights_only: bool = Field(default=False, validate_default=True)
+    save_weights_only: bool = Field(default=False)
     """When `True`, only the model's weights will be saved (model.save_weights)."""
 
-    save_last: Optional[Literal[True, False, "link"]] = Field(
-        default=True, validate_default=True
-    )
+    save_last: Optional[Literal[True, False, "link"]] = Field(default=True)
     """When `True`, saves a last.ckpt copy whenever a checkpoint file gets saved."""
 
-    save_top_k: int = Field(default=3, ge=1, le=10, validate_default=True)
+    save_top_k: int = Field(default=3, ge=-1, le=100)
     """If `save_top_k == kz, the best k models according to the quantity monitored
     will be saved. If `save_top_k == 0`, no models are saved. if `save_top_k == -1`,
     all models are saved."""
 
-    mode: Literal["min", "max"] = Field(default="min", validate_default=True)
+    mode: Literal["min", "max"] = Field(default="min")
     """One of {min, max}. If `save_top_k != 0`, the decision to overwrite the current
     save file is made based on either the maximization or the minimization of the
     monitored quantity. For 'val_acc', this should be 'max', for 'val_loss' this should
     be 'min', etc.
     """
 
-    auto_insert_metric_name: bool = Field(default=False, validate_default=True)
+    auto_insert_metric_name: bool = Field(default=False)
     """When `True`, the checkpoints filenames will contain the metric name."""
 
-    every_n_train_steps: Optional[int] = Field(
-        default=None, ge=1, le=10, validate_default=True
-    )
+    every_n_train_steps: Optional[int] = Field(default=None, ge=1, le=1000)
     """Number of training steps between checkpoints."""
 
-    train_time_interval: Optional[timedelta] = Field(
-        default=None, validate_default=True
-    )
+    train_time_interval: Optional[timedelta] = Field(default=None)
     """Checkpoints are monitored at the specified time interval."""
 
-    every_n_epochs: Optional[int] = Field(
-        default=None, ge=1, le=10, validate_default=True
-    )
+    every_n_epochs: Optional[int] = Field(default=None, ge=1, le=100)
     """Number of epochs between checkpoints."""
 
 
@@ -83,41 +73,40 @@ class EarlyStoppingModel(BaseModel):
 
     model_config = ConfigDict(
         validate_assignment=True,
+        validate_default=True,
     )
 
-    monitor: Literal["val_loss"] = Field(default="val_loss", validate_default=True)
+    monitor: Literal["val_loss"] = Field(default="val_loss")
     """Quantity to monitor."""
 
-    min_delta: float = Field(default=0.0, ge=0.0, le=1.0, validate_default=True)
+    min_delta: float = Field(default=0.0, ge=0.0, le=1.0)
     """Minimum change in the monitored quantity to qualify as an improvement, i.e. an
     absolute change of less than or equal to min_delta, will count as no improvement."""
 
-    patience: int = Field(default=3, ge=1, le=10, validate_default=True)
+    patience: int = Field(default=3, ge=1, le=10)
     """Number of checks with no improvement after which training will be stopped."""
 
-    verbose: bool = Field(default=False, validate_default=True)
+    verbose: bool = Field(default=False)
     """Verbosity mode."""
 
-    mode: Literal["min", "max", "auto"] = Field(default="min", validate_default=True)
+    mode: Literal["min", "max", "auto"] = Field(default="min")
     """One of {min, max, auto}."""
 
-    check_finite: bool = Field(default=True, validate_default=True)
+    check_finite: bool = Field(default=True)
     """When `True`, stops training when the monitored quantity becomes `NaN` or
     `inf`."""
 
-    stopping_threshold: Optional[float] = Field(default=None, validate_default=True)
+    stopping_threshold: Optional[float] = Field(default=None)
     """Stop training immediately once the monitored quantity reaches this threshold."""
 
-    divergence_threshold: Optional[float] = Field(default=None, validate_default=True)
+    divergence_threshold: Optional[float] = Field(default=None)
     """Stop training as soon as the monitored quantity becomes worse than this
     threshold."""
 
-    check_on_train_epoch_end: Optional[bool] = Field(
-        default=False, validate_default=True
-    )
+    check_on_train_epoch_end: Optional[bool] = Field(default=False)
     """Whether to run early stopping at the end of the training epoch. If this is
     `False`, then the check runs at the end of the validation."""
 
-    log_rank_zero_only: bool = Field(default=False, validate_default=True)
+    log_rank_zero_only: bool = Field(default=False)
     """When set `True`, logs the status of the early stopping callback only for rank 0
     process."""

careamics/config/configuration.py

@@ -3,11 +3,14 @@
 from __future__ import annotations
 
 import re
+from collections.abc import Callable
 from pprint import pformat
 from typing import Any, Literal, Union
 
+import numpy as np
 from bioimageio.spec.generic.v0_3 import CiteEntry
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+from pydantic.main import IncEx
 from typing_extensions import Self
 
 from careamics.config.algorithms import (
@@ -182,6 +185,50 @@ class Configuration(BaseModel):
 
         return name
 
+    @model_validator(mode="after")
+    def validate_n2v_mask_pixel_perc(self: Self) -> Self:
+        """
+        Validate that there will always be at least one blind-spot pixel in every patch.
+
+        The probability of creating a blind-spot pixel is a function of the chosen
+        masked pixel percentage and patch size.
+
+        Returns
+        -------
+        Self
+            Validated configuration.
+
+        Raises
+        ------
+        ValueError
+            If the probability of masking a pixel within a patch is less than 1 for the
+            chosen masked pixel percentage and patch size.
+        """
+        # No validation needed for non n2v algorithms
+        if not isinstance(self.algorithm_config, N2VAlgorithm):
+            return self
+
+        mask_pixel_perc = self.algorithm_config.n2v_config.masked_pixel_percentage
+        patch_size = self.data_config.patch_size
+        expected_area_per_pixel = 1 / (mask_pixel_perc / 100)
+
+        n_dims = 3 if self.algorithm_config.model.is_3D() else 2
+        patch_size_lower_bound = int(np.ceil(expected_area_per_pixel ** (1 / n_dims)))
+        required_patch_size = tuple(
+            2 ** int(np.ceil(np.log2(patch_size_lower_bound))) for _ in range(n_dims)
+        )
+        required_mask_pixel_perc = (1 / np.prod(patch_size)) * 100
+        if expected_area_per_pixel > np.prod(patch_size):
+            raise ValueError(
+                "The probability of creating a blind-spot pixel within a patch is "
+                f"below 1, for a patch size of {patch_size} with a masked pixel "
+                f"percentage of {mask_pixel_perc}%. Either increase the patch size to "
+                f"{required_patch_size} or increase the masked pixel percentage to "
+                f"at least {required_mask_pixel_perc}%."
+            )
+
+        return self
+
     @model_validator(mode="after")
     def validate_3D(self: Self) -> Self:
         """
@@ -297,17 +344,18 @@ class Configuration(BaseModel):
         self,
         *,
         mode: Literal["json", "python"] | str = "python",
-        include: Any | None = None,
-        exclude: Any | None = None,
+        include: IncEx | None = None,
+        exclude: IncEx | None = None,
         context: Any | None = None,
-        by_alias: bool = False,
+        by_alias: bool | None = False,
         exclude_unset: bool = False,
         exclude_defaults: bool = False,
         exclude_none: bool = True,
         round_trip: bool = False,
         warnings: bool | Literal["none", "warn", "error"] = True,
+        fallback: Callable[[Any], Any] | None = None,
         serialize_as_any: bool = False,
-    ) -> dict:
+    ) -> dict[str, Any]:
         """
         Override model_dump method in order to set default values.
 
@@ -337,6 +385,8 @@ class Configuration(BaseModel):
             representation.
         warnings : bool | Literal['none', 'warn', 'error'], default=True
             Whether to emit warnings.
+        fallback : Callable[[Any], Any] | None, default=None
+            A function to call when an unknown value is encountered.
         serialize_as_any : bool, default=False
             Whether to serialize all types as Any.
 
@@ -356,6 +406,7 @@ class Configuration(BaseModel):
             exclude_none=exclude_none,
             round_trip=round_trip,
             warnings=warnings,
+            fallback=fallback,
             serialize_as_any=serialize_as_any,
         )