Official CST-YOLO

This is the source code for the paper, "CST-YOLO: A Novel Method for Blood Cell Detection Based on Improved YOLOv7 and CNN-Swin Transformer", of which I am the first author.

Model

The model configuration (i.e., network construction) file is cst-yolo.yaml in the directory ./cfg/training/.

Recommended dependencies:

Python <= 3.8
Torch <= 1.7.1
CUDA <= 11.1

Training

The hyperparameter setting file is hyp.scratch.p5.yaml in the directory ./data/.

Single GPU training

python train.py --workers 8 --device 0 --batch-size 32 --data data/cbc.yaml --img 640 640 --cfg cfg/training/cst-yolo.yaml --weights '' --name cst-yolo --hyp data/hyp.scratch.p5.yaml

Multiple GPU training

python -m torch.distributed.launch --nproc_per_node 4 --master_port 9527 train.py --workers 8 --device 0,1,2,3 --sync-bn --batch-size 128 --data data/cbc.yaml --img 640 640 --cfg cfg/training/cst-yolo.yaml --weights '' --name rcs-yolo --hyp data/hyp.scratch.p5.yaml

Testing

python test.py --data data/cbc.yaml --img 640 --batch 32 --conf 0.001 --iou 0.65 --device 0 --weights runs/train/exp/weights/best.pt --name val

Evaluation

We trained and evaluated CST-YOLO on three blood cell detection datasets Blood Cell Count and Detection (BCCD), Complete Blood Count (CBC), and Blood Cell Detection (BCD). The 60 samples of the validation set duplicate those from the training set in the CBC dataset. Each image includes three types of blood cells: Red Blood Cells (RBCs), White Blood Cells (WBCs), and platelets.

Table 1    Number of examples in BCCD, CBC, and BCD.

Table 2.1    Performance comparison of YOLOv5x, YOLOv7 and CST-YOLO on the BCCD dataset. Results are [email protected] for each blood cell type and overall performance. The best results are shown in bold.

Table 2.2    Performance comparison of YOLOv5x, YOLOv7 and CST-YOLO on the CBC dataset. Results are [email protected] for each blood cell type and overall performance. The best results are shown in bold.

Table 2.3    Performance comparison of YOLOv5x, YOLOv7 and CST-YOLO on the BCD dataset. Results are [email protected] for each blood cell type and overall performance. The best results are shown in bold.

Ablation Study

The tables below show the effect on the performance of the proposed modules for the three blood cell datasets.

Table 3.1    Abation study of the proposed modules. Results are [email protected] on the BCCD dataset. The best results are shown in bold.

Table 3.2    Abation study of the proposed modules. Results are [email protected] on the CBC dataset. The best results are shown in bold.

Table 3.3    Abation study of the proposed modules. Results are [email protected] on the BCD dataset. The best results are shown in bold.

Generalizability Study

Table 4 compares the performance between the original YOLOv7 and CST-YOLO on the TinyPerson dataset in different domain from medical images for external validation. The experimental results demonstrate the generalizability effectivenss of CST-YOLO in small object detection on both medical (i.e., blood cell) and natural (i.e., tiny person) images.

Table 4    Performance comparison of YOLOv7 and CST-YOLO for all classes of the TinyPerson dataset. The best results are shown in bold.

Visual Illustrations for Model Explanation

Some heat maps are generated by Grad-CAM from the BCCD dataset. Examples of heat maps and failure cases are shown in Fig. 1 and Fig. 2, respectively.

Below is one example of failure cases because there aren't any Grad-CAM heatmaps solely generated for platelets. The model ability of dim object detection need be enhanced for future improvements.

Suggested Citation

Our manuscript has been uploaded on arXiv. Please cite our paper if you use code from this repository:

Plain Text

• IEEE Style M. Kang, C.-M. Ting, F. F. Ting, and R. C.-W. Phan, "Cst-yolo: A novel method for blood cell detection based on improved yolov7 and cnn-swin transformer," arXiv:2306.14590 [cs.CV], Jun. 2023.

• Nature Style Kang, M., Ting, C.-M., Ting, F. F. & Phan, R. C.-W. CST-YOLO: a novel method for blood cell detection based on improved YOLOv7 and CNN-swin transformer. Preprint at https://arxiv.org/abs/2306.14590 (2023).

• Springer Style Kang, M., Ting, C.-M., Ting, F. F., Phan, R.C.-W.: CST-YOLO: a novel method for blood cell detection based on improved YOLOv7 and CNN-swin transformer. arXiv preprint arXiv:2306.14590 (2023)

License

CST-YOLO is released under the GNU General Public License v3.0. Please see the LICENSE file for more information.

Copyright Notice

Many utility codes of our project base on the codes of YOLOv7 and Swin Transformer repositories.