CellViT: Vision Transformers for Precise Cell Segmentation and Classification

Key Features • Installation • Usage • Training • Inference • Examples • Roadmap • Citation

Update 08.08.2023:

:bangbang: We fixed a severe training bug and uploaded new checkpoints. Please make sure to pull all changes and redownload your CellViT checkpoints to get the best results :bangbang:

:ballot_box_with_check: Improved reproducability by providing config and log files for best models (CellViT-SAM-H and CellViT-256) and adopted PanNuke inference script for an easier evaluation

:ballot_box_with_check: Inference speed improved by x100 for postprocessing, added new preprocessing with CuCIM speedup

:ballot_box_with_check: Fixed bug in postprocessing that may insert doubled cells during cell-detection

:ballot_box_with_check: Added batch-size and mixed-precision options to inference cli to support RAM limited GPUs

:ballot_box_with_check: Extended configuration and added sweep configuration

Hörst, F., Rempe, M., Heine, L., Seibold, C., Keyl, J., Baldini, G., Ugurel, S., Siveke, J., Grünwald, B., Egger, J., & Kleesiek, J. (2023). CellViT: Vision Transformers for precise cell segmentation and classification. https://doi.org/10.48550/ARXIV.2306.15350

This repository contains the code implementation of CellViT, a deep learning-based method for automated instance segmentation of cell nuclei in digitized tissue samples. CellViT utilizes a Vision Transformer architecture and achieves state-of-the-art performance on the PanNuke dataset, a challenging nuclei instance segmentation benchmark.

If you intend to use anything from this repo, citation of the original publication given above is necessary

Key Features

State-of-the-Art Performance: CellViT outperforms existing methods for nuclei instance segmentation by a substantial margin, delivering superior results on the PanNuke dataset:
- Mean panoptic quality: 0.51
- F1-detection score: 0.83
Vision Transformer Encoder: The project incorporates pre-trained Vision Transformer (ViT) encoders, which are known for their effectiveness in various computer vision tasks. This choice enhances the segmentation performance of CellViT.
U-Net Architecture: CellViT adopts a U-Net-shaped encoder-decoder network structure, allowing for efficient and accurate nuclei instance segmentation. The network architecture facilitates both high-level and low-level feature extraction for improved segmentation results.
Weighted Sampling Strategy: To enhance the performance of CellViT, a novel weighted sampling strategy is introduced. This strategy improves the representation of challenging nuclei instances, leading to more accurate segmentation results.
Fast Inference on Gigapixel WSI: The framework provides fast inference results by utilizing a large inference patch size of $1024 \times 1024$ pixels, in contrast to the conventional $256$-pixel-sized patches. This approach enables efficient analysis of Gigapixel Whole Slide Images (WSI) and generates localizable deep features that hold potential value for downstream tasks. We provide a fast inference pipeline with connection to current Viewing Software such as QuPath

Visualization

Example

Installation

Clone the repository: git clone https://github.com/TIO-IKIM/CellViT.git
Create a conda environment with Python 3.9.7 version and install conda requirements: conda env create -f environment.yml. You can change the environment name by editing the name tag in the environment.yaml file. This step is necessary, as we need to install Openslide with binary files. This is easier with conda. Otherwise, installation from source needs to be performed and packages installed with pi
Activate environment: conda activate cellvit_env
Install torch (>=2.0) for your system, as described here. Preferred version is 2.0, see optional_dependencies for help. You can find all version here: https://pytorch.org/get-started/previous-versions/
Install optional dependencies pip install -r optional_dependencies.txt to get a speedup using NVIDIA-Clara and CuCIM for preprocessing during inference. Please select your CUDA versions. Help for installing cucim can be found online. Note Error: cannot import name CuImage from cucim If you get this error, install cucim from conda to get all binary files. First remove your previous dependeny with pip uninstall cupy-cuda117 and reinstall with conda install -c rapidsai cucim inside your conda environment. This process is time consuming, so you should be patient. Also follow their official guideline.

FAQ: Environment problems

ResolvePackageNotFound: -gcc

Fix: Comment out the gcc package in the environment.yml file

ResolvePackageNotFound: -libtiff==4.5.0=h6adf6a1_2, -openslide==3.4.1=h7773abc_6

Fix: Remove the version hash from environment.yml file, such that:

...
dependencies:
  ...
  - libtiff=4.5.0
  - openslide=3.4.1

pip:
...

PyDantic Validation Errors for the CLI

Please install the pydantic version specified (pydantic==1.10.4), otherwise validation errors could occur for the CLI.

Usage:

Project Structure

We are currently using the following folder structure:

├── base_ml               # Basic Machine Learning Code: CLI, Trainer, Experiment, ...
├── cell_segmentation     # Cell Segmentation training and inference files
│   ├── datasets          # Datasets (PyTorch)
│   ├── experiments       # Specific Experiment Code for different experiments
│   ├── inference         # Inference code for experiment statistics and plots
│   ├── trainer           # Trainer functions to train networks
│   ├── utils             # Utils code
│   └── run_xxx.py        # Run file to start an experiment
├── configs               # Config files
│   ├── examples          # Example config files with explanations
│   └── python            # Python configuration file for global Python settings
├── datamodel             # Datamodels of WSI, Patientes etc. (not ML specific)
├── docs                  # Documentation files (in addition to this main README.md)
├── models                # Machine Learning Models (PyTorch implementations)
│   ├── encoders          # Encoder networks (see ML structure below)
│   ├── pretrained        # Checkpoint of important pretrained models (needs to be downloaded from Google drive)
│   └── segmentation      # CellViT Code
├── preprocessing         # Preprocessing code
│   └── patch_extraction  # Code to extract patches from WSI

Training

The CLI for a ML-experiment to train the CellViT-Network is as follows (here the run_cellvit.py script is used):

usage: run_cellvit.py [-h] --config CONFIG [--gpu GPU] [--sweep | --agent AGENT | --checkpoint CHECKPOINT]

Start an experiment with given configuration file.

optional arguments:
  -h, --help            show this help message and exit
  --gpu GPU             Cuda-GPU ID (default: None)
  --sweep               Starting a sweep. For this the configuration file must be structured according to WandB sweeping. Compare
                        https://docs.wandb.ai/guides/sweeps and https://community.wandb.ai/t/nested-sweep-configuration/3369/3 for further
                        information. This parameter cannot be set in the config file! (default: False)
  --agent AGENT         Add a new agent to the sweep. Please pass the sweep ID as argument in the way entity/project/sweep_id, e.g.,
                        user1/test_project/v4hwbijh. The agent configuration can be found in the WandB dashboard for the running sweep in
                        the sweep overview tab under launch agent. Just paste the entity/project/sweep_id given there. The provided config
                        file must be a sweep config file.This parameter cannot be set in the config file! (default: None)
  --checkpoint CHECKPOINT
                        Path to a PyTorch checkpoint file. The file is loaded and continued to train with the provided settings. If this is
                        passed, no sweeps are possible. This parameter cannot be set in the config file! (default: None)

required named arguments:
  --config CONFIG       Path to a config file (default: None)

The important file is the configuration file, in which all paths are set, the model configuration is given and the hyperparameters or sweeps are defined. For each specific run file, there exists an example file in the ./configs/examples/cell_segmentation folder with the same naming as well as a configuration file that explains how to run WandB sweeps for hyperparameter search. All metrics defined in your trainer are logged to WandB. The WandB configuration needs to be set up in the configuration file, but also turned off by the user.

An example config file is given here with explanations here. For sweeps, we provide a sweep example file train_cellvit_sweep.yaml.

Pre-trained ViT models for training initialization can be downloaded from Google Drive: ViT-Models. Please check out the corresponding licenses before distribution and further usage! Note: We just used the teacher models for ViT-256.

:exclamation: If your training crashes at some point, you can continue from a checkpoint

Dataset preparation

We use a customized dataset structure for the PanNuke and the MoNuSeg dataset. The dataset structures are explained in pannuke.md and monuseg.md documentation files. We also provide preparation scripts in the cell_segmentation/datasets/ folder.

Evaluation

In our paper, we did not (!) use early stopping, but rather train all models for 130 to eliminate selection bias but have the largest possible database for training. Therefore, evaluation neeeds to be performed with the latest_checkpoint.pth model and not the best early stopping model. We provide to script to create evaluation results: inference_cellvit_experiment.py for PanNuke and inference_cellvit_monuseg.py for MoNuSeg.

:exclamation: We recently adapted the evaluation code and added a tag to the config files to select which checkpoint needs to be used. Please make sure to use the right checkpoint and select the appropriate dataset magnification.

Inference

Model checkpoints can be downloaded here:

License: Apache 2.0 with Commons Clause

Proved checkpoints have been trained on 90% of the data from all folds with the settings described in the publication.

Steps

The following steps are necessary for preprocessing:

Prepare WSI with our preprocessing pipeline
Run inference with the inference/cell_detection.py script

Results are stored at preprocessing locations

1. Preprocessing

In our Pre-Processing pipeline, we are able to extract quadratic patches from detected tissue areas, load annotation files (.json) and apply color normlizations. We make use of the popular OpenSlide library, but extended it with the RAPIDS cuCIM framework for an x8 speedup in patch-extraction. The documentation for the preprocessing can be found here.

Preprocessing is necessary to extract patches for our inference pipeline. We use squared patches of size 1024 pixels with an overlap of 64 px.

Please make sure that you select the following properties for our CellViT inference

Parameter	Value
patch_size	1024
patch_overlap	6.25

Resulting Dataset Structure

In general, the folder structure for a preprocessed dataset looks like this: The aim of pre-processing is to create one dataset per WSI in the following structure:

WSI_Name
├── annotation_masks      # thumbnails of extracted annotation masks
│   ├── all_overlaid.png  # all with same dimension as the thumbnail
│   ├── tumor.png
│   └── ...  
├── context               # context patches, if extracted
│   ├── 2                 # subfolder for each scale
│   │   ├── WSI_Name_row1_col1_context_2.png
│   │   ├── WSI_Name_row2_col1_context_2.png
│   │   └── ...
│   └── 4
│   │   ├── WSI_Name_row1_col1_context_2.png
│   │   ├── WSI_Name_row2_col1_context_2.png
│   │   └── ...
├── masks                 # Mask (numpy) files for each patch -> optional folder for segmentation
│   ├── WSI_Name_row1_col1.npy
│   ├── WSI_Name_row2_col1.npy
│   └── ...
├── metadata              # Metadata files for each patch
│   ├── WSI_Name_row1_col1.yaml
│   ├── WSI_Name_row2_col1.yaml
│   └── ...
├── patches               # Patches as .png files
│   ├── WSI_Name_row1_col1.png
│   ├── WSI_Name_row2_col1.png
│   └── ...
├── thumbnails            # Different kind of thumbnails
│   ├── thumbnail_mpp_5.png
│   ├── thumbnail_downsample_32.png
│   └── ...
├── tissue_masks          # Tissue mask images for checking
│   ├── mask.png          # all with same dimension as the thumbnail
│   ├── mask_nogrid.png
│   └── tissue_grid.png
├── mask.png              # tissue mask with green grid  
├── metadata.yaml         # WSI metdata for patch extraction
├── patch_metadata.json   # Patch metadata of WSI merged in one file
└── thumbnail.png         # WSI thumbnail

The cell detection and segmentation results are stored in a newly created cell_detection