YOLOv8-AM for Fracture Detection

YOLOv8-AM: YOLOv8 with Attention Mechanisms for Pediatric Wrist Fracture Detection

Architecture

Performance

Citation

If you find our paper useful in your research, please consider citing:

@article{chien2024yolov8am,
  title={YOLOv8-AM: YOLOv8 with Attention Mechanisms for Pediatric Wrist Fracture Detection},
   author={Chun-Tse Chien and Rui-Yang Ju and Kuang-Yi Chou and Enkaer Xieerke and Jen-Shiun Chiang},
  journal={arXiv preprint arXiv:2402.09329},
  year={2024}
}

Requirements

• Linux (Ubuntu)
• Python = 3.9
• Pytorch = 1.13.1
• NVIDIA GPU + CUDA CuDNN

Environment

  pip install -r requirements.txt

Dataset

Download the dataset

• You can download the GRAZPEDWRI-DX Dataset on this Link.

Split the dataset

• To split the dataset into training set, validation set, and test set, you should first put the image and annotatation into ./GRAZPEDWRI-DX/data/images, and ./GRAZPEDWRI-DX/data/labels.

• And then you can split the dataset as the following step:

    python split.py
  

• The dataset is divided into training, validation, and testing set (70-20-10 %) according to the key patient_id stored in dataset.csv. The script then will move the files into the relative folder as it is represented here below.

   GRAZPEDWRI-DX
      └── data   
           ├── meta.yaml
           ├── images
           │    ├── train
           │    │    ├── train_img1.png
           │    │    └── ...
           │    ├── valid
           │    │    ├── valid_img1.png
           │    │    └── ...
           │    └── test
           │         ├── test_img1.png
           │         └── ...
           └── labels
                ├── train
                │    ├── train_annotation1.txt
                │    └── ...
                ├── valid
                │    ├── valid_annotation1.txt
                │    └── ...
                └── test
                     ├── test_annotation1.txt
                     └── ...
  

The script will create 3 files: train_data.csv, valid_data.csv, and test_data.csv with the same structure of dataset.csv.

Data Augmentation

• Data augmentation of the training set using the addWeighted function doubles the size of the training set.

  python imgaug.py --input_img /path/to/input/train/ --output_img /path/to/output/train/ --input_label /path/to/input/labels/ --output_label /path/to/output/labels/

For example:

  python imgaug.py --input_img ./GRAZPEDWRI-DX/data/images/train/ --output_img ./GRAZPEDWRI-DX/data/images/train_aug/ --input_label ./GRAZPEDWRI-DX/data/labels/train/ --output_label ./GRAZPEDWRI-DX/data/labels/train_aug/

• The path of the processed file is shown below:

   GRAZPEDWRI-DX
      └── data   
           ├── meta.yaml
           ├── images
           │    ├── train
           │    │    ├── train_img1.png
           │    │    └── ...
           │    ├── train_aug
           │    │    ├── train_aug_img1.png
           │    │    └── ...
           │    ├── valid
           │    │    ├── valid_img1.png
           │    │    └── ...
           │    └── test
           │         ├── test_img1.png
           │         └── ...
           └── labels
                ├── train
                │    ├── train_annotation1.txt
                │    └── ...
                ├── train_aug
                │    ├── train_aug_annotation1.txt
                │    └── ...
                ├── valid
                │    ├── valid_annotation1.txt
                │    └── ...
                └── test
                     ├── test_annotation1.txt
                     └── ...
  

Methodology

• We have modified the model architecture of YOLOv8 by adding four types of attention modules, including Shuffle Attention (SA), Efficient Channel Attention (ECA), Global Attention Mechanism (GAM), and ResBlock Convolutional Block Attention Module (ResCBAM).

Train & Validate

• We have provided a training set, test set and validation set containing a single image that you can run directly by following the steps in the example below.
• Before training the model, make sure the path to the data in the ./GRAZPEDWRI-DX/data/meta.yaml file is correct.

  # patch: /path/to/GRAZPEDWRI-DX/data
  path: 'E:/GRAZPEDWRI-DX/data'
  train: 'images/train_aug'
  val: 'images/valid'
  test: 'images/test'

• Arguments

You can set the value in the ./ultralytics/cfg/default.yaml.

• Example Train & Val Steps (yolov8m):

  python start_train.py --model ./ultralytics/cfg/models/v8/yolov8m.yaml --data_dir ./GRAZPEDWRI-DX/data/meta.yaml

• Example Train & Val Steps (yolov8m_ECA):

  python start_train.py --model ./ultralytics/cfg/models/v8/yolov8m_ECA.yaml --data_dir ./GRAZPEDWRI-DX/data/meta.yaml

Related Works