/breast_ultrasound_lesion_segmentation_PyTorch

PyTorch implementation of segmentation models for breast ultrasound images.

Primary LanguagePythonMIT LicenseMIT

Breast Ultrasound Segmentation

This repository is no longer actively maintained.

Easy implementation for different models in PyTorch for breast ultrasound lesion segmentation.

Platform

Main expectation is to enable better model training & visualization in Google Colab, while the scripts can be employed for other GPU available environments.

Training Dataset

Breast Ultrasound Images Dataset (Dataset BUSI) Dataset Public Access

  • Number of patients: 600 females.
  • Number of images: 780 PNG images with an average size of 500*500.
  • Classes: Normal, Benign, Malignant.

Both ground truth masks and original images are provided.

Normal images are not used during training.

For preprocessing, you should rearrange the data into two separate folders "images" and "mask_ground_truth".

Models (work in progress)

  • UNet with ResNet backbone (ResUNet)
  • ResUNet with attention
  • Attention UNet
  • Multi-scale Attention Net (MA-Net)

General Workflow

Option 1:

Check train_with_Colab.ipynb and follow the instructions inside.

Option 2:

  1. To download the BUSI dataset, run python BUSI/BUSI_prepare_for_training.py. You can also prepare it differently.

  2. Change configuration settings in options/<your_config_filename>.json or use the default json file.

  3. To enable Tensorboard visualization, make sure tensorboard option is turned on in the config file. "tensorboard: true" At the project root, open Tensorboard server:

    mkdir exp_results
    mkdir exp_results/log
    tensorboard --logdir exp_results/log/
    

    The server will then open at http://localhost:6006

  4. To start training, run

    python train.py --config options/<your_config_filename>.json 
    --device 'index to GPU device' 
    --resume 'path/to/latest/checkpoint' or None
    

References

Template: PyTorch Template Project https://github.com/victoresque/pytorch-template#pytorch-template-project Library: Segmentation Models https://smp.readthedocs.io/en/latest/