docs: mod readme, fix notebooks

Author: okunator

README.md

@@ -24,11 +24,21 @@
 
 ## Introduction
 
-Contains multi-task encoder-decoder architectures along with dedicated post-processing methods for segmenting cell/nuclei instances. As the name suggests, this library is heavily inspired by [segmentation_models.pytorch](https://github.com/qubvel/segmentation_models.pytorch) library for semantic segmentation.
+**cellseg-models.pytorch** is a library built upon [PyTorch](https://pytorch.org/) that contains multi-task encoder-decoder architectures along with dedicated post-processing methods for segmenting cell/nuclei instances. As the name might suggest, this library is heavily inspired by [segmentation_models.pytorch](https://github.com/qubvel/segmentation_models.pytorch) library for semantic segmentation.
 
-<br><br>
+## Features
 
-![Architecture](./images/architecture_overview.png)
+- High level API to define cell/nuclei instance segmentation models.
+- 4 cell/nuclei instance segmentation models and more to come.
+- Open source datasets for training and benchmarking.
+- Pre-trained backbones/encoders from the [timm](https://github.com/rwightman/pytorch-image-models) library.
+- All the architectures can be augmented to **panoptic segmentation**.
+- A lot of flexibility to modify the components of the model architectures.
+- Sliding window inference for large images.
+- Multi-GPU inference.
+- Popular training losses and benchmarking metrics.
+- Simple model training with [pytorch-lightning](https://www.pytorchlightning.ai/).
+- Benchmarking utilities both for model latency & segmentation performance.
 
 ## Installation
 
@@ -44,17 +54,6 @@ pip install cellseg-models-pytorch
 pip install cellseg-models-pytorch[all]
 ```
 
-## Features
-
-- High level API to define cell/nuclei instance segmentation models.
-- 4 cell/nuclei instance segmentation models and more to come.
-- Pre-trained backbones/encoders from the [timm](https://github.com/rwightman/pytorch-image-models) library.
-- All the architectures can be augmented to output semantic segmentation outputs along with instance semgentation outputs (panoptic segmentation).
-- A lot of flexibility to modify the components of the model architectures.
-- Multi-GPU inference.
-- Popular training losses and benchmarking metrics.
-- Simple model training with [pytorch-lightning](https://www.pytorchlightning.ai/).
-
 ## Models
 
 | Model                      | Paper                                                                          |
@@ -109,10 +108,10 @@ y = model(x) # {"cellpose": [1, 2, 256, 256], "type": [1, 5, 256, 256], "sem": [
 ```python
 import cellseg_models_pytorch as csmp
 
-# two decoder branches.
+# the model will include two decoder branches.
 decoders = ("cellpose", "sem")
 
-# three segmentation heads from the decoders.
+# and in total three segmentation heads emerging from the decoders.
 heads = {
     "cellpose": {"cellpose": 2, "type": 5},
     "sem": {"sem": 3}
@@ -148,32 +147,33 @@ y = model(x) # {"cellpose": [1, 2, 256, 256], "type": [1, 5, 256, 256], "sem": [
 ```python
 import cellseg_models_pytorch as csmp
 
+# define the model
 model = csmp.models.hovernet_base(type_classes=5)
-# returns {"hovernet": [B, 2, H, W], "type": [B, 5, H, W], "inst": [B, 2, H, W]}
 
-# the final activations for each model output
+# define the final activations for each model output
 out_activations = {"hovernet": "tanh", "type": "softmax", "inst": "softmax"}
 
-# models perform the poorest at the image boundaries, with overlapping patches this
-# causes issues which can be overcome by adding smoothing to the prediction boundaries
+# define whether to weight down the predictions at the image boundaries
+# typically, models perform the poorest at the image boundaries and with
+# overlapping patches this causes issues which can be overcome by down-
+# weighting the prediction boundaries
 out_boundary_weights = {"hovernet": True, "type": False, "inst": False}
 
-# Sliding window inference for big images using overlapping patches
+# define the inferer
 inferer = csmp.inference.SlidingWindowInferer(
     model=model,
     input_folder="/path/to/images/",
     checkpoint_path="/path/to/model/weights/",
     out_activations=out_activations,
     out_boundary_weights=out_boundary_weights,
-    instance_postproc="hovernet", # THE POST-PROCESSING METHOD
+    instance_postproc="hovernet",               # THE POST-PROCESSING METHOD
+    normalization="percentile",                 # same normalization as in training
     patch_size=(256, 256),
     stride=128,
     padding=80,
     batch_size=8,
-    normalization="percentile", # same normalization as in training
 )
 
-# Run sliding window inference.
 inferer.infer()
 
 inferer.out_masks
@@ -182,9 +182,14 @@ inferer.out_masks
 
 ## Models API
 
+Generally, the model building API enables the effortless creation of hard-parameter sharing multi-task encoder-decoder CNN architectures. The general architectural schema is illustrated in the below image.
+
+<br><br>
+![Architecture](./images/architecture_overview.png)
+
 ### Class API
 
-The class API enables the most flexibility in defining different model architectures. It allows for defining a multitude of hard-parameter sharing multi-task encoder-decoder architectures with (relatively) low effort. The class API is borrowing a lot from [segmentation_models.pytorch](https://github.com/qubvel/segmentation_models.pytorch) models API.
+The class API enables the most flexibility in defining different model architectures. It borrows a lot from [segmentation_models.pytorch](https://github.com/qubvel/segmentation_models.pytorch) models API.
 
 **Model classes**:
 

cellseg_models_pytorch/utils/file_manager.py

@@ -424,7 +424,7 @@ def save_masks(
                 save_dir = fname.parent / "cells"
                 if not Path(save_dir).exists():
                     Path(save_dir).mkdir(parents=True, exist_ok=True)
-                # print(save_dir)
+
                 fn = save_dir / f"{fname.name}_cells"
                 FileHandler.write_gson(
                     fname=fn,

examples/lizard_nuclei_segmentation_cellpose.ipynb

@@ -289,10 +289,9 @@
     "\n",
     "Other important params include: \n",
     "- `out_activations` - Sets the output activation functions for each of the model outputs\n",
-    "- `out_boundary_weights` - Sets whether we will use a weight matrix to add less weight to boundaries of the prediction of the image. This can only be useful when inference is run for bigger images that are patched in overlapping patches (this can be done with the `SlidingWindowInferer`).\n",
+    "- `out_boundary_weights` - Sets whether we will use a weight matrix to add less weight to boundaries of the predictions. This can only be useful when inference is run for bigger images that are patched in overlapping patches (inference with overlapping patches can be done with the `SlidingWindowInferer`).\n",
     "- `normalization` - Should be set to the same one as during training.\n",
-    "- `n_images` - Run inference only for the 3 first images of inside the input folder.\n",
-    "- `batch_size` -This needs to be set to 1 since the input images have different sizes and the dataloader can't stack them."
+    "- `n_images` - Run inference only for the 50 first images of inside the input folder."
    ]
   },
   {

examples/pannuke_nuclei_segmentation_stardist.ipynb

@@ -178,7 +178,7 @@
     "\n",
     "For the nuclei type masks we will monitor mIoU metric during training.\n",
     "\n",
-    "The optimizer used here is [AdamP](https://arxiv.org/abs/2006.08217)."
+    "The optimizer used here is [AdamW](https://arxiv.org/abs/1711.05101).."
    ]
   },
   {
@@ -206,7 +206,7 @@
     "    model=model,\n",
     "    branch_losses={\"dist\": \"ssim_mse\", \"stardist\": \"ssim_mse\", \"type\": \"ce_dice\"},\n",
     "    branch_metrics={\"dist\": [None], \"stardist\":[None], \"type\": [\"miou\"]},\n",
-    "    optimizer=\"adamp\",\n",
+    "    optimizer=\"adamw\",\n",
     "    lookahead=False,\n",
     ")\n",
     "\n",
@@ -279,12 +279,11 @@
     "\n",
     "Other important params include: \n",
     "- `out_activations` - Sets the output activation functions for each of the model outputs\n",
-    "- `out_boundary_weights` - Sets whether we will use a weight matrix to add less weight to boundaries of the prediction of the image. This can only be useful when inference is run for bigger images that are patched in overlapping patches (this can be done with the `SlidingWindowInferer`).\n",
+    "- `out_boundary_weights` - Sets whether we will use a weight matrix to add less weight to boundaries of the predictions. This can only be useful when inference is run for bigger images that are patched in overlapping patches (inference with overlapping patches can be done with the `SlidingWindowInferer`).\n",
     "- `normalization` - Should be set to the same one as during training.\n",
     "- `n_images` - Run inference only for the 50 first images of inside the input folder.\n",
-    "- `use_mask` - Use a mask to get cell type classifications for only to the same pixels as in the instance segmentation.\n",
     "\n",
-    "**NOTE**: Another important thing to note here, is that the `\"stardist\"` post-proc method is not the original one introduced in the [Stardist](https://github.com/stardist/stardist) but rather a workaround that's is not as optimal as the original one but still does the job. You can use the original by setting it to `\"stardist_orig\"`, however, this requires also the original `stardist` library that can be installed with `pip install stardist`."
+    "**NOTE**: Another important thing to note here, is that the `\"stardist\"` post-proc method is not the original one introduced in the [Stardist](https://github.com/stardist/stardist) paper. It is a python rewrite of the original one and can be even twice as fast as the orig one with only neglible differneces in the output. However, if you like, you can use the original by setting `instance_postproc` to `\"stardist_orig\"`. It should be noted that the original version requires also the original `stardist` library that can be installed with `pip install stardist`."
    ]
   },
   {
@@ -304,14 +303,13 @@
     "inferer = csmp.inference.ResizeInferer(\n",
     "    model=experiment,\n",
     "    input_folder=save_dir / \"test\" / \"images\",\n",
-    "    out_activations={\"dist\": \"tanh\", \"stardist\": None, \"type\": \"softmax\"},\n",
+    "    out_activations={\"dist\": None, \"stardist\": None, \"type\": \"softmax\"},\n",
     "    out_boundary_weights={\"dist\": False, \"stardist\": False, \"type\": False},\n",
     "    resize=(256, 256), # Not actually resizing anything,\n",
     "    instance_postproc=\"stardist\",\n",
     "    save_intermediate=True, # save intermediate soft masks for visualization\n",
     "    normalization=\"percentile\", # same normalization as during training\n",
     "    batch_size=8,\n",
-    "    use_mask=True,\n",
     "    n_images=50 # Use only the 50 first images of the folder\n",
     ")\n",
     "inferer.infer()"