docs: typo fixes in notebooks

okunator · okunator · commit fffe6cc62c20 · 2024-07-05T16:23:27.000+03:00
diff --git a/examples/pannuke_nuclei_segmentation_cppnet.ipynb b/examples/pannuke_nuclei_segmentation_cppnet.ipynb
@@ -41,10 +41,12 @@
    ],
    "source": [
     "# version info\n",
-    "import torch\n",
+    "from platform import python_version\n",
+    "\n",
     "import lightning\n",
+    "import torch\n",
+    "\n",
     "import cellseg_models_pytorch\n",
-    "from platform import python_version\n",
     "\n",
     "print(\"torch version:\", torch.__version__)\n",
     "print(\"lightning version:\", lightning.__version__)\n",
@@ -86,6 +88,7 @@
    ],
    "source": [
     "from pathlib import Path\n",
+    "\n",
     "from cellseg_models_pytorch.datamodules import PannukeDataModule\n",
     "\n",
     "# fold1 and fold2 are used for training, fold3 is used for validation\n",
@@ -145,17 +148,18 @@
     }
    ],
    "source": [
-    "import numpy as np\n",
     "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
     "from skimage.color import label2rgb\n",
     "\n",
-    "# filehandler contains methods to read and write images and masks\n",
-    "from cellseg_models_pytorch.utils import FileHandler\n",
     "from cellseg_models_pytorch.transforms.functional import (\n",
-    "    gen_stardist_maps,\n",
     "    gen_dist_maps,\n",
+    "    gen_stardist_maps,\n",
     ")\n",
     "\n",
+    "# filehandler contains methods to read and write images and masks\n",
+    "from cellseg_models_pytorch.utils import FileHandler\n",
+    "\n",
     "img_dir = save_dir / \"train\" / \"images\"\n",
     "mask_dir = save_dir / \"train\" / \"labels\"\n",
     "imgs = sorted(img_dir.glob(\"*\"))\n",
@@ -216,13 +220,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from typing import Dict, List, Optional, Tuple\n",
+    "\n",
+    "import lightning.pytorch as pl\n",
     "import torch\n",
     "import torch.nn as nn\n",
     "import torch.optim as optim\n",
-    "import torchmetrics\n",
-    "import lightning.pytorch as pl\n",
-    "from copy import deepcopy\n",
-    "from typing import List, Tuple, Dict, Optional\n",
     "\n",
     "\n",
     "class SegmentationExperiment(pl.LightningModule):\n",
@@ -232,7 +235,6 @@
     "        multitask_loss: Dict[str, nn.Module],\n",
     "        optimizer: optim.Optimizer,\n",
     "        scheduler: optim.lr_scheduler._LRScheduler,\n",
-    "        branch_metrics: Optional[Dict[str, List[torchmetrics.Metric]]] = None,\n",
     "        optimizer_kwargs: Optional[Dict[str, float]] = None,\n",
     "        scheduler_kwargs: Optional[Dict[str, float]] = None,\n",
     "        **kwargs,\n",
@@ -249,16 +251,10 @@
     "        self.optimizer_kwargs = optimizer_kwargs or {}\n",
     "        self.scheduler_kwargs = scheduler_kwargs or {}\n",
     "\n",
-    "        self.branch_metrics = branch_metrics\n",
     "        self.criterion = multitask_loss\n",
     "\n",
     "        self.save_hyperparameters(ignore=\"model\")\n",
     "\n",
-    "        metrics = self.configure_metrics()\n",
-    "        self.train_metrics = deepcopy(metrics)\n",
-    "        self.val_metrics = deepcopy(metrics)\n",
-    "        self.test_metrics = deepcopy(metrics)\n",
-    "\n",
     "    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:\n",
     "        \"\"\"Forward pass.\"\"\"\n",
     "        return self.model(x)\n",
@@ -298,16 +294,9 @@
     "        # forward backward pass\n",
     "        loss, soft_masks, targets = self.step(batch)\n",
     "\n",
-    "        # compute training metrics\n",
-    "        metrics = self.compute_metrics(soft_masks, targets, \"train\")\n",
-    "\n",
     "        # log the loss\n",
     "        self.log(\"train_loss\", loss, on_step=True, on_epoch=False, prog_bar=True)\n",
     "\n",
-    "        # log the metrics\n",
-    "        for metric_name, metric in metrics.items():\n",
-    "            self.log(metric_name, metric, on_step=True, on_epoch=False, prog_bar=True)\n",
-    "\n",
     "        return loss\n",
     "\n",
     "    def validation_step(\n",
@@ -317,16 +306,9 @@
     "        # forward pass\n",
     "        loss, soft_masks, targets = self.step(batch)\n",
     "\n",
-    "        # compute validation metrics\n",
-    "        metrics = self.compute_metrics(soft_masks, targets, \"val\")\n",
-    "\n",
     "        # log the loss\n",
     "        self.log(\"val_loss\", loss, on_step=False, on_epoch=True, prog_bar=False)\n",
     "\n",
-    "        # log the metrics\n",
-    "        for metric_name, metric in metrics.items():\n",
-    "            self.log(metric_name, metric, on_step=True, on_epoch=True, prog_bar=True)\n",
-    "\n",
     "        return loss\n",
     "\n",
     "    def test_step(\n",
@@ -336,50 +318,11 @@
     "        # forward pass\n",
     "        loss, soft_masks, targets = self.step(batch)\n",
     "\n",
-    "        # compute validation metrics\n",
-    "        metrics = self.compute_metrics(soft_masks, targets, \"test\")\n",
-    "\n",
     "        # log the loss\n",
     "        self.log(\"test_loss\", loss, on_step=False, on_epoch=True, prog_bar=False)\n",
     "\n",
-    "        # log the metrics\n",
-    "        for metric_name, metric in metrics.items():\n",
-    "            self.log(metric_name, metric, on_step=False, on_epoch=True, prog_bar=False)\n",
-    "\n",
     "        return loss\n",
     "\n",
-    "    def compute_metrics(\n",
-    "        self,\n",
-    "        preds: Dict[str, torch.Tensor],\n",
-    "        targets: Dict[str, torch.Tensor],\n",
-    "        phase: str,\n",
-    "    ) -> Dict[str, torch.Tensor]:\n",
-    "        \"\"\"Compute metrics for logging.\"\"\"\n",
-    "        if phase == \"train\":\n",
-    "            metrics_dict = self.train_metrics\n",
-    "        elif phase == \"val\":\n",
-    "            metrics_dict = self.val_metrics\n",
-    "        elif phase == \"test\":\n",
-    "            metrics_dict = self.test_metrics\n",
-    "\n",
-    "        ret = {}\n",
-    "        for metric_name, metric in metrics_dict.items():\n",
-    "            if metric is not None:\n",
-    "                branch = metric_name.split(\"_\")[0]\n",
-    "                ret[metric_name] = metric(preds[branch], targets[branch])\n",
-    "\n",
-    "        return ret\n",
-    "\n",
-    "    def configure_metrics(self) -> Dict[str, torchmetrics.Metric]:\n",
-    "        \"\"\"Configure the metrics for the model.\"\"\"\n",
-    "        # We can put all the metrics in a ModuleDict and return it.\n",
-    "        metrics = nn.ModuleDict()\n",
-    "        for branch, metric_list in self.branch_metrics.items():\n",
-    "            for metric in metric_list:\n",
-    "                metrics[f\"{branch}_{metric}\"] = metric\n",
-    "\n",
-    "        return metrics\n",
-    "\n",
     "    def configure_optimizers(self) -> List[optim.Optimizer]:\n",
     "        \"\"\"Configure the optimizers for the model.\"\"\"\n",
     "        opt = self.optimizer(self.parameters(), **self.optimizer_kwargs)\n",
@@ -411,9 +354,6 @@
     "\n",
     "**Branch losses.** For each output of the model, we define a joint-loss function. These losses are summed together during backprop to form a multi-task loss. In `CPP-Net`, we get two outputs for the `stardist` output instead of one. The reason is that `CPP-Net` was designed to refine the `stardist`-output for better segmentation performance, thus, the model returns the both the regular `stardist` and the refined `stardist_refined` output. For both of these outputs we set the masked `MAE`-loss. For the `\"dist\"`-outputs we use the masked `BCELoss` (binary-cross-entropy) with `MSE`, and for the `\"type\"`-output we will use the masked multiclass categorical `CELoss` with `DiceLoss`.\n",
     "\n",
-    "**Logging metrics.**\n",
-    "For the nuclei type masks we will monitor Jaccard-index i.e. the mIoU metric during training. The metric is averaged over the class-specific mIoUs. The metrics are from the `torchmetrics` library. Also, we will monitor the MSE-metric for the `stardist`-outputs.\n",
-    "\n",
     "**Optimizer and scheduler.**\n",
     "The optimizer used here is [AdamW](https://arxiv.org/abs/1711.05101). The learning rate is scheduled with the `ReduceLROnPlateau` schedule. The learning rate will be set to 0.0003. \n",
     "\n",
@@ -438,18 +378,17 @@
    ],
    "source": [
     "import torch.optim as optim\n",
-    "from torchmetrics import JaccardIndex, MeanSquaredError\n",
     "\n",
-    "from cellseg_models_pytorch.models import cppnet_base_multiclass\n",
     "from cellseg_models_pytorch.losses import (\n",
     "    MAE,\n",
     "    MSE,\n",
-    "    DiceLoss,\n",
     "    BCELoss,\n",
     "    CELoss,\n",
+    "    DiceLoss,\n",
     "    JointLoss,\n",
     "    MultiTaskLoss,\n",
     ")\n",
+    "from cellseg_models_pytorch.models import cppnet_base_multiclass\n",
     "\n",
     "# seed the experiment for reproducibility\n",
     "pl.seed_everything(42)\n",
@@ -481,21 +420,6 @@
     "    },  # weights for each branch\n",
     ")\n",
     "\n",
-    "# Define the metrics that will be logged for each model output.\n",
-    "# We will log the Jaccard Index for the type maps and the Mean Squared Error\n",
-    "# for the horizontal and vertical gradients output. For the type maps, the metric\n",
-    "# is the average computed over each class separately ('macro').\n",
-    "branch_metrics = {\n",
-    "    \"dist\": [None],\n",
-    "    \"type\": [\n",
-    "        JaccardIndex(\n",
-    "            task=\"multiclass\", average=\"macro\", num_classes=6, compute_on_cpu=True\n",
-    "        )\n",
-    "    ],\n",
-    "    \"stardist\": [MeanSquaredError(compute_on_cpu=True)],\n",
-    "    \"stardist_refined\": [MeanSquaredError(compute_on_cpu=True)],\n",
-    "}\n",
-    "\n",
     "# Initialize the optimizer.\n",
     "# We will be using the AdamW optimizer from the torch.optim library with learning rate of 0.0003.\n",
     "adamw = optim.AdamW\n",
@@ -511,7 +435,6 @@
     "experiment = SegmentationExperiment(\n",
     "    model=model,\n",
     "    multitask_loss=multitask_loss,\n",
-    "    branch_metrics=branch_metrics,\n",
     "    optimizer=adamw,\n",
     "    scheduler=scheduler,\n",
     "    optimizer_kwargs=optimizer_kwargs,\n",
@@ -560,7 +483,7 @@
     "    monitor=\"val_loss\",\n",
     "    mode=\"min\",\n",
     ")\n",
-    "ckpt_callback.CHECKPOINT_NAME_LAST = f\"cppnet_last\"\n",
+    "ckpt_callback.CHECKPOINT_NAME_LAST = \"cppnet_last\"\n",
     "callbacks.append(ckpt_callback)\n",
     "\n",
     "# Lightning training\n",
@@ -1002,14 +925,14 @@
     "inferer = ResizeInferer(\n",
     "    model=experiment,\n",
     "    input_path=save_dir / \"valid\" / \"images\",\n",
-    "    checkpoint_path=Path.home() / \"pannuke\" / \"cppnet\" / \"cppnet_last.ckpt\",\n",
     "    out_activations=out_acts,\n",
     "    out_boundary_weights=out_weights,\n",
     "    resize=(256, 256),  # Not actually resizing anything,\n",
     "    instance_postproc=\"stardist\",\n",
     "    batch_size=8,\n",
     "    n_images=50,  # Use only the 50 first images of the folder,\n",
     "    normalization=\"percentile\",\n",
+    "    checkpoint_path=trainer.checkpoint_callback.best_model_path,\n",
     ")\n",
     "inferer.infer()"
    ]
@@ -1050,9 +973,10 @@
     }
    ],
    "source": [
+    "import matplotlib.patches as mpatches\n",
     "import numpy as np\n",
+    "\n",
     "from cellseg_models_pytorch.utils import draw_thing_contours\n",
-    "import matplotlib.patches as mpatches\n",
     "\n",
     "fig, ax = plt.subplots(5, 2, figsize=(10, 17))\n",
     "ax = ax.flatten()\n",
diff --git a/examples/pannuke_nuclei_segmentation_omnipose.ipynb b/examples/pannuke_nuclei_segmentation_omnipose.ipynb
@@ -23,7 +23,8 @@
    "outputs": [],
    "source": [
     "# !pip install cellseg-models-pytorch\n",
-    "# !pip install accelerate"
+    "# !pip install accelerate\n",
+    "# !pip install torchmetrics"
    ]
   },
   {
@@ -43,10 +44,12 @@
     }
    ],
    "source": [
-    "import torch\n",
+    "from platform import python_version\n",
+    "\n",
     "import accelerate\n",
+    "import torch\n",
+    "\n",
     "import cellseg_models_pytorch\n",
-    "from platform import python_version\n",
     "\n",
     "print(\"torch version:\", torch.__version__)\n",
     "print(\"accelerate version:\", accelerate.__version__)\n",
@@ -88,6 +91,7 @@
    ],
    "source": [
     "from pathlib import Path\n",
+    "\n",
     "from cellseg_models_pytorch.datamodules import PannukeDataModule\n",
     "\n",
     "fold_split = {\"fold1\": \"train\", \"fold2\": \"train\", \"fold3\": \"valid\"}\n",
@@ -148,13 +152,14 @@
     }
    ],
    "source": [
-    "import numpy as np\n",
     "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
     "from skimage.color import label2rgb\n",
     "\n",
+    "from cellseg_models_pytorch.transforms.functional import gen_omni_flow_maps\n",
+    "\n",
     "# filehandler contains methods to read and write images and masks\n",
     "from cellseg_models_pytorch.utils import FileHandler\n",
-    "from cellseg_models_pytorch.transforms.functional import gen_omni_flow_maps\n",
     "\n",
     "img_dir = save_dir / \"train\" / \"images\"\n",
     "mask_dir = save_dir / \"train\" / \"labels\"\n",
@@ -238,15 +243,14 @@
    "outputs": [],
    "source": [
     "import accelerate\n",
-    "from accelerate.utils import set_seed\n",
-    "\n",
     "import torch.nn as nn\n",
     "import torch.nn.functional as F\n",
+    "from accelerate.utils import set_seed\n",
     "from torchmetrics import JaccardIndex, MeanSquaredError\n",
     "from tqdm import tqdm\n",
     "\n",
+    "from cellseg_models_pytorch.losses import CELoss, DiceLoss, JointLoss, MultiTaskLoss\n",
     "from cellseg_models_pytorch.models import omnipose_base\n",
-    "from cellseg_models_pytorch.losses import DiceLoss, CELoss, JointLoss, MultiTaskLoss\n",
     "\n",
     "\n",
     "# Quick wrapper for MAE loss to make it fit the JointLoss API\n",
@@ -961,10 +965,10 @@
     }
    ],
    "source": [
-    "import numpy as np\n",
-    "from cellseg_models_pytorch.utils import draw_thing_contours\n",
     "import matplotlib.patches as mpatches\n",
+    "import numpy as np\n",
     "\n",
+    "from cellseg_models_pytorch.utils import draw_thing_contours\n",
     "\n",
     "fig, ax = plt.subplots(5, 2, figsize=(10, 17))\n",
     "ax = ax.flatten()\n",
