careamics 0.0.19__py3-none-any.whl

careamics/dataset_ng/demos/demo_custom_image_stack.ipynb

@@ -0,0 +1,736 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Creating a custom `ImageStack`\n",
+    "\n",
+    "You might want to write a custom `ImageStack` class if you have data stored in a format\n",
+    "that is chunked or capable of sub-file access, i.e. you want to be able to extract \n",
+    "patches during the training loop without having to load all the data into RAM at once. \n",
+    "The image stack has to follow the python `Protocol` defined in [patch_extractor/image_stack/image_stack_protocol.py](patch_extractor/image_stack/image_stack_protocol.py).\n",
+    "\n",
+    "To use a custom `ImageStack` with the `CAREamicsDataset` we will also have to write an\n",
+    "image stack loader function, that has a protocol defined in [src/careamics/dataset_ng/patch_extractor/image_stack_loader.py](patch_extractor/image_stack_loader.py). It is a callable with the function signature:\n",
+    "\n",
+    "```python\n",
+    "# example signature\n",
+    "def custom_image_stack_loader(\n",
+    "    source: Any, axes: str, *args: Any, **kwargs: Any\n",
+    ") -> Sequence[ImageStack]: ...\n",
+    "```\n",
+    "\n",
+    "In this demo, we will create a custom image stack and image stack loader for data saved\n",
+    "in a hdf5 file."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from collections.abc import Sequence\n",
+    "from pathlib import Path\n",
+    "from typing import Union\n",
+    "\n",
+    "import h5py\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import tifffile\n",
+    "from careamics_portfolio import PortfolioManager\n",
+    "from numpy.typing import DTypeLike, NDArray\n",
+    "\n",
+    "from careamics.config import create_care_configuration\n",
+    "from careamics.dataset_ng.dataset import Mode\n",
+    "from careamics.dataset_ng.factory import create_dataset"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Downloading and re-saving data\n",
+    "\n",
+    "We will resave some data as HDF5 for the purpose of this demo.\n",
+    "\n",
+    "First we download some data that is available using `careamics_portfolio`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# instantiate data portfolio manager and download the data\n",
+    "data_path = Path(\"./data\")\n",
+    "\n",
+    "portfolio = PortfolioManager()\n",
+    "download = portfolio.denoising.CARE_U2OS.download(data_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "root_path = data_path / \"denoising-CARE_U2OS.unzip\" / \"data\" / \"U2OS\"\n",
+    "train_path = root_path / \"train\" / \"low\"\n",
+    "target_path = root_path / \"train\" / \"GT\"\n",
+    "test_path = root_path / \"test\" / \"low\"\n",
+    "test_target_path = root_path / \"test\" / \"GT\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# checking the train input and target files we have\n",
+    "print(list(train_path.glob(\"*.tif\")))\n",
+    "print(list(target_path.glob(\"*.tif\")))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Save as HDF5\n",
+    "\n",
+    "We will save all the images in a HDF5 file, the input images under a \"train\" path and \n",
+    "target images under a \"target\" path, and all the images will have their original file \n",
+    "name."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "hdf5_file_path = data_path / \"CARE_U2OS-train.h5\"\n",
+    "\n",
+    "if not hdf5_file_path.is_file():\n",
+    "    with h5py.File(name=hdf5_file_path, mode=\"w\") as file:\n",
+    "        train_group = file.create_group(\"train_input\")\n",
+    "        target_group = file.create_group(\"train_target\")\n",
+    "        test_group = file.create_group(\"test_input\")\n",
+    "        test_target_group = file.create_group(\"test_target\")\n",
+    "        for path in train_path.glob(\"*.tif\"):\n",
+    "            image = tifffile.imread(path)\n",
+    "            train_group.create_dataset(name=path.stem, data=image)\n",
+    "        for path in target_path.glob(\"*.tif\"):\n",
+    "            image = tifffile.imread(path)\n",
+    "            target_group.create_dataset(name=path.stem, data=image)\n",
+    "        for path in test_path.glob(\"*.tif\"):\n",
+    "            image = tifffile.imread(path)\n",
+    "            test_group.create_dataset(name=path.stem, data=image)\n",
+    "        for path in test_target_path.glob(\"*.tif\"):\n",
+    "            image = tifffile.imread(path)\n",
+    "            test_target_group.create_dataset(name=path.stem, data=image)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Defining the image stack\n",
+    "\n",
+    "An ImageStack must have the attributes: `data_shape`, `data_dtype` and `source` and the\n",
+    "method `extract_patch`.\n",
+    "\n",
+    "The `data_shape` attribute should be shape the data would have once reshaped to match the axes \n",
+    "`SC(Z)YX`.\n",
+    "\n",
+    "The `data_dtype` attribute is the data type of the underlying array.\n",
+    "\n",
+    "The `source` attribute should have the type `Path`, it will be returned alongside the patches by the\n",
+    "`CAREamicsDataset` and can be used as a way to identify where the data came from. In the\n",
+    "future it may be used as a way to automatically save predictions to disk.\n",
+    "\n",
+    "The `extract_patch` method needs to return a patch for a given `sample_index`, `coords` \n",
+    "and `patch_size` that has the axes `SC(Z)YX`. So, for our HDF5 case the patches need to \n",
+    "be reshaped when the `extract_patch_method` is called."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from careamics.dataset_ng.patch_extractor.image_stack.zarr_image_stack import (\n",
+    "    _reshaped_array_shape,\n",
+    ")\n",
+    "\n",
+    "from careamics.dataset.dataset_utils import reshape_array\n",
+    "\n",
+    "\n",
+    "class HDF5ImageStack:\n",
+    "\n",
+    "    def __init__(self, image_data: h5py.Dataset, axes: str):\n",
+    "        self._image_data = image_data\n",
+    "        self._original_axes = axes\n",
+    "        self._original_data_shape = image_data.shape\n",
+    "        self.data_shape = _reshaped_array_shape(\n",
+    "            self._original_axes, self._image_data.shape\n",
+    "        )\n",
+    "\n",
+    "    @property\n",
+    "    def data_dtype(self) -> DTypeLike:\n",
+    "        return self._image_data.dtype\n",
+    "\n",
+    "    @property\n",
+    "    def source(self) -> Path:\n",
+    "        return Path(self._image_data.file.filename + str(self._image_data.name))\n",
+    "\n",
+    "    # this method is almost an exact copy of the ZarrImageStack.extract patch\n",
+    "    def extract_patch(\n",
+    "        self, sample_idx: int, coords: Sequence[int], patch_size: Sequence[int]\n",
+    "    ) -> NDArray:\n",
+    "        # original axes assumed to be any subset of STCZYX (containing YX), in any order\n",
+    "        # arguments must be transformed to index data in original axes order\n",
+    "        # to do this: loop through original axes and append correct index/slice\n",
+    "        #   for each case: STCZYX\n",
+    "        #   Note: if any axis is not present in original_axes it is skipped.\n",
+    "\n",
+    "        # guard for no S and T in original axes\n",
+    "        if (\"S\" not in self._original_axes) and (\"T\" not in self._original_axes):\n",
+    "            if sample_idx not in [0, -1]:\n",
+    "                raise IndexError(\n",
+    "                    f\"Sample index {sample_idx} out of bounds for S axes with size \"\n",
+    "                    f\"{self.data_shape[0]}\"\n",
+    "                )\n",
+    "\n",
+    "        patch_slice: list[Union[int, slice]] = []\n",
+    "        for d in self._original_axes:\n",
+    "            if d == \"S\":\n",
+    "                patch_slice.append(self._get_S_index(sample_idx))\n",
+    "            elif d == \"T\":\n",
+    "                patch_slice.append(self._get_T_index(sample_idx))\n",
+    "            elif d == \"C\":\n",
+    "                patch_slice.append(slice(None, None))\n",
+    "            elif d == \"Z\":\n",
+    "                patch_slice.append(slice(coords[0], coords[0] + patch_size[0]))\n",
+    "            elif d == \"Y\":\n",
+    "                y_idx = 0 if \"Z\" not in self._original_axes else 1\n",
+    "                patch_slice.append(\n",
+    "                    slice(coords[y_idx], coords[y_idx] + patch_size[y_idx])\n",
+    "                )\n",
+    "            elif d == \"X\":\n",
+    "                x_idx = 1 if \"Z\" not in self._original_axes else 2\n",
+    "                patch_slice.append(\n",
+    "                    slice(coords[x_idx], coords[x_idx] + patch_size[x_idx])\n",
+    "                )\n",
+    "            else:\n",
+    "                raise ValueError(f\"Unrecognised axis '{d}', axes should be in STCZYX.\")\n",
+    "\n",
+    "        patch = self._image_data[tuple(patch_slice)]\n",
+    "        patch_axes = self._original_axes.replace(\"S\", \"\").replace(\"T\", \"\")\n",
+    "        return reshape_array(patch, patch_axes)[0]  # remove first sample dim\n",
+    "\n",
+    "    def _get_T_index(self, sample_idx: int) -> int:\n",
+    "        \"\"\"Get T index given `sample_idx`.\"\"\"\n",
+    "        if \"T\" not in self._original_axes:\n",
+    "            raise ValueError(\"No 'T' axis specified in original data axes.\")\n",
+    "        axis_idx = self._original_axes.index(\"T\")\n",
+    "        dim = self._original_data_shape[axis_idx]\n",
+    "\n",
+    "        # new S' = S*T\n",
+    "        # T_idx = S_idx' // T_size\n",
+    "        # S_idx = S_idx' % T_size\n",
+    "        # - floor divide finds the row\n",
+    "        # - modulus finds how far along the row i.e. the column\n",
+    "        return sample_idx % dim\n",
+    "\n",
+    "    def _get_S_index(self, sample_idx: int) -> int:\n",
+    "        \"\"\"Get S index given `sample_idx`.\"\"\"\n",
+    "        if \"S\" not in self._original_axes:\n",
+    "            raise ValueError(\"No 'S' axis specified in original data axes.\")\n",
+    "        if \"T\" in self._original_axes:\n",
+    "            T_axis_idx = self._original_axes.index(\"T\")\n",
+    "            T_dim = self._original_data_shape[T_axis_idx]\n",
+    "\n",
+    "            # new S' = S*T\n",
+    "            # T_idx = S_idx' // T_size\n",
+    "            # S_idx = S_idx' % T_size\n",
+    "            # - floor divide finds the row\n",
+    "            # - modulus finds how far along the row i.e. the column\n",
+    "            return sample_idx // T_dim\n",
+    "        else:\n",
+    "            return sample_idx"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Now define the image loader\n",
+    "\n",
+    "The loader needs to have the first two arguments be `source` and `axes`, then any \n",
+    "additional kwargs are allowed. However, note that the additional kwargs have to be \n",
+    "shared by both the input and the target when the dataset is initialized.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# A image stack loader\n",
+    "# both the input and target image stacks must be contained within the same HDF5 file\n",
+    "def hdf5_image_stack_loader(\n",
+    "    source: Sequence[str], axes: str, file: h5py.File\n",
+    ") -> Sequence[HDF5ImageStack]:\n",
+    "    image_stacks: list[HDF5ImageStack] = []\n",
+    "    for data_path in source:\n",
+    "        if data_path not in file:\n",
+    "            raise KeyError(f\"Data does not exist at path '{data_path}'\")\n",
+    "        image_data = file[data_path]\n",
+    "        if not isinstance(image_data, h5py.Dataset):\n",
+    "            raise TypeError(f\"HDF5 node at path '{data_path}' is not a Dataset.\")\n",
+    "        image_stacks.append(HDF5ImageStack(image_data, axes=axes))\n",
+    "    return image_stacks"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# This is an alternative hdf5 image stack loader\n",
+    "# The input and target files can be contained in separate hdf5 files\n",
+    "# An HDF5Source typed dict has to be defined\n",
+    "# this is to allow both the file and the data paths to be combined in a single argument\n",
+    "\n",
+    "from typing import TypedDict\n",
+    "\n",
+    "\n",
+    "class HDF5Source(TypedDict):\n",
+    "    file: h5py.File\n",
+    "    data_path: str\n",
+    "\n",
+    "\n",
+    "def hdf5_image_stack_loader_alt(\n",
+    "    source: Sequence[HDF5Source], axes: str\n",
+    ") -> Sequence[HDF5ImageStack]:\n",
+    "    image_stacks: list[HDF5ImageStack] = []\n",
+    "    for image_stack_source in source:\n",
+    "        data_path = image_stack_source[\"data_path\"]\n",
+    "        file = image_stack_source[\"file\"]\n",
+    "        if data_path not in file:\n",
+    "            raise KeyError(f\"Data does not exist at path '{data_path}'\")\n",
+    "        image_data = file[data_path]\n",
+    "        if not isinstance(image_data, h5py.Dataset):\n",
+    "            raise TypeError(f\"HDF5 node at path '{data_path}' is not a Dataset.\")\n",
+    "        image_stacks.append(HDF5ImageStack(image_data, axes=axes))\n",
+    "    return image_stacks"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Now we test it\n",
+    "\n",
+    "### create a configuration for the data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_files = sorted(train_path.glob(\"*.tif\"))\n",
+    "train_target_files = sorted(target_path.glob(\"*.tif\"))\n",
+    "\n",
+    "config = create_care_configuration(\n",
+    "    experiment_name=\"care_U20S\",\n",
+    "    data_type=\"custom\",\n",
+    "    axes=\"YX\",\n",
+    "    patch_size=[128, 128],\n",
+    "    batch_size=32,\n",
+    "    num_epochs=50,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "hdf5_file = h5py.File(hdf5_file_path, mode=\"r\")\n",
+    "\n",
+    "inputs = sorted([f\"train_input/{key}\" for key in hdf5_file[\"train_input\"].keys()])\n",
+    "targets = sorted([f\"train_target/{key}\" for key in hdf5_file[\"train_target\"].keys()])\n",
+    "test_inputs = sorted([f\"test_input/{key}\" for key in hdf5_file[\"test_input\"].keys()])\n",
+    "test_targets = sorted([f\"test_target/{key}\" for key in hdf5_file[\"test_target\"].keys()])\n",
+    "\n",
+    "dataset = create_dataset(\n",
+    "    config=config.data_config,\n",
+    "    mode=Mode.TRAINING,\n",
+    "    inputs=inputs,\n",
+    "    targets=targets,\n",
+    "    in_memory=False,\n",
+    "    image_stack_loader=hdf5_image_stack_loader,\n",
+    "    image_stack_loader_kwargs={\"file\": hdf5_file},\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Index the dataset and display the result"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, axes = plt.subplots(1, 2)\n",
+    "train_input, target = dataset[0]\n",
+    "axes[0].imshow(train_input.data[0])\n",
+    "axes[0].set_title(\"Input\")\n",
+    "axes[1].imshow(target.data[0])\n",
+    "axes[1].set_title(\"Target\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# input and target are ImageRegionData objects\n",
+    "train_input, target"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Test the alternative image stack loader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "hdf5_file = h5py.File(hdf5_file_path, mode=\"r\")\n",
+    "\n",
+    "data_keys = sorted(hdf5_file[\"train_input\"].keys())\n",
+    "\n",
+    "# for the alternative image stack loader we have to construct a list of dicts\n",
+    "# because we defined the source type to be a HDF5Source typed dict\n",
+    "inputs: list[HDF5Source] = [\n",
+    "    {\"data_path\": f\"train_input/{key}\", \"file\": hdf5_file} for key in data_keys\n",
+    "]\n",
+    "targets: list[HDF5Source] = [\n",
+    "    {\"data_path\": f\"train_target/{key}\", \"file\": hdf5_file} for key in data_keys\n",
+    "]\n",
+    "\n",
+    "dataset = create_dataset(\n",
+    "    config=config.data_config,\n",
+    "    mode=Mode.TRAINING,\n",
+    "    inputs=inputs,\n",
+    "    targets=targets,\n",
+    "    in_memory=False,\n",
+    "    image_stack_loader=hdf5_image_stack_loader_alt,\n",
+    "    # now we don't have any additional kwargs\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# display the first item\n",
+    "# note this will be a different patch because of the random patching\n",
+    "fig, axes = plt.subplots(1, 2)\n",
+    "train_input, target = dataset[0]\n",
+    "axes[0].imshow(train_input.data[0])\n",
+    "axes[0].set_title(\"Input\")\n",
+    "axes[1].imshow(target.data[0])\n",
+    "axes[1].set_title(\"Target\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Now let's run N2V training pipeline and see how it performs"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Creating the lightning data module for training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from careamics.config.inference_model import InferenceConfig\n",
+    "\n",
+    "from careamics.lightning.dataset_ng.data_module import CareamicsDataModule\n",
+    "\n",
+    "hdf5_file = h5py.File(hdf5_file_path, mode=\"r\")\n",
+    "\n",
+    "train_data_keys = sorted(hdf5_file[\"train_input\"].keys())\n",
+    "\n",
+    "inputs: list[HDF5Source] = [\n",
+    "    {\"data_path\": f\"train_input/{key}\", \"file\": hdf5_file} for key in train_data_keys\n",
+    "]\n",
+    "targets: list[HDF5Source] = [\n",
+    "    {\"data_path\": f\"train_target/{key}\", \"file\": hdf5_file} for key in train_data_keys\n",
+    "]\n",
+    "\n",
+    "test_data_keys = sorted(hdf5_file[\"test_input\"].keys())\n",
+    "test_inputs: list[HDF5Source] = [\n",
+    "    {\"data_path\": f\"test_input/{key}\", \"file\": hdf5_file} for key in test_data_keys\n",
+    "]\n",
+    "test_targets: list[HDF5Source] = [\n",
+    "    {\"data_path\": f\"test_target/{key}\", \"file\": hdf5_file} for key in test_data_keys\n",
+    "]\n",
+    "config = create_care_configuration(\n",
+    "    experiment_name=\"care_U20S\",\n",
+    "    data_type=\"custom\",\n",
+    "    axes=\"YX\",\n",
+    "    patch_size=[128, 128],\n",
+    "    batch_size=32,\n",
+    "    num_epochs=50,\n",
+    ")\n",
+    "train_data_module = CareamicsDataModule(\n",
+    "    data_config=config.data_config,\n",
+    "    train_data=inputs,\n",
+    "    train_data_target=targets,\n",
+    "    val_data=inputs,\n",
+    "    val_data_target=targets,\n",
+    "    image_stack_loader=hdf5_image_stack_loader_alt,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Creating the model and the trainer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pytorch_lightning import Trainer\n",
+    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
+    "\n",
+    "from careamics.lightning.dataset_ng.lightning_modules import CAREModule\n",
+    "\n",
+    "root = Path(\"care_stack_loader\")\n",
+    "\n",
+    "# TODO: replace with N2V!!!\n",
+    "model = CAREModule(config.algorithm_config)\n",
+    "\n",
+    "callbacks = [\n",
+    "    ModelCheckpoint(\n",
+    "        dirpath=root / \"checkpoints\",\n",
+    "        filename=\"care_baseline\",\n",
+    "        save_last=True,\n",
+    "        monitor=\"val_loss\",\n",
+    "        mode=\"min\",\n",
+    "    )\n",
+    "]\n",
+    "\n",
+    "trainer = Trainer(max_epochs=50, default_root_dir=root, callbacks=callbacks)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Training the model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer.fit(model, datamodule=train_data_module)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Creating the inference data module"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "inference_config = InferenceConfig(\n",
+    "    model_config=config,\n",
+    "    data_type=\"custom\",\n",
+    "    tile_size=(128, 128),\n",
+    "    tile_overlap=(32, 32),\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(\n",
+    "    data_config=inference_config,\n",
+    "    pred_data=test_inputs,\n",
+    "    image_stack_loader=hdf5_image_stack_loader_alt,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Running the prediction on the test set"
+   ]
+  },
+  {
+   "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",
+    "predictions = trainer.predict(model, datamodule=inf_data_module)\n",
+    "tile_infos = imageregions_to_tileinfos(predictions)\n",
+    "prediction = convert_outputs(tile_infos, tiled=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Displaying the predictions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from careamics.utils.metrics import psnr, scale_invariant_psnr\n",
+    "\n",
+    "# Show two images\n",
+    "noises = [tifffile.imread(f) for f in sorted(test_path.glob(\"*.tif\"))]\n",
+    "gts = [tifffile.imread(f) for f in sorted(test_target_path.glob(\"*.tif\"))]\n",
+    "\n",
+    "fig, ax = plt.subplots(3, 3, figsize=(7, 7))\n",
+    "fig.tight_layout()\n",
+    "\n",
+    "for i in range(3):\n",
+    "    pred_image = prediction[i].squeeze()\n",
+    "    psnr_noisy = psnr(\n",
+    "        gts[i],\n",
+    "        noises[i],\n",
+    "        data_range=gts[i].max() - gts[i].min(),\n",
+    "    )\n",
+    "    psnr_result = psnr(\n",
+    "        gts[i],\n",
+    "        pred_image,\n",
+    "        data_range=gts[i].max() - gts[i].min(),\n",
+    "    )\n",
+    "\n",
+    "    scale_invariant_psnr_result = scale_invariant_psnr(gts[i], pred_image)\n",
+    "\n",
+    "    ax[i, 0].imshow(noises[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[i], cmap=\"gray\")\n",
+    "    ax[i, 2].title.set_text(\"Ground-truth\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Calculating the metrics on the test set"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "psnrs = np.zeros((len(prediction), 1))\n",
+    "scale_invariant_psnrs = np.zeros((len(prediction), 1))\n",
+    "\n",
+    "for i, (pred, gt) in enumerate(zip(prediction, 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",
+    ")"
+   ]
+  }
+ ],
+ "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.21"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}