Attempt at fixing #396

.github/workflows/codecov.yml

@@ -2,17 +2,14 @@ name: codecov
 
 on:
   push:
-    branches:
-      - main
     tags:
       - "v*"
-  pull_request:
-  workflow_dispatch:
 
 jobs:
-  test:
+  codecov:
     name: ${{ matrix.platform }} (${{ matrix.python-version }})
     runs-on: ${{ matrix.platform }}
+    if: startsWith(github.ref, 'refs/tags/') && github.event_name != 'schedule'
     strategy:
       fail-fast: false
       matrix:
@@ -38,7 +35,7 @@ jobs:
         run: |
           python -m pip install -U pip
           # if running a cron job, we add the --pre flag to test against pre-releases
-          python -m pip install .[dev] ${{ github.event_name == 'schedule' && '--pre' || ''  }}
+          python -m pip install .[dev,czi] ${{ github.event_name == 'schedule' && '--pre' || ''  }}
 
       - name: 🧪 Run Tests
         run: pytest --color=yes --cov --cov-config=pyproject.toml --cov-report=xml --cov-report=term-missing 

.pre-commit-config.yaml

@@ -14,7 +14,7 @@ repos:
       - id: validate-pyproject
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.11.12
+    rev: v0.12.2
     hooks:
       - id: ruff
         exclude: "^src/careamics/lvae_training/.*|^src/careamics/models/lvae/.*|^scripts/.*"
@@ -26,7 +26,7 @@ repos:
       - id: black
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.16.0
+    rev: v1.16.1
     hooks:
       - id: mypy
         files: "^src/"
@@ -40,7 +40,7 @@ repos:
 
   # check docstrings
   - repo: https://github.com/numpy/numpydoc
-    rev: v1.8.0
+    rev: v1.9.0
     hooks:
       - id: numpydoc-validation
         exclude: "^src/careamics/dataset_ng/.*|^src/careamics/lvae_training/.*|^src/careamics/models/lvae/.*|^src/careamics/losses/lvae/.*|^src/careamics/lightning/dataset_ng/data_module.py|^scripts/.*"

pyproject.toml

@@ -11,8 +11,7 @@ build-backend = "hatchling.build"
 source = "vcs"
 
 [tool.hatch.build.targets.wheel]
-only-include = ["src"]
-sources = ["src"]
+packages = ["src/careamics"]
 
 # https://peps.python.org/pep-0621/
 [project]
@@ -44,7 +43,7 @@ dependencies = [
     'numpy<2.0.0',
     'torch>=2.0,<=2.7.1',
     'torchvision<=0.22.1',
-    'bioimageio.core==0.7',
+    'bioimageio.core==0.9.0',
     'tifffile<=2025.5.10',
     'psutil<=7.0.0',
     'pydantic>=2.11,<=2.12',
@@ -55,7 +54,6 @@ dependencies = [
     'zarr<3.0.0',
     'pillow<=11.2.1',
     'matplotlib<=3.10.3',
-    'xarray<2025.3.0',
 ]
 
 [project.optional-dependencies]

src/careamics/careamist.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Literal, Optional, Union, overload
+from typing import Any, Literal, Union, overload
 
 import numpy as np
 from numpy.typing import NDArray
@@ -34,7 +34,6 @@
     PredictDataModule,
     ProgressBarCallback,
     TrainDataModule,
-    VAEModule,
     create_predict_datamodule,
 )
 from careamics.model_io import export_to_bmz, load_pretrained
@@ -85,25 +84,25 @@ class CAREamist:
     def __init__(  # numpydoc ignore=GL08
         self,
         source: Union[Path, str],
-        work_dir: Optional[Union[Path, str]] = None,
-        callbacks: Optional[list[Callback]] = None,
+        work_dir: Union[Path, str] | None = None,
+        callbacks: list[Callback] | None = None,
         enable_progress_bar: bool = True,
     ) -> None: ...
 
     @overload
     def __init__(  # numpydoc ignore=GL08
         self,
         source: Configuration,
-        work_dir: Optional[Union[Path, str]] = None,
-        callbacks: Optional[list[Callback]] = None,
+        work_dir: Union[Path, str] | None = None,
+        callbacks: list[Callback] | None = None,
         enable_progress_bar: bool = True,
     ) -> None: ...
 
     def __init__(
         self,
         source: Union[Path, str, Configuration],
-        work_dir: Optional[Union[Path, str]] = None,
-        callbacks: Optional[list[Callback]] = None,
+        work_dir: Union[Path, str] | None = None,
+        callbacks: list[Callback] | None = None,
         enable_progress_bar: bool = True,
     ) -> None:
         """
@@ -158,10 +157,13 @@ def __init__(
                     algorithm_config=self.cfg.algorithm_config,
                 )
             elif isinstance(self.cfg.algorithm_config, VAEBasedAlgorithm):
-                self.model = VAEModule(
-                    algorithm_config=self.cfg.algorithm_config,
+                # self.model = VAEModule(
+                #     algorithm_config=self.cfg.algorithm_config,
+                # )
+                raise NotImplementedError(
+                    "VAE based algorithms are not yet available with the CAREamist API."
+                    " Use the Lightning API for these algorithms (see documentation)."
                 )
-                # raise NotImplementedError("VAE based algorithms are not implemented.")
             else:
                 raise NotImplementedError("Architecture not supported.")
 
@@ -233,11 +235,11 @@ def __init__(
         )
 
         # place holder for the datamodules
-        self.train_datamodule: Optional[TrainDataModule] = None
-        self.pred_datamodule: Optional[PredictDataModule] = None
+        self.train_datamodule: TrainDataModule | None = None
+        self.pred_datamodule: PredictDataModule | None = None
 
     def _define_callbacks(
-        self, callbacks: Optional[list[Callback]], enable_progress_bar: bool
+        self, callbacks: list[Callback] | None, enable_progress_bar: bool
     ) -> None:
         """Define the callbacks for the training loop.
 
@@ -299,11 +301,11 @@ def stop_training(self) -> None:
     def train(
         self,
         *,
-        datamodule: Optional[TrainDataModule] = None,
-        train_source: Optional[Union[Path, str, NDArray]] = None,
-        val_source: Optional[Union[Path, str, NDArray]] = None,
-        train_target: Optional[Union[Path, str, NDArray]] = None,
-        val_target: Optional[Union[Path, str, NDArray]] = None,
+        datamodule: TrainDataModule | None = None,
+        train_source: Union[Path, str, NDArray] | None = None,
+        val_source: Union[Path, str, NDArray] | None = None,
+        train_target: Union[Path, str, NDArray] | None = None,
+        val_target: Union[Path, str, NDArray] | None = None,
         use_in_memory: bool = True,
         val_percentage: float = 0.1,
         val_minimum_split: int = 1,
@@ -454,9 +456,9 @@ def _train_on_datamodule(self, datamodule: TrainDataModule) -> None:
     def _train_on_array(
         self,
         train_data: NDArray,
-        val_data: Optional[NDArray] = None,
-        train_target: Optional[NDArray] = None,
-        val_target: Optional[NDArray] = None,
+        val_data: NDArray | None = None,
+        train_target: NDArray | None = None,
+        val_target: NDArray | None = None,
         val_percentage: float = 0.1,
         val_minimum_split: int = 5,
     ) -> None:
@@ -495,9 +497,9 @@ def _train_on_array(
     def _train_on_path(
         self,
         path_to_train_data: Union[Path, str],
-        path_to_val_data: Optional[Union[Path, str]] = None,
-        path_to_train_target: Optional[Union[Path, str]] = None,
-        path_to_val_target: Optional[Union[Path, str]] = None,
+        path_to_val_data: Union[Path, str] | None = None,
+        path_to_train_target: Union[Path, str] | None = None,
+        path_to_val_target: Union[Path, str] | None = None,
         use_in_memory: bool = True,
         val_percentage: float = 0.1,
         val_minimum_split: int = 1,
@@ -560,13 +562,13 @@ def predict(  # numpydoc ignore=GL08
         source: Union[Path, str],
         *,
         batch_size: int = 1,
-        tile_size: Optional[tuple[int, ...]] = None,
-        tile_overlap: Optional[tuple[int, ...]] = (48, 48),
-        axes: Optional[str] = None,
-        data_type: Optional[Literal["tiff", "custom"]] = None,
+        tile_size: tuple[int, ...] | None = None,
+        tile_overlap: tuple[int, ...] | None = (48, 48),
+        axes: str | None = None,
+        data_type: Literal["tiff", "custom"] | None = None,
         tta_transforms: bool = False,
-        dataloader_params: Optional[dict] = None,
-        read_source_func: Optional[Callable] = None,
+        dataloader_params: dict | None = None,
+        read_source_func: Callable | None = None,
         extension_filter: str = "",
     ) -> Union[list[NDArray], NDArray]: ...
 
@@ -576,26 +578,26 @@ def predict(  # numpydoc ignore=GL08
         source: NDArray,
         *,
         batch_size: int = 1,
-        tile_size: Optional[tuple[int, ...]] = None,
-        tile_overlap: Optional[tuple[int, ...]] = (48, 48),
-        axes: Optional[str] = None,
-        data_type: Optional[Literal["array"]] = None,
+        tile_size: tuple[int, ...] | None = None,
+        tile_overlap: tuple[int, ...] | None = (48, 48),
+        axes: str | None = None,
+        data_type: Literal["array"] | None = None,
         tta_transforms: bool = False,
-        dataloader_params: Optional[dict] = None,
+        dataloader_params: dict | None = None,
     ) -> Union[list[NDArray], NDArray]: ...
 
     def predict(
         self,
         source: Union[PredictDataModule, Path, str, NDArray],
         *,
         batch_size: int = 1,
-        tile_size: Optional[tuple[int, ...]] = None,
-        tile_overlap: Optional[tuple[int, ...]] = (48, 48),
-        axes: Optional[str] = None,
-        data_type: Optional[Literal["array", "tiff", "custom"]] = None,
+        tile_size: tuple[int, ...] | None = None,
+        tile_overlap: tuple[int, ...] | None = (48, 48),
+        axes: str | None = None,
+        data_type: Literal["array", "tiff", "custom"] | None = None,
         tta_transforms: bool = False,
-        dataloader_params: Optional[dict] = None,
-        read_source_func: Optional[Callable] = None,
+        dataloader_params: dict | None = None,
+        read_source_func: Callable | None = None,
         extension_filter: str = "",
         **kwargs: Any,
     ) -> Union[list[NDArray], NDArray]:
@@ -664,7 +666,8 @@ def predict(
             or self.cfg.data_config.image_stds is None
         ):
             raise ValueError(
-                "Mean and std must be provided in the configuration. \n Likely reason is that model was not trained"
+                "Mean and std must be provided in the configuration. \n Has the model "
+                "is that the model was not trained."
             )
 
         # tile size for UNets
@@ -717,18 +720,18 @@ def predict_to_disk(
         source: Union[PredictDataModule, Path, str],
         *,
         batch_size: int = 1,
-        tile_size: Optional[tuple[int, ...]] = None,
-        tile_overlap: Optional[tuple[int, ...]] = (48, 48),
-        axes: Optional[str] = None,
-        data_type: Optional[Literal["tiff", "custom"]] = None,
+        tile_size: tuple[int, ...] | None = None,
+        tile_overlap: tuple[int, ...] | None = (48, 48),
+        axes: str | None = None,
+        data_type: Literal["tiff", "custom"] | None = None,
         tta_transforms: bool = False,
-        dataloader_params: Optional[dict] = None,
-        read_source_func: Optional[Callable] = None,
+        dataloader_params: dict | None = None,
+        read_source_func: Callable | None = None,
         extension_filter: str = "",
         write_type: Literal["tiff", "custom"] = "tiff",
-        write_extension: Optional[str] = None,
-        write_func: Optional[WriteFunc] = None,
-        write_func_kwargs: Optional[dict[str, Any]] = None,
+        write_extension: str | None = None,
+        write_func: WriteFunc | None = None,
+        write_func_kwargs: dict[str, Any] | None = None,
         prediction_dir: Union[Path, str] = "predictions",
         **kwargs,
     ) -> None:
@@ -898,8 +901,8 @@ def export_to_bmz(
         authors: list[dict],
         general_description: str,
         data_description: str,
-        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 the BioImage Model Zoo format.

src/careamics/cli/conf.py

@@ -3,7 +3,7 @@
 import sys
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Annotated, Optional
+from typing import Annotated
 
 import click
 import typer
@@ -135,10 +135,10 @@ def care(  # numpydoc ignore=PR01
         ),
     ] = "mae",
     n_channels_in: Annotated[
-        Optional[int], typer.Option(help="Number of channels in")
+        int | None, typer.Option(help="Number of channels in")
     ] = None,
     n_channels_out: Annotated[
-        Optional[int], typer.Option(help="Number of channels out")
+        int | None, typer.Option(help="Number of channels out")
     ] = None,
     logger: Annotated[
         click.Choice,
@@ -222,10 +222,10 @@ def n2n(  # numpydoc ignore=PR01
         ),
     ] = "mae",
     n_channels_in: Annotated[
-        Optional[int], typer.Option(help="Number of channels in")
+        int | None, typer.Option(help="Number of channels in")
     ] = None,
     n_channels_out: Annotated[
-        Optional[int], typer.Option(help="Number of channels out")
+        int | None, typer.Option(help="Number of channels out")
     ] = None,
     logger: Annotated[
         click.Choice,
@@ -300,7 +300,7 @@ def n2v(  # numpydoc ignore=PR01
     ] = True,
     use_n2v2: Annotated[bool, typer.Option(help="Whether to use N2V2")] = False,
     n_channels: Annotated[
-        Optional[int], typer.Option(help="Number of channels (in and out)")
+        int | None, typer.Option(help="Number of channels (in and out)")
     ] = None,
     roi_size: Annotated[int, typer.Option(help="N2V pixel manipulation area.")] = 11,
     masked_pixel_percentage: Annotated[