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

careamics/dataset_ng/demos/mouse_nuclei_demo.ipynb

@@ -0,0 +1,292 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "from careamics_portfolio import PortfolioManager\n",
+    "\n",
+    "from careamics.config.configuration_factories import create_n2v_configuration\n",
+    "from careamics.config.support import SupportedTransform\n",
+    "from careamics.lightning.callbacks import HyperParametersCallback\n",
+    "from careamics.lightning.dataset_ng.data_module import CareamicsDataModule\n",
+    "from careamics.lightning.dataset_ng.lightning_modules import N2VModule"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "portfolio = PortfolioManager()\n",
+    "files = portfolio.denoiseg.MouseNuclei_n20.download()\n",
+    "files.sort()\n",
+    "\n",
+    "# load images\n",
+    "train_data = np.load(files[1])[\"X_train\"]\n",
+    "print(f\"Train data shape: {train_data.shape}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "indices = [34, 293, 571, 783]\n",
+    "\n",
+    "fig, ax = plt.subplots(2, 2, figsize=(8, 8))\n",
+    "ax[0, 0].imshow(train_data[indices[0]], cmap=\"gray\")\n",
+    "ax[0, 0].set_title(f\"Image {indices[0]}\")\n",
+    "ax[0, 0].set_xticks([])\n",
+    "ax[0, 0].set_yticks([])\n",
+    "\n",
+    "ax[0, 1].imshow(train_data[indices[1]], cmap=\"gray\")\n",
+    "ax[0, 1].set_title(f\"Image {indices[1]}\")\n",
+    "ax[0, 1].set_xticks([])\n",
+    "ax[0, 1].set_yticks([])\n",
+    "\n",
+    "ax[1, 0].imshow(train_data[indices[2]], cmap=\"gray\")\n",
+    "ax[1, 0].set_title(f\"Image {indices[2]}\")\n",
+    "ax[1, 0].set_xticks([])\n",
+    "ax[1, 0].set_yticks([])\n",
+    "\n",
+    "ax[1, 1].imshow(train_data[indices[3]], cmap=\"gray\")\n",
+    "ax[1, 1].set_title(f\"Image {indices[3]}\")\n",
+    "ax[1, 1].set_xticks([])\n",
+    "ax[1, 1].set_yticks([])\n",
+    "\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "config = create_n2v_configuration(\n",
+    "    experiment_name=\"mouse_nuclei_n2v\",\n",
+    "    data_type=\"array\",\n",
+    "    axes=\"SYX\",\n",
+    "    patch_size=(64, 64),\n",
+    "    batch_size=16,\n",
+    "    num_epochs=10,\n",
+    ")\n",
+    "\n",
+    "print(config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Ensuring that transforms are set\n",
+    "config.data_config.transforms =[\n",
+    "    {\n",
+    "        \"name\": SupportedTransform.XY_FLIP.value,\n",
+    "        \"flip_x\": True,\n",
+    "        \"flip_y\": True,\n",
+    "    },\n",
+    "    {\n",
+    "        \"name\": SupportedTransform.XY_RANDOM_ROTATE90.value,\n",
+    "    },\n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sklearn.model_selection import train_test_split\n",
+    "\n",
+    "train_data, val_data = train_test_split(train_data, test_size=0.1, random_state=42)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_data_module = CareamicsDataModule(\n",
+    "    data_config=config.data_config,\n",
+    "    train_data=train_data,\n",
+    "    val_data=val_data,\n",
+    ")\n",
+    "\n",
+    "model = N2VModule(config.algorithm_config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pytorch_lightning import Trainer\n",
+    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
+    "from pytorch_lightning.loggers import WandbLogger\n",
+    "\n",
+    "root = Path(\"nuclei_n2v\")\n",
+    "callbacks = [\n",
+    "    ModelCheckpoint(\n",
+    "        dirpath=root / \"checkpoints\",\n",
+    "        filename=\"nuclei_new_lightning_module\",\n",
+    "        save_last=True,\n",
+    "        monitor=\"val_loss\",\n",
+    "        mode=\"min\",\n",
+    "    ),\n",
+    "    HyperParametersCallback(config)\n",
+    "]\n",
+    "logger = WandbLogger(\n",
+    "    project=\"nuclei-n2v\", name=\"nuclei_new_lightning_module\"\n",
+    ")\n",
+    "\n",
+    "trainer = Trainer(\n",
+    "    max_epochs=10,\n",
+    "    default_root_dir=root,\n",
+    "    callbacks=callbacks,\n",
+    "    logger=logger\n",
+    ")\n",
+    "trainer.fit(model, datamodule=train_data_module)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from careamics.config.inference_model import InferenceConfig\n",
+    "from careamics.dataset_ng.legacy_interoperability import imageregions_to_tileinfos\n",
+    "from careamics.prediction_utils import convert_outputs\n",
+    "\n",
+    "train_data = np.load(files[1])[\"X_train\"]\n",
+    "\n",
+    "config = InferenceConfig(\n",
+    "    model_config=config,\n",
+    "    data_type=\"array\",\n",
+    "    tile_size=(64, 64),\n",
+    "    tile_overlap=(32, 32),\n",
+    "    axes=\"SYX\",\n",
+    "    batch_size=1,\n",
+    "    image_means=train_data_module.train_dataset.input_stats.means,\n",
+    "    image_stds=train_data_module.train_dataset.input_stats.stds\n",
+    ")\n",
+    "\n",
+    "inf_data_module = CareamicsDataModule(\n",
+    "    data_config=config,\n",
+    "    pred_data=train_data\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "predictions = trainer.predict(model, datamodule=inf_data_module)\n",
+    "tile_infos = imageregions_to_tileinfos(predictions)\n",
+    "predictions = convert_outputs(tile_infos, tiled=True)\n",
+    "predictions = np.stack(predictions).squeeze()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "files = portfolio.denoiseg.MouseNuclei_n0.download()\n",
+    "files.sort()\n",
+    "\n",
+    "gt_data = np.load(files[1])[\"X_train\"]\n",
+    "print(f\"GT data shape: {gt_data.shape}\")\n",
+    "print(f\"Predictions shape: {predictions.shape}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from careamics.utils.metrics import scale_invariant_psnr\n",
+    "\n",
+    "indices = [389, 621]\n",
+    "\n",
+    "for i in indices:\n",
+    "    # compute psnr\n",
+    "    psnr_noisy = scale_invariant_psnr(gt_data[i], train_data[i])\n",
+    "    psnr_denoised = scale_invariant_psnr(gt_data[i], predictions[i].squeeze())\n",
+    "\n",
+    "    # plot images\n",
+    "    fig, ax = plt.subplots(1, 3, figsize=(10, 10))\n",
+    "    ax[0].imshow(train_data[i], cmap=\"gray\")\n",
+    "    ax[0].set_title(f\"Noisy Image\\nPSNR: {psnr_noisy:.2f}\")\n",
+    "    ax[0].set_xticks([])\n",
+    "    ax[0].set_yticks([])\n",
+    "\n",
+    "    ax[1].imshow(predictions[i].squeeze(), cmap=\"gray\")\n",
+    "    ax[1].set_title(f\"Denoised Image\\nPSNR: {psnr_denoised:.2f}\")\n",
+    "    ax[1].set_xticks([])\n",
+    "    ax[1].set_yticks([])\n",
+    "\n",
+    "    ax[2].imshow(gt_data[i], cmap=\"gray\")\n",
+    "    ax[2].set_title(\"GT Image\")\n",
+    "    ax[2].set_xticks([])\n",
+    "    ax[2].set_yticks([])\n",
+    "\n",
+    "    plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "psnrs = np.zeros(gt_data.shape[0])\n",
+    "\n",
+    "for i in range(gt_data.shape[0]):\n",
+    "    psnrs[i] = scale_invariant_psnr(gt_data[i], predictions[i].squeeze())\n",
+    "\n",
+    "print(f\"PSNR: {np.mean(psnrs):.2f} ± {np.std(psnrs):.2f}\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.20"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

careamics/dataset_ng/factory.py

@@ -0,0 +1,408 @@
+from collections.abc import Sequence
+from enum import Enum
+from pathlib import Path
+from typing import Any, Optional, Union
+
+from numpy.typing import NDArray
+from typing_extensions import ParamSpec
+
+from careamics.config import DataConfig, InferenceConfig
+from careamics.config.support import SupportedData
+from careamics.dataset_ng.patch_extractor import ImageStackLoader, PatchExtractor
+from careamics.dataset_ng.patch_extractor.image_stack import (
+    GenericImageStack,
+    ImageStack,
+    InMemoryImageStack,
+    ZarrImageStack,
+)
+from careamics.dataset_ng.patch_extractor.patch_extractor_factory import (
+    create_array_extractor,
+    create_custom_file_extractor,
+    create_custom_image_stack_extractor,
+    create_ome_zarr_extractor,
+    create_tiff_extractor,
+)
+from careamics.file_io.read import ReadFunc
+
+from .dataset import CareamicsDataset, Mode
+
+P = ParamSpec("P")
+
+
+# Enum class used to determine which loading functions should be used
+class DatasetType(Enum):
+    """Labels for the dataset based on the underlying data and how it is loaded."""
+
+    ARRAY = "array"
+    IN_MEM_TIFF = "in_mem_tiff"
+    LAZY_TIFF = "lazy_tiff"
+    IN_MEM_CUSTOM_FILE = "in_mem_custom_file"
+    OME_ZARR = "ome_zarr"
+    CUSTOM_IMAGE_STACK = "custom_image_stack"
+
+
+# bit of a mess of if-else statements
+def determine_dataset_type(
+    data_type: SupportedData,
+    in_memory: bool,
+    read_func: Optional[ReadFunc] = None,
+    image_stack_loader: Optional[ImageStackLoader] = None,
+) -> DatasetType:
+    """Determine what the dataset type should be based on the input arguments.
+
+    Parameters
+    ----------
+    data_type : SupportedData
+        The underlying datatype.
+    in_memory : bool
+        Whether all the data should be loaded into memory. This is argument is ignored
+        unless the `data_type` is "tiff" or "custom".
+    read_func : ReadFunc, optional
+        A function that can that can be used to load custom data. This argument is
+        ignored unless the `data_type` is "custom".
+    image_stack_loader : ImageStackLoader, optional
+        A function for custom image stack loading. This argument is ignored unless the
+        `data_type` is "custom".
+
+    Returns
+    -------
+    DatasetType
+        The Dataset type.
+
+    Raises
+    ------
+    NotImplementedError
+        For lazy-loading (`in_memory=False`) of a custom file type.
+    ValueError
+        If the `data_type` is "custom" but both `read_func` and `image_stack_loader` are
+        None.
+    ValueError
+        If the `data_type` is unrecognized.
+    """
+    if data_type == SupportedData.ARRAY:
+        # TODO: ignoring in_memory arg, error if False?
+        return DatasetType.ARRAY
+    elif data_type == SupportedData.TIFF:
+        if in_memory:
+            return DatasetType.IN_MEM_TIFF
+        else:
+            return DatasetType.LAZY_TIFF
+    elif data_type == SupportedData.CUSTOM:
+        if read_func is not None:
+            if in_memory:
+                return DatasetType.IN_MEM_CUSTOM_FILE
+            else:
+                raise NotImplementedError(
+                    "Lazy loading has not been implemented for custom file types yet."
+                )
+        elif image_stack_loader is not None:
+            # TODO: ignoring im_memory arg
+            return DatasetType.CUSTOM_IMAGE_STACK
+        else:
+            raise ValueError(
+                "Found `data_type='custom'` but no `read_func` or `image_stack_loader` "
+                "has been provided."
+            )
+    # TODO: ZARR
+    else:
+        raise ValueError(f"Unrecognized `data_type`, '{data_type}'.")
+
+
+# convenience function but should use `create_dataloader` function instead
+# For lazy loading custom batch sampler also needs to be set.
+def create_dataset(
+    config: Union[DataConfig, InferenceConfig],
+    mode: Mode,
+    inputs: Any,
+    targets: Any,
+    in_memory: bool,
+    read_func: Optional[ReadFunc] = None,
+    read_kwargs: Optional[dict[str, Any]] = None,
+    image_stack_loader: Optional[ImageStackLoader] = None,
+    image_stack_loader_kwargs: Optional[dict[str, Any]] = None,
+) -> CareamicsDataset[ImageStack]:
+    """
+    Convenience function to create the CAREamicsDataset.
+
+    Parameters
+    ----------
+    config : DataConfig or InferenceConfig
+        The data configuration.
+    mode : Mode
+        Whether to create the dataset in "training", "validation" or "predicting" mode.
+    inputs : Any
+        The input sources to the dataset.
+    targets : Any, optional
+        The target sources to the dataset.
+    in_memory : bool
+        Whether all the data should be loaded into memory. This is argument is ignored
+        unless the `data_type` in `config` is "tiff" or "custom".
+    read_func : ReadFunc, optional
+        A function that can that can be used to load custom data. This argument is
+        ignored unless the `data_type` in the `config` is "custom".
+    read_kwargs : dict of {str, Any}, optional
+        Additional key-word arguments to pass to the `read_func`.
+    image_stack_loader : ImageStackLoader, optional
+        A function for custom image stack loading. This argument is ignored unless the
+        `data_type` in the `config` is "custom".
+    image_stack_loader_kwargs : {str, Any}, optional
+        Additional key-word arguments to pass to the `image_stack_loader`.
+
+    Returns
+    -------
+    CareamicsDataset[ImageStack]
+        The CAREamicsDataset
+
+    Raises
+    ------
+    ValueError
+        For an unrecognized `data_type` in the `config`.
+    """
+    data_type = SupportedData(config.data_type)
+    dataset_type = determine_dataset_type(
+        data_type, in_memory, read_func, image_stack_loader
+    )
+    if dataset_type == DatasetType.ARRAY:
+        return create_array_dataset(config, mode, inputs, targets)
+    elif dataset_type == DatasetType.IN_MEM_TIFF:
+        return create_tiff_dataset(config, mode, inputs, targets)
+    # TODO: Lazy tiff
+    elif dataset_type == DatasetType.IN_MEM_CUSTOM_FILE:
+        if read_kwargs is None:
+            read_kwargs = {}
+        assert read_func is not None  # should be true from `determine_dataset_type`
+        return create_custom_file_dataset(
+            config, mode, inputs, targets, read_func=read_func, read_kwargs=read_kwargs
+        )
+    elif dataset_type == DatasetType.CUSTOM_IMAGE_STACK:
+        if image_stack_loader_kwargs is None:
+            image_stack_loader_kwargs = {}
+        assert image_stack_loader is not None  # should be true
+        return create_custom_image_stack_dataset(
+            config,
+            mode,
+            inputs,
+            targets,
+            image_stack_loader,
+            **image_stack_loader_kwargs,
+        )
+    else:
+        raise ValueError(f"Unrecognized dataset type, {dataset_type}.")
+
+
+def create_array_dataset(
+    config: Union[DataConfig, InferenceConfig],
+    mode: Mode,
+    inputs: Sequence[NDArray[Any]],
+    targets: Optional[Sequence[NDArray[Any]]],
+) -> CareamicsDataset[InMemoryImageStack]:
+    """
+    Create a CAREamicsDataset from array data.
+
+    Parameters
+    ----------
+    config : DataConfig or InferenceConfig
+        The data configuration.
+    mode : Mode
+        Whether to create the dataset in "training", "validation" or "predicting" mode.
+    inputs : Any
+        The input sources to the dataset.
+    targets : Any, optional
+        The target sources to the dataset.
+
+    Returns
+    -------
+    CareamicsDataset[InMemoryImageStack]
+        A CAREamicsDataset
+    """
+    input_extractor = create_array_extractor(source=inputs, axes=config.axes)
+    target_extractor: Optional[PatchExtractor[InMemoryImageStack]]
+    if targets is not None:
+        target_extractor = create_array_extractor(source=targets, axes=config.axes)
+    else:
+        target_extractor = None
+    return CareamicsDataset(config, mode, input_extractor, target_extractor)
+
+
+def create_tiff_dataset(
+    config: Union[DataConfig, InferenceConfig],
+    mode: Mode,
+    inputs: Sequence[Path],
+    targets: Optional[Sequence[Path]],
+) -> CareamicsDataset[InMemoryImageStack]:
+    """
+    Create a CAREamicsDataset from tiff files that will be all loaded into memory.
+
+    Parameters
+    ----------
+    config : DataConfig or InferenceConfig
+        The data configuration.
+    mode : Mode
+        Whether to create the dataset in "training", "validation" or "predicting" mode.
+    inputs : Any
+        The input sources to the dataset.
+    targets : Any, optional
+        The target sources to the dataset.
+
+    Returns
+    -------
+    CareamicsDataset[InMemoryImageStack]
+        A CAREamicsDataset
+    """
+    input_extractor = create_tiff_extractor(
+        source=inputs,
+        axes=config.axes,
+    )
+    target_extractor: Optional[PatchExtractor[InMemoryImageStack]]
+    if targets is not None:
+        target_extractor = create_tiff_extractor(source=targets, axes=config.axes)
+    else:
+        target_extractor = None
+    dataset = CareamicsDataset(config, mode, input_extractor, target_extractor)
+    return dataset
+
+
+def create_ome_zarr_dataset(
+    config: Union[DataConfig, InferenceConfig],
+    mode: Mode,
+    inputs: Sequence[Path],
+    targets: Optional[Sequence[Path]],
+) -> CareamicsDataset[ZarrImageStack]:
+    """
+    Create a dataset from OME ZARR files.
+
+    Parameters
+    ----------
+    config : DataConfig or InferenceConfig
+        The data configuration.
+    mode : Mode
+        Whether to create the dataset in "training", "validation" or "predicting" mode.
+    inputs : Any
+        The input sources to the dataset.
+    targets : Any, optional
+        The target sources to the dataset.
+
+    Returns
+    -------
+    CareamicsDataset[ZarrImageStack]
+        A CAREamicsDataset
+    """
+
+    input_extractor = create_ome_zarr_extractor(source=inputs, axes=config.axes)
+    target_extractor: Optional[PatchExtractor[ZarrImageStack]]
+    if targets is not None:
+        target_extractor = create_ome_zarr_extractor(source=targets, axes=config.axes)
+    else:
+        target_extractor = None
+    dataset = CareamicsDataset(config, mode, input_extractor, target_extractor)
+    return dataset
+
+
+def create_custom_file_dataset(
+    config: Union[DataConfig, InferenceConfig],
+    mode: Mode,
+    inputs: Sequence[Path],
+    targets: Optional[Sequence[Path]],
+    *,
+    read_func: ReadFunc,
+    read_kwargs: dict[str, Any],
+) -> CareamicsDataset[InMemoryImageStack]:
+    """
+    Create a CAREamicsDataset from custom files that will be all loaded into memory.
+
+    Parameters
+    ----------
+    config : DataConfig or InferenceConfig
+        The data configuration.
+    mode : Mode
+        Whether to create the dataset in "training", "validation" or "predicting" mode.
+    inputs : Any
+        The input sources to the dataset.
+    targets : Any, optional
+        The target sources to the dataset.
+    read_func : Optional[ReadFunc], optional
+        A function that can that can be used to load custom data. This argument is
+        ignored unless the `data_type` is "custom".
+    image_stack_loader : Optional[ImageStackLoader], optional
+        A function for custom image stack loading. This argument is ignored unless the
+        `data_type` is "custom".
+
+    Returns
+    -------
+    CareamicsDataset[InMemoryImageStack]
+        A CAREamicsDataset
+    """
+    input_extractor = create_custom_file_extractor(
+        source=inputs, axes=config.axes, read_func=read_func, read_kwargs=read_kwargs
+    )
+    target_extractor: Optional[PatchExtractor[InMemoryImageStack]]
+    if targets is not None:
+        target_extractor = create_custom_file_extractor(
+            source=targets,
+            axes=config.axes,
+            read_func=read_func,
+            read_kwargs=read_kwargs,
+        )
+    else:
+        target_extractor = None
+    dataset = CareamicsDataset(config, mode, input_extractor, target_extractor)
+    return dataset
+
+
+def create_custom_image_stack_dataset(
+    config: Union[DataConfig, InferenceConfig],
+    mode: Mode,
+    inputs: Any,
+    targets: Optional[Any],
+    image_stack_loader: ImageStackLoader[P, GenericImageStack],
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> CareamicsDataset[GenericImageStack]:
+    """
+    Create a CAREamicsDataset from a custom `ImageStack` class.
+
+    The custom `ImageStack` class can be loaded using the `image_stack_loader` function.
+
+    Parameters
+    ----------
+    config : DataConfig or InferenceConfig
+        The data configuration.
+    mode : Mode
+        Whether to create the dataset in "training", "validation" or "predicting" mode.
+    inputs : Any
+        The input sources to the dataset.
+    targets : Any, optional
+        The target sources to the dataset.
+    image_stack_loader : ImageStackLoader
+        A function for custom image stack loading. This argument is ignored unless the
+        `data_type` is "custom".
+    *args : Any
+        Positional arguments to pass to the `image_stack_loader`.
+    **kwargs : Any
+        Key-word arguments to pass to the `image_stack_loader`.
+
+    Returns
+    -------
+    CareamicsDataset[GenericImageStack]
+        A CAREamicsDataset
+    """
+    input_extractor = create_custom_image_stack_extractor(
+        inputs,
+        config.axes,
+        image_stack_loader,
+        *args,
+        **kwargs,
+    )
+    target_extractor: Optional[PatchExtractor[GenericImageStack]]
+    if targets is not None:
+        target_extractor = create_custom_image_stack_extractor(
+            targets,
+            config.axes,
+            image_stack_loader,
+            *args,
+            **kwargs,
+        )
+    else:
+        target_extractor = None
+    dataset = CareamicsDataset(config, mode, input_extractor, target_extractor)
+    return dataset