careamics 0.0.19__py3-none-any.whl

careamics/dataset_ng/demos/care_U2OS_demo.ipynb

@@ -0,0 +1,330 @@
+{
+ "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 import create_care_configuration\n",
+    "from careamics.dataset_ng.legacy_interoperability import imageregions_to_tileinfos\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 CAREModule\n",
+    "from careamics.prediction_utils import convert_outputs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Set seed for reproducibility\n",
+    "from pytorch_lightning import seed_everything\n",
+    "\n",
+    "seed_everything(42)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Load data and set paths to it"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# instantiate data portfolio manager and download the data\n",
+    "root_path = Path(\"./data\")\n",
+    "\n",
+    "portfolio = PortfolioManager()\n",
+    "download = portfolio.denoising.CARE_U2OS.download(root_path)\n",
+    "\n",
+    "root_path = root_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": "markdown",
+   "metadata": {},
+   "source": [
+    "### Create config"
+   ]
+  },
+  {
+   "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=\"tiff\",\n",
+    "    axes=\"YX\",\n",
+    "    patch_size=(128, 128),\n",
+    "    batch_size=32,\n",
+    "    num_epochs=50,\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=config.data_config,\n",
+    "    train_data=train_path,\n",
+    "    train_data_target=target_path,\n",
+    "    val_data=test_path,\n",
+    "    val_data_target=test_target_path,\n",
+    ")\n",
+    "\n",
+    "model = CAREModule(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(2, 8, figsize=(10, 3))\n",
+    "\n",
+    "ax[0][0].set_title(\"Train batch\")\n",
+    "ax[1][0].set_title(\"Train target\")\n",
+    "for i in range(8):\n",
+    "    ax[0][i].imshow(train_batch[0].data[i][0].numpy(), cmap=\"gray\")\n",
+    "    ax[1][i].imshow(train_batch[1].data[i][0].numpy(), cmap=\"gray\")\n",
+    "\n",
+    "\n",
+    "fig, ax = plt.subplots(2, 8, figsize=(10, 3))\n",
+    "ax[0][0].set_title(\"Val batch\")\n",
+    "ax[1][0].set_title(\"Val target\")\n",
+    "for i in range(8):\n",
+    "    ax[0][i].imshow(val_batch[0].data[i][0].numpy(), cmap=\"gray\")\n",
+    "    ax[1][i].imshow(val_batch[1].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(\"care_baseline\")\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",
+    "    HyperParametersCallback(config),\n",
+    "]\n",
+    "\n",
+    "wandb_logger = WandbLogger(project=\"care-U2OS\", name=\"new-dataset\")\n",
+    "\n",
+    "trainer = Trainer(\n",
+    "    max_epochs=50, default_root_dir=root, callbacks=callbacks, logger=wandb_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.config.inference_model import InferenceConfig\n",
+    "\n",
+    "config = InferenceConfig(\n",
+    "    model_config=config,\n",
+    "    data_type=\"tiff\",\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=config, pred_data=test_path\n",
+    ")"
+   ]
+  },
+  {
+   "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",
+    "prediction = 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",
+    "# 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",
+    "# images to show\n",
+    "images = [0, 1, 2]\n",
+    "\n",
+    "fig, ax = plt.subplots(3, 3, figsize=(15, 15))\n",
+    "fig.tight_layout()\n",
+    "\n",
+    "for i in range(3):\n",
+    "    pred_image = prediction[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(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",
+    ")\n",
+    "print(\"Target PSNR: 31.53 +/- 3.71\")"
+   ]
+  }
+ ],
+ "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
+}