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

careamics/dataset_ng/dataset.py

@@ -0,0 +1,229 @@
+from collections.abc import Sequence
+from enum import Enum
+from pathlib import Path
+from typing import Any, Generic, Literal, NamedTuple, Optional, Union
+
+import numpy as np
+from numpy.typing import NDArray
+from torch.utils.data import Dataset
+from tqdm.auto import tqdm
+
+from careamics.config.data.ng_data_model import NGDataConfig
+from careamics.config.support.supported_patching_strategies import (
+    SupportedPatchingStrategy,
+)
+from careamics.config.transformations import NormalizeModel
+from careamics.dataset.dataset_utils.running_stats import WelfordStatistics
+from careamics.dataset.patching.patching import Stats
+from careamics.dataset_ng.patch_extractor import GenericImageStack, PatchExtractor
+from careamics.dataset_ng.patching_strategies import (
+    FixedRandomPatchingStrategy,
+    PatchingStrategy,
+    PatchSpecs,
+    RandomPatchingStrategy,
+    TilingStrategy,
+    WholeSamplePatchingStrategy,
+)
+from careamics.transforms import Compose
+
+
+class Mode(str, Enum):
+    TRAINING = "training"
+    VALIDATING = "validating"
+    PREDICTING = "predicting"
+
+
+class ImageRegionData(NamedTuple):
+    data: NDArray
+    source: Union[str, Literal["array"]]
+    data_shape: Sequence[int]
+    dtype: str  # dtype should be str for collate
+    axes: str
+    region_spec: PatchSpecs
+
+
+InputType = Union[Sequence[NDArray[Any]], Sequence[Path]]
+
+
+class CareamicsDataset(Dataset, Generic[GenericImageStack]):
+    def __init__(
+        self,
+        data_config: NGDataConfig,
+        mode: Mode,
+        input_extractor: PatchExtractor[GenericImageStack],
+        target_extractor: Optional[PatchExtractor[GenericImageStack]] = None,
+    ):
+        self.config = data_config
+        self.mode = mode
+
+        self.input_extractor = input_extractor
+        self.target_extractor = target_extractor
+
+        self.patching_strategy = self._initialize_patching_strategy()
+
+        self.input_stats, self.target_stats = self._initialize_statistics()
+
+        self.transforms = self._initialize_transforms()
+
+    def _initialize_patching_strategy(self) -> PatchingStrategy:
+        patching_strategy: PatchingStrategy
+        if self.mode == Mode.TRAINING:
+            if self.config.patching.name != SupportedPatchingStrategy.RANDOM:
+                raise ValueError(
+                    f"Only `random` patching strategy supported during training, got "
+                    f"{self.config.patching.name}."
+                )
+
+            patching_strategy = RandomPatchingStrategy(
+                data_shapes=self.input_extractor.shape,
+                patch_size=self.config.patching.patch_size,
+                seed=self.config.seed,
+            )
+        elif self.mode == Mode.VALIDATING:
+            if self.config.patching.name != SupportedPatchingStrategy.RANDOM:
+                raise ValueError(
+                    f"Only `random` patching strategy supported during training, got "
+                    f"{self.config.patching.name}."
+                )
+
+            patching_strategy = FixedRandomPatchingStrategy(
+                data_shapes=self.input_extractor.shape,
+                patch_size=self.config.patching.patch_size,
+                seed=self.config.seed,
+            )
+        elif self.mode == Mode.PREDICTING:
+            if (
+                self.config.patching.name != SupportedPatchingStrategy.TILED
+                and self.config.patching.name != SupportedPatchingStrategy.WHOLE
+            ):
+                raise ValueError(
+                    f"Only `tiled` and `whole` patching strategy supported during "
+                    f"training, got {self.config.patching.name}."
+                )
+            elif self.config.patching.name == SupportedPatchingStrategy.TILED:
+                patching_strategy = TilingStrategy(
+                    data_shapes=self.input_extractor.shape,
+                    tile_size=self.config.patching.patch_size,
+                    overlaps=self.config.patching.overlaps,
+                )
+            else:
+                patching_strategy = WholeSamplePatchingStrategy(
+                    data_shapes=self.input_extractor.shape
+                )
+        else:
+            raise ValueError(f"Unrecognised dataset mode {self.mode}.")
+
+        return patching_strategy
+
+    def _initialize_transforms(self) -> Optional[Compose]:
+        normalize = NormalizeModel(
+            image_means=self.input_stats.means,
+            image_stds=self.input_stats.stds,
+            target_means=self.target_stats.means,
+            target_stds=self.target_stats.stds,
+        )
+        if self.mode == Mode.TRAINING:
+            # TODO: initialize normalization separately depending on configuration
+            return Compose(transform_list=[normalize] + list(self.config.transforms))
+
+        # TODO: add TTA
+        return Compose(transform_list=[normalize])
+
+    def _calculate_stats(
+        self, data_extractor: PatchExtractor[GenericImageStack]
+    ) -> Stats:
+        image_stats = WelfordStatistics()
+        n_patches = self.patching_strategy.n_patches
+
+        for idx in tqdm(range(n_patches), desc="Computing statistics"):
+            patch_spec = self.patching_strategy.get_patch_spec(idx)
+            patch = data_extractor.extract_patch(
+                data_idx=patch_spec["data_idx"],
+                sample_idx=patch_spec["sample_idx"],
+                coords=patch_spec["coords"],
+                patch_size=patch_spec["patch_size"],
+            )
+            # TODO: statistics accept SCYX format, while patch is CYX
+            image_stats.update(patch[None, ...], sample_idx=idx)
+
+        image_means, image_stds = image_stats.finalize()
+        return Stats(image_means, image_stds)
+
+    # TODO: add running stats
+    def _initialize_statistics(self) -> tuple[Stats, Stats]:
+        if self.config.image_means is not None and self.config.image_stds is not None:
+            input_stats = Stats(self.config.image_means, self.config.image_stds)
+        else:
+            input_stats = self._calculate_stats(self.input_extractor)
+
+        target_stats = Stats((), ())
+
+        if self.config.target_means is not None and self.config.target_stds is not None:
+            target_stats = Stats(self.config.target_means, self.config.target_stds)
+        elif self.target_extractor is not None:
+            target_stats = self._calculate_stats(self.target_extractor)
+
+        return input_stats, target_stats
+
+    def __len__(self):
+        return self.patching_strategy.n_patches
+
+    def _create_image_region(
+        self, patch: np.ndarray, patch_spec: PatchSpecs, extractor: PatchExtractor
+    ) -> ImageRegionData:
+        data_idx = patch_spec["data_idx"]
+        source = extractor.image_stacks[data_idx].source
+        return ImageRegionData(
+            data=patch,
+            source=str(source),
+            dtype=str(extractor.image_stacks[data_idx].data_dtype),
+            data_shape=extractor.image_stacks[data_idx].data_shape,
+            # TODO: should it be axes of the original image instead?
+            axes=self.config.axes,
+            region_spec=patch_spec,
+        )
+
+    def __getitem__(
+        self, index: int
+    ) -> Union[tuple[ImageRegionData], tuple[ImageRegionData, ImageRegionData]]:
+        patch_spec = self.patching_strategy.get_patch_spec(index)
+        input_patch = self.input_extractor.extract_patch(
+            data_idx=patch_spec["data_idx"],
+            sample_idx=patch_spec["sample_idx"],
+            coords=patch_spec["coords"],
+            patch_size=patch_spec["patch_size"],
+        )
+
+        target_patch = (
+            self.target_extractor.extract_patch(
+                data_idx=patch_spec["data_idx"],
+                sample_idx=patch_spec["sample_idx"],
+                coords=patch_spec["coords"],
+                patch_size=patch_spec["patch_size"],
+            )
+            if self.target_extractor is not None
+            else None
+        )
+
+        if self.transforms is not None:
+            if self.target_extractor is not None:
+                input_patch, target_patch = self.transforms(input_patch, target_patch)
+            else:
+                # TODO: compose doesn't return None for target patch anymore
+                #   so have to do this annoying if else
+                (input_patch,) = self.transforms(input_patch, target_patch)
+                target_patch = None
+
+        input_data = self._create_image_region(
+            patch=input_patch, patch_spec=patch_spec, extractor=self.input_extractor
+        )
+
+        if target_patch is not None and self.target_extractor is not None:
+            target_data = self._create_image_region(
+                patch=target_patch,
+                patch_spec=patch_spec,
+                extractor=self.target_extractor,
+            )
+            return input_data, target_data
+        else:
+            return (input_data,)

careamics/dataset_ng/demos/bsd68_demo.ipynb

@@ -0,0 +1,361 @@
+{
+ "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",
+    "import tifffile\n",
+    "from careamics_portfolio import PortfolioManager\n",
+    "\n",
+    "from careamics.config.configuration_factories import (\n",
+    "    _create_ng_data_configuration,\n",
+    "    create_n2v_configuration,\n",
+    ")\n",
+    "from careamics.config.data import NGDataConfig\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": [
+    "# Set seeds for reproducibility\n",
+    "from pytorch_lightning import seed_everything\n",
+    "\n",
+    "seed = 42\n",
+    "seed_everything(seed)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Load data and set paths to it"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# instantiate data portfolio manage and download the data\n",
+    "root_path = Path(\"./data\")\n",
+    "\n",
+    "portfolio = PortfolioManager()\n",
+    "files = portfolio.denoising.N2V_BSD68.download(root_path)\n",
+    "\n",
+    "# create paths for the data\n",
+    "data_path = Path(root_path / \"denoising-N2V_BSD68.unzip/BSD68_reproducibility_data\")\n",
+    "train_path = data_path / \"train\"\n",
+    "val_path = data_path / \"val\"\n",
+    "test_path = data_path / \"test\" / \"images\"\n",
+    "gt_path = data_path / \"test\" / \"gt\"\n",
+    "\n",
+    "# list train, val and test files\n",
+    "train_files = sorted(train_path.rglob(\"*.tiff\"))\n",
+    "val_files = sorted(val_path.rglob(\"*.tiff\"))\n",
+    "test_files = sorted(test_path.rglob(\"*.tiff\"))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Visualize a single train and val image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# load training and validation image and show them side by side\n",
+    "single_train_image = tifffile.imread(train_files[0])[0]\n",
+    "single_val_image = tifffile.imread(val_files[0])[0]\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 2, figsize=(10, 5))\n",
+    "ax[0].imshow(single_train_image, cmap=\"gray\")\n",
+    "ax[0].set_title(\"Training Image\")\n",
+    "ax[1].imshow(single_val_image, cmap=\"gray\")\n",
+    "ax[1].set_title(\"Validation Image\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Create config"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "config = create_n2v_configuration(\n",
+    "    experiment_name=\"bsd68_n2v\",\n",
+    "    data_type=\"tiff\",\n",
+    "    axes=\"SYX\",\n",
+    "    patch_size=(64, 64),\n",
+    "    batch_size=64,\n",
+    "    num_epochs=100,\n",
+    ")\n",
+    "\n",
+    "# TODO until the NGDataConfig is accepted by the Confiugration, these are separte\n",
+    "ng_data_config = _create_ng_data_configuration(\n",
+    "    data_type=config.data_config.data_type,\n",
+    "    axes=config.data_config.axes,\n",
+    "    patch_size=config.data_config.patch_size,\n",
+    "    batch_size=config.data_config.batch_size,\n",
+    "    augmentations=config.data_config.transforms,\n",
+    "    train_dataloader_params=config.data_config.train_dataloader_params,\n",
+    "    val_dataloader_params=config.data_config.val_dataloader_params,\n",
+    "    seed=seed,\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Create Lightning datamodule and model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_data_module = CareamicsDataModule(\n",
+    "    data_config=ng_data_config,\n",
+    "    train_data=train_files,\n",
+    "    val_data=val_files,\n",
+    ")\n",
+    "\n",
+    "model = N2VModule(config.algorithm_config)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Manually initialize the datamodule and visualize single train and val batches"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_data_module.setup(\"fit\")\n",
+    "train_data_module.setup(\"validate\")\n",
+    "\n",
+    "train_batch = next(iter(train_data_module.train_dataloader()))\n",
+    "val_batch = next(iter(train_data_module.val_dataloader()))\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 8, figsize=(10, 5))\n",
+    "ax[0].set_title(\"Training Batch\")\n",
+    "for i in range(8):\n",
+    "    ax[i].imshow(train_batch[0].data[i][0].numpy(), cmap=\"gray\")\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 8, figsize=(10, 5))\n",
+    "ax[0].set_title(\"Validation Batch\")\n",
+    "for i in range(8):\n",
+    "    ax[i].imshow(val_batch[0].data[i][0].numpy(), cmap=\"gray\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Train the model"
+   ]
+  },
+  {
+   "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(\"bsd68_n2v\")\n",
+    "callbacks = [\n",
+    "    ModelCheckpoint(\n",
+    "        dirpath=root / \"checkpoints\",\n",
+    "        filename=\"bsd68_new_lightning_module\",\n",
+    "        save_last=True,\n",
+    "        monitor=\"val_loss\",\n",
+    "        mode=\"min\",\n",
+    "    ),\n",
+    "    HyperParametersCallback(config),\n",
+    "]\n",
+    "logger = WandbLogger(project=\"bsd68-n2v\", name=\"bsd68_new_lightning_module\")\n",
+    "\n",
+    "trainer = Trainer(\n",
+    "    max_epochs=50, default_root_dir=root, callbacks=callbacks, logger=logger\n",
+    ")\n",
+    "trainer.fit(model, datamodule=train_data_module)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Create an inference config and datamodule"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from careamics.dataset_ng.legacy_interoperability import imageregions_to_tileinfos\n",
+    "from careamics.prediction_utils import convert_outputs\n",
+    "\n",
+    "config = NGDataConfig(\n",
+    "    data_type=\"tiff\",\n",
+    "    patching={\n",
+    "        \"name\": \"tiled\",\n",
+    "        \"patch_size\": (128, 128),\n",
+    "        \"overlaps\": (32, 32),\n",
+    "    },\n",
+    "    axes=\"YX\",\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(data_config=config, pred_data=test_files)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Convert outputs to the legacy format and stitch the tiles"
+   ]
+  },
+  {
+   "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)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Visualize predictions and count metrics"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from careamics.utils.metrics import psnr, scale_invariant_psnr\n",
+    "\n",
+    "noises = [tifffile.imread(f) for f in sorted(test_path.glob(\"*.tiff\"))]\n",
+    "gts = [tifffile.imread(f) for f in sorted(gt_path.glob(\"*.tiff\"))]\n",
+    "\n",
+    "images = [0, 1, 2]\n",
+    "fig, ax = plt.subplots(3, 3, figsize=(15, 15))\n",
+    "fig.tight_layout()\n",
+    "\n",
+    "for i in range(3):\n",
+    "    pred_image = predictions[images[i]].squeeze()\n",
+    "    psnr_noisy = psnr(\n",
+    "        gts[images[i]],\n",
+    "        noises[images[i]],\n",
+    "        data_range=gts[images[i]].max() - gts[images[i]].min(),\n",
+    "    )\n",
+    "    psnr_result = psnr(\n",
+    "        gts[images[i]],\n",
+    "        pred_image,\n",
+    "        data_range=gts[images[i]].max() - gts[images[i]].min(),\n",
+    "    )\n",
+    "\n",
+    "    scale_invariant_psnr_result = scale_invariant_psnr(gts[images[i]], pred_image)\n",
+    "\n",
+    "    ax[i, 0].imshow(noises[images[i]], cmap=\"gray\")\n",
+    "    ax[i, 0].title.set_text(f\"Noisy\\nPSNR: {psnr_noisy:.2f}\")\n",
+    "\n",
+    "    ax[i, 1].imshow(pred_image, cmap=\"gray\")\n",
+    "    ax[i, 1].title.set_text(\n",
+    "        f\"Prediction\\nPSNR: {psnr_result:.2f}\\n\"\n",
+    "        f\"Scale invariant PSNR: {scale_invariant_psnr_result:.2f}\"\n",
+    "    )\n",
+    "\n",
+    "    ax[i, 2].imshow(gts[images[i]], cmap=\"gray\")\n",
+    "    ax[i, 2].title.set_text(\"Ground-truth\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "psnrs = np.zeros((len(predictions), 1))\n",
+    "scale_invariant_psnrs = np.zeros((len(predictions), 1))\n",
+    "\n",
+    "for i, (pred, gt) in enumerate(zip(predictions, gts, strict=False)):\n",
+    "    psnrs[i] = psnr(gt, pred.squeeze(), data_range=gt.max() - gt.min())\n",
+    "    scale_invariant_psnrs[i] = scale_invariant_psnr(gt, pred.squeeze())\n",
+    "\n",
+    "print(f\"PSNR: {psnrs.mean():.2f} +/- {psnrs.std():.2f}\")\n",
+    "print(\n",
+    "    f\"Scale invariant PSNR: \"\n",
+    "    f\"{scale_invariant_psnrs.mean():.2f} +/- {scale_invariant_psnrs.std():.2f}\"\n",
+    ")\n",
+    "print(\"Reported PSNR: 27.71\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "czi",
+   "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.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}