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

careamics/careamist.py

@@ -11,7 +11,7 @@ from pytorch_lightning.callbacks import (
     EarlyStopping,
     ModelCheckpoint,
 )
-from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
+from pytorch_lightning.loggers import CSVLogger, TensorBoardLogger, WandbLogger
 
 from careamics.config import Configuration, FCNAlgorithmConfig, load_configuration
 from careamics.config.support import (
@@ -20,7 +20,8 @@ from careamics.config.support import (
     SupportedData,
     SupportedLogger,
 )
-from careamics.dataset.dataset_utils import reshape_array
+from careamics.dataset.dataset_utils import list_files, reshape_array
+from careamics.file_io import WriteFunc, get_write_func
 from careamics.lightning import (
     FCNModule,
     HyperParametersCallback,
@@ -32,10 +33,11 @@ from careamics.lightning import (
 from careamics.model_io import export_to_bmz, load_pretrained
 from careamics.prediction_utils import convert_outputs
 from careamics.utils import check_path_exists, get_logger
+from careamics.utils.lightning_utils import read_csv_logger
 
 logger = get_logger(__name__)
 
-LOGGER_TYPES = Optional[Union[TensorBoardLogger, WandbLogger]]
+LOGGER_TYPES = list[Union[TensorBoardLogger, WandbLogger, CSVLogger]]
 
 
 class CAREamist:
@@ -144,6 +146,7 @@ class CAREamist:
 
         # path to configuration file or model
         else:
+            # TODO: update this check so models can be downloaded directly from BMZ
             source = check_path_exists(source)
 
             # configuration file
@@ -169,18 +172,29 @@ class CAREamist:
         self._define_callbacks(callbacks)
 
         # instantiate logger
+        csv_logger = CSVLogger(
+            name=self.cfg.experiment_name,
+            save_dir=self.work_dir / "csv_logs",
+        )
+
         if self.cfg.training_config.has_logger():
             if self.cfg.training_config.logger == SupportedLogger.WANDB:
-                self.experiment_logger: LOGGER_TYPES = WandbLogger(
-                    name=self.cfg.experiment_name,
-                    save_dir=self.work_dir / Path("logs"),
-                )
+                experiment_logger: LOGGER_TYPES = [
+                    WandbLogger(
+                        name=self.cfg.experiment_name,
+                        save_dir=self.work_dir / Path("wandb_logs"),
+                    ),
+                    csv_logger,
+                ]
             elif self.cfg.training_config.logger == SupportedLogger.TENSORBOARD:
-                self.experiment_logger = TensorBoardLogger(
-                    save_dir=self.work_dir / Path("logs"),
-                )
+                experiment_logger = [
+                    TensorBoardLogger(
+                        save_dir=self.work_dir / Path("tb_logs"),
+                    ),
+                    csv_logger,
+                ]
         else:
-            self.experiment_logger = None
+            experiment_logger = [csv_logger]
 
         # instantiate trainer
         self.trainer = Trainer(
@@ -194,7 +208,7 @@ class CAREamist:
             gradient_clip_algorithm=self.cfg.training_config.gradient_clip_algorithm,
             callbacks=self.callbacks,
             default_root_dir=self.work_dir,
-            logger=self.experiment_logger,
+            logger=experiment_logger,
         )
 
         # place holder for the datamodules
@@ -519,7 +533,7 @@ class CAREamist:
         *,
         batch_size: int = 1,
         tile_size: Optional[tuple[int, ...]] = None,
-        tile_overlap: tuple[int, ...] = (48, 48),
+        tile_overlap: Optional[tuple[int, ...]] = (48, 48),
         axes: Optional[str] = None,
         data_type: Optional[Literal["tiff", "custom"]] = None,
         tta_transforms: bool = False,
@@ -535,7 +549,7 @@ class CAREamist:
         *,
         batch_size: int = 1,
         tile_size: Optional[tuple[int, ...]] = None,
-        tile_overlap: tuple[int, ...] = (48, 48),
+        tile_overlap: Optional[tuple[int, ...]] = (48, 48),
         axes: Optional[str] = None,
         data_type: Optional[Literal["array"]] = None,
         tta_transforms: bool = False,
@@ -546,7 +560,7 @@ class CAREamist:
         self,
         source: Union[PredictDataModule, Path, str, NDArray],
         *,
-        batch_size: Optional[int] = None,
+        batch_size: int = 1,
         tile_size: Optional[tuple[int, ...]] = None,
         tile_overlap: Optional[tuple[int, ...]] = (48, 48),
         axes: Optional[str] = None,
@@ -567,7 +581,7 @@ class CAREamist:
         configuration parameters will be used, with the `patch_size` instead of
         `tile_size`.
 
-        Test-time augmentation (TTA) can be switched off using the `tta_transforms`
+        Test-time augmentation (TTA) can be switched on using the `tta_transforms`
         parameter. The TTA augmentation applies all possible flip and 90 degrees
         rotations to the prediction input and averages the predictions. TTA augmentation
         should not be used if you did not train with these augmentations.
@@ -580,7 +594,7 @@ class CAREamist:
 
         Parameters
         ----------
-        source : CAREamicsPredData, pathlib.Path, str or numpy.ndarray
+        source : PredictDataModule, pathlib.Path, str or numpy.ndarray
             Data to predict on.
         batch_size : int, default=1
             Batch size for prediction.
@@ -668,15 +682,195 @@ class CAREamist:
         )
         return convert_outputs(predictions, self.pred_datamodule.tiled)
 
+    def predict_to_disk(
+        self,
+        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,
+        tta_transforms: bool = False,
+        dataloader_params: Optional[dict] = None,
+        read_source_func: Optional[Callable] = 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,
+        prediction_dir: Union[Path, str] = "predictions",
+        **kwargs,
+    ) -> None:
+        """
+        Make predictions on the provided data and save outputs to files.
+
+        The predictions will be saved in a new directory 'predictions' within the set
+        working directory. The directory stucture within the 'predictions' directory
+        will match that of the source directory.
+
+        The `source` must be from files and not arrays. The file names of the
+        predictions will match those of the source. If there is more than one sample
+        within a file, the samples will be saved to seperate files. The file names of
+        samples will have the name of the corresponding source file but with the sample
+        index appended. E.g. If the the source file name is 'images.tiff' then the first
+        sample's prediction will be saved with the file name "image_0.tiff".
+        Input can be a PredictDataModule instance, a path to a data file, or a numpy
+        array.
+
+        If `data_type`, `axes` and `tile_size` are not provided, the training
+        configuration parameters will be used, with the `patch_size` instead of
+        `tile_size`.
+
+        Test-time augmentation (TTA) can be switched on using the `tta_transforms`
+        parameter. The TTA augmentation applies all possible flip and 90 degrees
+        rotations to the prediction input and averages the predictions. TTA augmentation
+        should not be used if you did not train with these augmentations.
+
+        Note that if you are using a UNet model and tiling, the tile size must be
+        divisible in every dimension by 2**d, where d is the depth of the model. This
+        avoids artefacts arising from the broken shift invariance induced by the
+        pooling layers of the UNet. If your image has less dimensions, as it may
+        happen in the Z dimension, consider padding your image.
+
+        Parameters
+        ----------
+        source : PredictDataModule or pathlib.Path, str
+            Data to predict on.
+        batch_size : int, default=1
+            Batch size for prediction.
+        tile_size : tuple of int, optional
+            Size of the tiles to use for prediction.
+        tile_overlap : tuple of int, default=(48, 48)
+            Overlap between tiles.
+        axes : str, optional
+            Axes of the input data, by default None.
+        data_type : {"array", "tiff", "custom"}, optional
+            Type of the input data.
+        tta_transforms : bool, default=True
+            Whether to apply test-time augmentation.
+        dataloader_params : dict, optional
+            Parameters to pass to the dataloader.
+        read_source_func : Callable, optional
+            Function to read the source data.
+        extension_filter : str, default=""
+            Filter for the file extension.
+        write_type : {"tiff", "custom"}, default="tiff"
+            The data type to save as, includes custom.
+        write_extension : str, optional
+            If a known `write_type` is selected this argument is ignored. For a custom
+            `write_type` an extension to save the data with must be passed.
+        write_func : WriteFunc, optional
+            If a known `write_type` is selected this argument is ignored. For a custom
+            `write_type` a function to save the data must be passed. See notes below.
+        write_func_kwargs : dict of {str: any}, optional
+            Additional keyword arguments to be passed to the save function.
+        prediction_dir : Path | str, default="predictions"
+            The path to save the prediction results to. If `prediction_dir` is not
+            absolute, the directory will be assumed to be relative to the pre-set
+            `work_dir`. If the directory does not exist it will be created.
+        **kwargs : Any
+            Unused.
+
+        Raises
+        ------
+        ValueError
+            If `write_type` is custom and `write_extension` is None.
+        ValueError
+            If `write_type` is custom and `write_fun is None.
+        ValueError
+            If `source` is not `str`, `Path` or `PredictDataModule`
+        """
+        if write_func_kwargs is None:
+            write_func_kwargs = {}
+
+        if Path(prediction_dir).is_absolute():
+            write_dir = Path(prediction_dir)
+        else:
+            write_dir = self.work_dir / prediction_dir
+        write_dir.mkdir(exist_ok=True, parents=True)
+
+        # guards for custom types
+        if write_type == SupportedData.CUSTOM:
+            if write_extension is None:
+                raise ValueError(
+                    "A `write_extension` must be provided for custom write types."
+                )
+            if write_func is None:
+                raise ValueError(
+                    "A `write_func` must be provided for custom write types."
+                )
+        else:
+            write_func = get_write_func(write_type)
+            write_extension = SupportedData.get_extension(write_type)
+
+        # extract file names
+        source_path: Union[Path, str, NDArray]
+        source_data_type: Literal["array", "tiff", "custom"]
+        if isinstance(source, PredictDataModule):
+            source_path = source.pred_data
+            source_data_type = source.data_type
+            extension_filter = source.extension_filter
+        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(
+                SupportedData(source_data_type)
+            )
+        else:
+            raise ValueError(f"Unsupported source type: '{type(source)}'.")
+
+        if source_data_type == "array":
+            raise ValueError(
+                "Predicting to disk is not supported for input type 'array'."
+            )
+        assert isinstance(source_path, (Path, str))  # because data_type != "array"
+        source_path = Path(source_path)
+
+        file_paths = list_files(source_path, source_data_type, extension_filter)
+
+        # predict and write each file in turn
+        for file_path in file_paths:
+            # source_path is relative to original source path...
+            # should mirror original directory structure
+            prediction = self.predict(
+                source=file_path,
+                batch_size=batch_size,
+                tile_size=tile_size,
+                tile_overlap=tile_overlap,
+                axes=axes,
+                data_type=data_type,
+                tta_transforms=tta_transforms,
+                dataloader_params=dataloader_params,
+                read_source_func=read_source_func,
+                extension_filter=extension_filter,
+                **kwargs,
+            )
+            # TODO: cast to float16?
+            write_data = np.concatenate(prediction)
+
+            # create directory structure and write path
+            if not source_path.is_file():
+                file_write_dir = write_dir / file_path.parent.relative_to(source_path)
+            else:
+                file_write_dir = write_dir
+            file_write_dir.mkdir(parents=True, exist_ok=True)
+            write_path = (file_write_dir / file_path.name).with_suffix(write_extension)
+
+            # write data
+            write_func(file_path=write_path, img=write_data)
+
     def export_to_bmz(
         self,
         path_to_archive: Union[Path, str],
         friendly_model_name: str,
         input_array: NDArray,
         authors: list[dict],
-        general_description: str = "",
+        general_description: str,
+        data_description: str,
+        covers: Optional[list[Union[Path, str]]] = None,
         channel_names: Optional[list[str]] = None,
-        data_description: Optional[str] = None,
+        model_version: str = "0.1.0",
     ) -> None:
         """Export the model to the BioImage Model Zoo format.
 
@@ -706,11 +900,15 @@ class CAREamist:
         authors : list of dict
             List of authors of the model.
         general_description : str
-            General description of the model, used in the metadata of the BMZ archive.
-        channel_names : list of str, optional
-            Channel names, by default None.
-        data_description : str, optional
-            Description of the data, by default None.
+            General description of the model used in the BMZ metadata.
+        data_description : str
+            Description of the data the model was trained on.
+        covers : list of pathlib.Path or str, default=None
+            Paths to the cover images.
+        channel_names : list of str, default=None
+            Channel names.
+        model_version : str, default="0.1.0"
+            Version of the model.
         """
         # TODO: add in docs that it is expected that input_array dimensions match
         # those in data_config
@@ -729,9 +927,21 @@ class CAREamist:
             path_to_archive=path_to_archive,
             model_name=friendly_model_name,
             general_description=general_description,
+            data_description=data_description,
             authors=authors,
             input_array=input_array,
             output_array=output,
+            covers=covers,
             channel_names=channel_names,
-            data_description=data_description,
+            model_version=model_version,
         )
+
+    def get_losses(self) -> dict[str, list]:
+        """Return data that can be used to plot train and validation loss curves.
+
+        Returns
+        -------
+        dict of str: list
+            Dictionary containing the losses for each epoch.
+        """
+        return read_csv_logger(self.cfg.experiment_name, self.work_dir / "csv_logs")

careamics/cli/conf.py

@@ -3,7 +3,7 @@
 import sys
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Tuple
+from typing import Optional
 
 import click
 import typer
@@ -17,6 +17,7 @@ from ..config import (
     create_n2v_configuration,
     save_configuration,
 )
+from .utils import handle_2D_3D_callback
 
 WORK_DIR = Path.cwd()
 
@@ -92,26 +93,6 @@ def conf_options(  # numpydoc ignore=PR01
     ctx.obj = ConfOptions(dir, name, force, print)
 
 
-def patch_size_callback(value: Tuple[int, int, int]) -> Tuple[int, ...]:
-    """
-    Callback for --patch-size option.
-
-    Parameters
-    ----------
-    value : (int, int, int)
-        Patch size value.
-
-    Returns
-    -------
-    (int, int, int) | (int, int)
-        If the last element in `value` is -1 the tuple is reduced to the first two
-        values.
-    """
-    if value[2] == -1:
-        return value[:2]
-    return value
-
-
 # TODO: Need to decide how to parse model kwargs
 #   - Could be json style string to be loaded as dict e.g. {"depth": 3}
 #        - Cons: Annoying to type, easily have syntax errors
@@ -132,7 +113,7 @@ def care(  # numpydoc ignore=PR01
                 "is not 3D pass the last value as -1 e.g. --patch-size 64 64 -1)."
             ),
             click_type=click.Tuple([int, int, int]),
-            callback=patch_size_callback,
+            callback=handle_2D_3D_callback,
         ),
     ],
     batch_size: Annotated[int, typer.Option(help="Batch size.")],
@@ -154,8 +135,12 @@ def care(  # numpydoc ignore=PR01
             help="Loss function to use.",
         ),
     ] = "mae",
-    n_channels_in: Annotated[int, typer.Option(help="Number of channels in")] = 1,
-    n_channels_out: Annotated[int, typer.Option(help="Number of channels out")] = -1,
+    n_channels_in: Annotated[
+        Optional[int], typer.Option(help="Number of channels in")
+    ] = None,
+    n_channels_out: Annotated[
+        Optional[int], typer.Option(help="Number of channels out")
+    ] = None,
     logger: Annotated[
         click.Choice,
         typer.Option(
@@ -215,7 +200,7 @@ def n2n(  # numpydoc ignore=PR01
                 "is not 3D pass the last value as -1 e.g. --patch-size 64 64 -1)."
             ),
             click_type=click.Tuple([int, int, int]),
-            callback=patch_size_callback,
+            callback=handle_2D_3D_callback,
         ),
     ],
     batch_size: Annotated[int, typer.Option(help="Batch size.")],
@@ -237,8 +222,12 @@ def n2n(  # numpydoc ignore=PR01
             help="Loss function to use.",
         ),
     ] = "mae",
-    n_channels_in: Annotated[int, typer.Option(help="Number of channels in")] = 1,
-    n_channels_out: Annotated[int, typer.Option(help="Number of channels out")] = -1,
+    n_channels_in: Annotated[
+        Optional[int], typer.Option(help="Number of channels in")
+    ] = None,
+    n_channels_out: Annotated[
+        Optional[int], typer.Option(help="Number of channels out")
+    ] = None,
     logger: Annotated[
         click.Choice,
         typer.Option(
@@ -295,7 +284,7 @@ def n2v(  # numpydoc ignore=PR01
                 "is not 3D pass the last value as -1 e.g. --patch-size 64 64 -1)."
             ),
             click_type=click.Tuple([int, int, int]),
-            callback=patch_size_callback,
+            callback=handle_2D_3D_callback,
         ),
     ],
     batch_size: Annotated[int, typer.Option(help="Batch size.")],
@@ -312,8 +301,8 @@ def n2v(  # numpydoc ignore=PR01
     ] = True,
     use_n2v2: Annotated[bool, typer.Option(help="Whether to use N2V2")] = False,
     n_channels: Annotated[
-        int, typer.Option(help="Number of channels (in and out)")
-    ] = 1,
+        Optional[int], 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[
         float, typer.Option(help="Percentage of pixels masked in each patch.")

careamics/cli/main.py

@@ -9,21 +9,20 @@ its implementation is contained in the conf.py file.
 from pathlib import Path
 from typing import Optional
 
+import click
 import typer
 from typing_extensions import Annotated
 
 from ..careamist import CAREamist
 from . import conf
+from .utils import handle_2D_3D_callback
 
 app = typer.Typer(
     help="Run CAREamics algorithms from the command line, including Noise2Void "
-    "and its many variants and cousins"
-)
-app.add_typer(
-    conf.app,
-    name="conf",
-    # callback=conf.conf_options
+    "and its many variants and cousins",
+    pretty_exceptions_show_locals=False,
 )
+app.add_typer(conf.app, name="conf")
 
 
 @app.command()
@@ -102,7 +101,7 @@ def train(  # numpydoc ignore=PR01
         typer.Option(
             "--work-dir",
             "-wd",
-            help=("Path to working directory in which to save checkpoints and " "logs"),
+            help=("Path to working directory in which to save checkpoints and logs"),
             exists=True,
             file_okay=False,
             dir_okay=True,
@@ -123,10 +122,112 @@ def train(  # numpydoc ignore=PR01
 
 
 @app.command()
-def predict():  # numpydoc ignore=PR01
+def predict(  # numpydoc ignore=PR01
+    model: Annotated[
+        Path,
+        typer.Argument(
+            help="Path to a configuration file or a trained model.",
+            exists=True,
+            file_okay=True,
+            dir_okay=False,
+        ),
+    ],
+    source: Annotated[
+        Path,
+        typer.Argument(
+            help="Path to the training data. Can be a directory or single file.",
+            exists=True,
+            file_okay=True,
+            dir_okay=True,
+        ),
+    ],
+    batch_size: Annotated[int, typer.Option(help="Batch size.")] = 1,
+    tile_size: Annotated[
+        Optional[click.Tuple],
+        typer.Option(
+            help=(
+                "Size of the tiles to use for prediction, (if the data "
+                "is not 3D pass the last value as -1 e.g. --tile_size 64 64 -1)."
+            ),
+            click_type=click.Tuple([int, int, int]),
+            callback=handle_2D_3D_callback,
+        ),
+    ] = None,
+    tile_overlap: Annotated[
+        click.Tuple,
+        typer.Option(
+            help=(
+                "Overlap between tiles, (if the data is not 3D pass the last value as "
+                "-1 e.g. --tile_overlap 64 64 -1)."
+            ),
+            click_type=click.Tuple([int, int, int]),
+            callback=handle_2D_3D_callback,
+        ),
+    ] = (48, 48, -1),
+    axes: Annotated[
+        Optional[str],
+        typer.Option(
+            help="Axes of the input data. If unused the data is assumed to have the "
+            "same axes as the original training data."
+        ),
+    ] = None,
+    data_type: Annotated[
+        click.Choice,
+        typer.Option(click_type=click.Choice(["tiff"]), help="Type of the input data."),
+    ] = "tiff",
+    tta_transforms: Annotated[
+        bool,
+        typer.Option(
+            "--tta-transforms/--no-tta-transforms",
+            "-t/-T",
+            help="Whether to apply test-time augmentation.",
+        ),
+    ] = False,
+    write_type: Annotated[
+        click.Choice,
+        typer.Option(
+            click_type=click.Choice(["tiff"]), help="Type of the output data."
+        ),
+    ] = "tiff",
+    # TODO: could make dataloader_params as json, necessary?
+    work_dir: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--work-dir",
+            "-wd",
+            help=("Path to working directory."),
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+        ),
+    ] = None,
+    prediction_dir: Annotated[
+        Path,
+        typer.Option(
+            "--prediction-dir",
+            "-pd",
+            help=(
+                "Directory to save predictions to. If not an abosulte path it will be "
+                "relative to the set working directory."
+            ),
+            file_okay=False,
+            dir_okay=True,
+        ),
+    ] = Path("predictions"),
+):
     """Create and save predictions from CAREamics models."""
-    # TODO: Need a save predict to workdir function
-    raise NotImplementedError
+    engine = CAREamist(source=model, work_dir=work_dir)
+    engine.predict_to_disk(
+        source=source,
+        batch_size=batch_size,
+        tile_size=tile_size,
+        tile_overlap=tile_overlap,
+        axes=axes,
+        data_type=data_type,
+        tta_transforms=tta_transforms,
+        write_type=write_type,
+        prediction_dir=prediction_dir,
+    )
 
 
 def run():

careamics/cli/utils.py

@@ -0,0 +1,29 @@
+"""Utility functions for the CAREamics CLI."""
+
+from typing import Optional, Tuple
+
+
+def handle_2D_3D_callback(
+    value: Optional[Tuple[int, int, int]]
+) -> Optional[Tuple[int, ...]]:
+    """
+    Callback for options that require 2D or 3D inputs.
+
+    In the case of 2D, the 3rd element should be set to -1.
+
+    Parameters
+    ----------
+    value : (int, int, int)
+        Tile size value.
+
+    Returns
+    -------
+    (int, int, int) | (int, int)
+        If the last element in `value` is -1 the tuple is reduced to the first two
+        values.
+    """
+    if value is None:
+        return value
+    if value[2] == -1:
+        return value[:2]
+    return value

careamics/config/__init__.py

@@ -18,6 +18,7 @@ __all__ = [
     "clear_custom_models",
     "GaussianMixtureNMConfig",
     "MultiChannelNMConfig",
+    "LVAELossConfig",
 ]
 from .architectures import CustomModel, clear_custom_models, register_model
 from .callback_model import CheckpointModel
@@ -34,6 +35,7 @@ from .configuration_model import (
 from .data_model import DataConfig
 from .fcn_algorithm_model import FCNAlgorithmConfig
 from .inference_model import InferenceConfig
+from .loss_model import LVAELossConfig
 from .nm_model import GaussianMixtureNMConfig, MultiChannelNMConfig
 from .training_model import TrainingConfig
 from .vae_algorithm_model import VAEAlgorithmConfig