UniCell: Universal Cell Nucleus Classification via Prompt Learning, AAAI 2024

-python=3.8

-pytorch=1.12.0

-torchvision=0.13.0

step 0. Install mmcv and mmengine using mim

pip install -U openmim
mim install mmengine==0.7.2
mim install mmcv==2.0.0
pip install timm==0.6.13
pip install ftfy==6.1.1
pip install regex==2023.3.23
pip install einops==0.6.0

step 1. Git clone the repository

git clone https://github.com/lhaof/UniCell.git

step 2. Install UniCell

cd UniCell/mmdetection
python -m pip install -e .

step 3. Install SAHI

cd projects/UniCell/tools/sahi
python -m pip install -e .

step 0. Download the dataset from google drive and unzip it to UniCell/dataset/.

step 1. Modify the dataset_path in projects/UniCell/tools/prepare_fourdataset_4Dataset_CMOL.py and run it to generate the dataset.

python prepare_fourdataset_4Dataset_CMOL.py

Modify the dataset_path in projects/UniCell/configs/nuclei_det_multihead_cmol.py and run the following command to train the model.

python ./tools/train.py tools/train.py projects/UniCell/configs/UniCell_CMOL.py\
	--work-dir=${SAVE_DIR}

bash ./tools/dist_train.sh projects/UniCell/configs/UniCell_CMOL.py\
    ${GPU_NUM} --work-dir=${SAVE_DIR}

step 0. Download the universal model from google drive.

step 1. Modify the path_to_dataset and checkpoint path in projects/UniCell/tools/inference_multihead.py and run it to test the model.

python inference_multihead.py

We use the entire training set for training, and use the final model for testing after completing 160k iterations.

step 0. Transfer your dataset into the format of HoverNet CoNSeP and put it in the UniCell/dataset/ directory.

step 1. Modify the dataset_path, categories, datasets in projects/UniCell/tools/prepare_fourdataset_4Dataset_CMOL.py and run it to generate the dataset.

Noted that the dataset name in trans_to_patch function should be the same as the dataset name in datasets.

step 2. Modify the METAINFO in projects/UniCell/configs/nuclei_det_multihead_cmol.py and make sure that the num_classes is equal to the number of categories in your dataset.

(You can register a new dataset type instead if you are familiar with mmdetection.)

step 3. Modify the num_classes in projects/UniCell/configs/UniCell_CMOL.py.

step 4. Modify the dataset_names and category_names in projects/UniCell/configs/UniCell_CMOL.py (text_cfg).

Modify the mask_map, the key corresponds to the dataset, and the value corresponds to the category that requires mask.

step 5. Training. Follow the steps in Training.

We thank the following projects for their valuable contributions to this work.

• MMDetection
• SAHI