careamics 0.0.19__py3-none-any.whl

careamics/dataset_ng/demos/bsd68_zarr_demo.ipynb

@@ -0,0 +1,453 @@
+{
+ "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 zarr\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": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import platform\n",
+    "\n",
+    "import torch\n",
+    "\n",
+    "print(platform.processor() in ('arm', 'arm64') and torch.backends.mps.is_available())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Create zarr dataset"
+   ]
+  },
+  {
+   "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 = root_path / \"denoising-N2V_BSD68.unzip/BSD68_reproducibility_data\"\n",
+    "# train_path = data_path / \"train\" / \"DCNN400_train_gaussian25.tiff\"\n",
+    "# val_path = data_path / \"val\" / \"DCNN400_validation_gaussian25.tiff\"\n",
+    "\n",
+    "# root = zarr.create_group(root_path / \"bsd68.zarr\")\n",
+    "# print(f\"Creating zarr at: {root.store_path}\")\n",
+    "\n",
+    "# # add train images to train group\n",
+    "# train = root.create_group('train')\n",
+    "# train_img = tifffile.imread(train_path)\n",
+    "\n",
+    "# for i in range(train_img.shape[0]):\n",
+    "\n",
+    "#     img = train_img[i]\n",
+    "#     name = f\"img_{i:04d}\"\n",
+    "#     train.create_array(name=name, data=img, chunks=(128, 128))\n",
+    "\n",
+    "# # add validation images to train group\n",
+    "# val = root.create_group('val')\n",
+    "# val_img = tifffile.imread(val_path)\n",
+    "\n",
+    "# for i in range(val_img.shape[0]):\n",
+    "\n",
+    "#     img = val_img[i]\n",
+    "#     name = f\"img_{i:04d}\"\n",
+    "#     val.create_array(name=name, data=img, chunks=(128, 128))\n",
+    "\n",
+    "\n",
+    "# # add test gt to zarr\n",
+    "# test_gt_f = sorted([f for f in (data_path / \"test\" / \"gt\").glob(\"*.tiff\")])\n",
+    "# test_gt_z = root.create_group('test_gt')\n",
+    "# for i in test_gt_f:\n",
+    "#     img = tifffile.imread(i)\n",
+    "#     name = i.stem\n",
+    "#     test_gt_z.create_array(name=name, data=img, chunks=(128, 128))\n",
+    "\n",
+    "\n",
+    "# test_raw_f = sorted([f for f in (data_path / \"test\" / \"images\").glob(\"*.tiff\")])\n",
+    "# test_raw_z = root.create_group('test_raw')\n",
+    "# for i in test_raw_f:\n",
+    "#     img = tifffile.imread(i)\n",
+    "#     name = i.stem\n",
+    "#     test_raw_z.create_array(name=name, data=img, chunks=(128, 128))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# data path\n",
+    "g = zarr.open(Path(\"data\") / \"bsd68.zarr\")\n",
+    "train_path = str(g[\"train\"].store_path)\n",
+    "val_path = str(g[\"val\"].store_path)"
+   ]
+  },
+  {
+   "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",
+    "list_train_arrays = list(g[\"train\"].array_keys())\n",
+    "list_val_arrays = list(g[\"val\"].array_keys())\n",
+    "\n",
+    "single_train_image = g[\"train\"][list_train_arrays[0]]\n",
+    "single_val_image = g[\"val\"][list_val_arrays[0]]\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": [
+    "epochs = 1\n",
+    "steps = 50\n",
+    "batch_size = 1\n",
+    "\n",
+    "config = create_n2v_configuration(\n",
+    "    experiment_name=\"bsd68_n2v\",\n",
+    "    data_type=\"custom\",\n",
+    "    axes=\"YX\",\n",
+    "    patch_size=(64, 64),\n",
+    "    batch_size=batch_size,\n",
+    "    num_epochs=epochs,\n",
+    "    num_steps=steps,\n",
+    ")\n",
+    "\n",
+    "# TODO until the NGDataConfig is accepted by the Configuration, these are separate\n",
+    "ng_data_config = _create_ng_data_configuration(\n",
+    "    data_type=\"zarr\", # specific to NG Dataset\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",
+    "\n",
+    "ng_data_config.set_means_and_stds(\n",
+    "    image_means=[0], image_stds=[1]\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_path],\n",
+    "    val_data=[val_path],\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",
+    "tdm = train_data_module.train_dataloader()\n",
+    "vdm = train_data_module.val_dataloader()\n",
+    "train_batch = next(iter(tdm))\n",
+    "val_batch = next(iter(vdm))\n",
+    "\n",
+    "fig, ax = plt.subplots(1, batch_size, figsize=(10, 5))\n",
+    "ax[0].set_title(\"Training Batch\")\n",
+    "for i in range(batch_size):\n",
+    "    ax[i].imshow(train_batch.data[i][0].numpy(), cmap=\"gray\")\n",
+    "\n",
+    "fig, ax = plt.subplots(1, batch_size, figsize=(10, 5))\n",
+    "ax[0].set_title(\"Validation Batch\")\n",
+    "for i in range(batch_size):\n",
+    "    ax[i].imshow(val_batch.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",
+    "\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",
+    "\n",
+    "trainer = Trainer(\n",
+    "    max_epochs=epochs,\n",
+    "    default_root_dir=root,\n",
+    "    callbacks=callbacks,\n",
+    "    limit_train_batches=steps\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",
+    "test_files = [str(g[\"test_raw\"].store_path)]\n",
+    "\n",
+    "\n",
+    "config = NGDataConfig(\n",
+    "    data_type=\"zarr\",\n",
+    "    patching={\n",
+    "        \"name\": \"tiled\",\n",
+    "        \"patch_size\": (128, 128),\n",
+    "        \"overlaps\": (48, 48),\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_str = sorted(g[\"test_raw\"].array_keys())\n",
+    "gts_str = sorted(g[\"test_gt\"].array_keys())\n",
+    "\n",
+    "noises = [\n",
+    "    g[\"test_raw\"][arr] for arr in noises_str\n",
+    "]\n",
+    "gts = [\n",
+    "    g[\"test_gt\"][arr] for arr in gts_str\n",
+    "]\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",
+    "    gts_arrs = gts[images[i]][...]\n",
+    "    noises_arrs = noises[images[i]][...]\n",
+    "\n",
+    "    pred_image = predictions[images[i]].squeeze()\n",
+    "    psnr_noisy = psnr(\n",
+    "        gts_arrs,\n",
+    "        noises_arrs,\n",
+    "        data_range=gts_arrs.max() - gts_arrs.min(),\n",
+    "    )\n",
+    "    psnr_result = psnr(\n",
+    "        gts_arrs,\n",
+    "        pred_image,\n",
+    "        data_range=gts_arrs.max() - gts_arrs.min(),\n",
+    "    )\n",
+    "\n",
+    "    scale_invariant_psnr_result = scale_invariant_psnr(gts_arrs, 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_arrs, 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",
+    "    gt = gt[...]\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\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "careamics (3.13.7)",
+   "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.13.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}