CellViT

---

[!IMPORTANT]

If you just want to run inference, checkout the PyPI package and the corresponding GitHub repo

[!IMPORTANT]

This repository has been superseded by our new implementation!

Please visit our new repository at CellViT++. The new version includes significant improvements in performance, usability, and features.

---

---

CellViT: Vision Transformers for Precise Cell Segmentation and Classification

---

Update 10.01.2025:

CellViT++ is out now!

---

---

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

Installation

1. Clone the repository: git clone https://github.com/TIO-IKIM/CellViT.git

2. 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

3. Activate environment: conda activate cellvit_env

4. 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/

5. 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:

…

Training

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

…

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:

• CellViT-SAM-H
• CellViT-256
• CellViT-SAM-H-x20
• CellViT-256-x20

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:

1. Prepare WSI with our preprocessing pipeline
2. 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:

…

The cell detection and segmentation results are stored in a newly created cell_detection folder for each WSI.

2. Cell detection script

If the data is prepared, use the cell_detection.py script inside the cell_segmentation/inference folder to perform inference:

…

Single WSI

For processing a single WSI file, you need to select the process_wsi (python3 cell_detection.py process_wsi) subcommand with the following structure:

usage: cell_detection.py process_wsi --wsi_path WSI_PATH --patched_slide_path PATCHED_SLIDE_PATH

Process a single WSI file

arguments:
  -h, --help            show this help message and exit
  --wsi_path WSI_PATH   Path to WSI file
  --patched_slide_path PATCHED_SLIDE_PATH
                        Path to patched WSI file (specific WSI file, not parent path of patched slide dataset)

Multiple WSI

To process an entire dataset, select process_dataset (python3 cell_detection.py process_dataset):

…

3. Example

We provide an example TCGA file to show the performance and usage of our algorithms. Files and scripts can be found in the [example](example