LKCell🔬

使用大卷积核的高效细胞核实例分割

Ziwei Cui ^1*, Jingfeng Yao ^1*, Lunbin Zeng ¹, Juan Yang ², Wenyu Liu ¹, Xinggang Wang ^1,📧

¹ School of Electronic Information and Communications, Huazhong University of Science and Technology
² Department of Cardiology, Huanggang Central Hospital

(* 同等贡献, 📧 通讯作者)

主要特性 • 安装 • 使用 • 训练 • 推理 • 引用

主要特性

点击即可在 🤗 Hugging Face Space 上试用 LKCell！

本仓库包含 LKCell 的代码实现，这是一种基于深度学习的自动化细胞核实例分割方法，用于数字化组织样本。LKCell 采用基于大卷积核的架构，并在 PanNuke 数据集上取得了最先进的性能，该数据集是一个具有挑战性的细胞核实例分割基准。

安装

git clone https://github.com/hustvl/LKCell.git
conda create -n lkcell
conda activate lkcell
pip install -r requirements.txt

注意：(1) 首选 torch 版本为 2.0；(2) 如果在安装 depthwise-conv2d-implicit-gemm==0.0.0 时遇到问题，请遵循此处的说明。

项目结构

我们目前使用以下文件夹结构：

├── base_ml               # Basic Machine Learning Code: 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_cellvit.py    # Run file to start an experiment
├── config                # Python configuration file for global Python settings            
├── docs                  # Documentation files (in addition to this main README.md
├── models                # Machine Learning Models (PyTorch implementations)
│   └── segmentation      # LKCell Code
├── datamodel             # Code of dataclass :Graph Data , WSI object , ...
├── preprocessing         # Code of preprocessing : Encoding , Patch Extraction , ...

训练

数据集准备

我们为 PanNuke 和 MoNuSeg 数据集使用了自定义的数据集结构。数据集结构在 pannuke.md 和 monuseg.md 文档文件中进行了说明。我们还在 cell_segmentation/datasets/ 文件夹中提供了准备脚本。

训练脚本

用于训练 LKCell-Network 的 ML 实验的 CLI 如下（此处使用了 run_cellvit.py 脚本）：

usage: run_cellvit.py [-h] --config CONFIG [--gpu GPU] [--sweep | --agent AGENT | --checkpoint CHECKPOINT]
Start an experiment with given configuration file.

python ./cell_segmentation/run_cellvit.py  --config ./config.yaml

重要的文件是配置文件，其中设置了所有路径，给出了模型配置，并定义了超参数或扫描。

预训练的 UnirepLKNet 模型 可用于训练初始化，可从 Google Drive 下载：UnirepLKNet-Models。

评估

在我们的论文中，我们没有（!）使用早停，而是将所有模型训练 100 个 epoch 以消除选择偏差，但拥有尽可能大的训练数据库。因此，评估需要使用 latest_checkpoint.pth 模型，而不是最佳早停模型。我们提供了一个脚本来创建评估结果：PanNuke 使用 inference_cellvit_experiment.py，MoNuSeg 使用 inference_cellvit_monuseg.py。

推理

模型检查点可在此处下载，您可以选择从 Google Drive 或 HuggingFace 下载：

Google Drive
- LKCell-L 🚀
- LKCell-B
HuggingFace
- LKCell-L 🚀
- LKCell-B

您可以点击 🤗 Hugging Face Space 快速执行模型推理。

致谢

本项目基于 CellViT 和 UniRepLKNet 构建。感谢这些出色的仓库！

引用

@misc{cui2024lkcellefficientcellnuclei,
      title={LKCell: Efficient Cell Nuclei Instance Segmentation with Large Convolution Kernels}, 
      author={Ziwei Cui and Jingfeng Yao and Lunbin Zeng and Juan Yang and Wenyu Liu and Xinggang Wang},
      year={2024},
      eprint={2407.18054},
      archivePrefix={arXiv},
      primaryClass={eess.IV},
      url={https://arxiv.org/abs/2407.18054}, 
}