UHNSeg-Quiz

U-Net and Vision Transformer for automatic blood cell segmentation

Installation

This repository is built on nnUNet. The installation contains two simple steps

Install latest pytorch
git clone https://github.com/JunMa11/UHNSeg-Quiz
Enter the folder UHNSeg-Quiz and run pip install -e .

Data preparation and preprocessing

Please organize the images and masks as follows:

data/nnUNet_raw/
├── Dataset705_SingleBloodCell
	├── imagesTr
		├── img_000_0000.png
		├── img_001_0000.png
		├── ...
	├── labelsTr
		├── img_000.png
		├── img_001.png
		├── ...
	├── imagesTs-Internal
		├── Internal_img_000_0000.png
		├── Internal_img_001_0000.png
		├── ...
	├── imagesTs-External

		├── External_img_000_0000.png
		├── External_img_001_0000.png
		├── ...
	├── dataset.json

Set the path

Run

nnUNetv2_plan_and_preprocess -d 705

Model training

To train the 2D U-Net model, run:

nnUNetv2_train 705 2d all

To train the 2D Efficient ViT model, run:

nnUNetv2_train 705 2d all -tr nnUNetTrainerEffViTl1Seg

Model Inference

Download the model checkpoints to data/nnUNet_results/Dataset705_SingleBloodCell

Inference the testing images with 2D U-Net model, run:

# internal validation
nnUNetv2_predict -i data/nnUNet_raw/Dataset705_SingleBloodCell/imagesTs-Internal -o data/nnUNet_raw/Dataset705_SingleBloodCell/UNet-seg-imagesTs-Internal -d 705 -c 2d -f all 

# External validation
nnUNetv2_predict -i data/nnUNet_raw/Dataset705_SingleBloodCell/imagesTs-External -o data/nnUNet_raw/Dataset705_SingleBloodCell/UNet-seg-imagesTs-External -d 705 -c 2d -f all -tr nnUNetTrainerEffViTl1Seg

To train the 2D Efficient ViT model, run:

# internal validation
nnUNetv2_predict -i data/nnUNet_raw/Dataset705_SingleBloodCell/imagesTs-Internal -o data/nnUNet_raw/Dataset705_SingleBloodCell/ViT-seg-imagesTs-Internal -d 705 -c 2d -f all 

# External validation
nnUNetv2_predict -i data/nnUNet_raw/Dataset705_SingleBloodCell/imagesTs-External -o data/nnUNet_raw/Dataset705_SingleBloodCell/ViT-seg-imagesTs-External -d 705 -c 2d -f all -tr nnUNetTrainerEffViTl1Seg

Visualized segmentation results

Acknowledgements

The implementation is built on nnUNet and EfficientViT. Thanks for making the source code publicly available.

JunMa11 / UHNSeg-Quiz