The effectiveness of Deep Neural Networks (DNNs) heavily relies on the abundance and accuracy of available training data. However, collecting and annotating data on a large scale is often both costly and time-intensive, particularly in medical cases where practitioners are already occupied with their duties. Moreover, ensuring that the model remains robust across various scenarios of image capture is crucial in medical domains, especially when dealing with ultrasound images that vary based on the settings of different devices and the manual operation of the transducer.
To address this challenge, we introduce a novel pipeline called MEDDAP, which leverages Stable Diffusion (SD) models to augment existing small datasets by automatically generating new informative labeled samples. Pretrained checkpoints for SD are typically based on natural images, and training them for medical images requires significant GPU resources due to their heavy parameters. To overcome this challenge, we introduce USLoRA (Ultrasound Low-Rank Adaptation), a novel fine-tuning method tailored specifically for ultrasound applications. USLoRA allows for selective fine-tuning of weights within SD, requiring fewer than 0.1% of parameters compared to fully fine-tuning only the UNet portion of SD.
To enhance dataset diversity, we incorporate different adjectives into the generation process prompts, thereby desensitizing the classifiers to intensity changes across different images. This approach is inspired by clinicians' decision-making processes regarding breast tumors, where tumor shape often plays a more crucial role than intensity. In conclusion, our pipeline not only outperforms classifiers trained on the original dataset but also demonstrates superior performance when encountering unseen datasets.
Please refer to the paper for more technical details.
- Install necessary python libraries:
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .
cd examples/text_to_image
pip install -r requirements.txtaccelerate config -----> set the config based on your devicecreate account on hugging face and copy the access token for writing mode and paste the token when it is asked
pip install --upgrade huggingface_hub
huggingface-cli login Make Dataset in Hugging face webpage in your account with your preferred name Run the below code in terminal
python make_dataset_breast.py --train_dir <path_to_train_directory>\
--classes 'benign','malignant','normal'\
--prompt_structure "an ultrasound photo of {class_name} tumor in breast"\
--dataset_path <path_to_dataset_hugging_face>\
--token <writing_token_hugging_face>Remmember if you want to use this pipeline for different application, in the code, I change the structure of prompt for "normal' class. Instead of writing "an ultrasound photo of normal tumor in breast", I wrote "an ultrasound photo of no tumor in breast". This happens in "Make New Expanded Datasets" section too.
you can chnge the hyperparameters based on your problem but remember that it works better with batch size equal to 1.
export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export DATASET_NAME=<path_to_dataset_hugging_face>
accelerate launch --mixed_precision="fp16" train_text_to_image_lora.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--dataset_name=$DATASET_NAME --caption_column="text" \
--resolution=224 --random_flip \
--train_batch_size=1 \
--num_train_epochs=100 --checkpointing_steps=5000 \
--learning_rate=1e-04 --lr_scheduler="constant" --lr_warmup_steps=0 \
--seed=42 \
--output_dir=<path_of_checkpoints> \
--validation_prompt="an ultrasound photo of benign tumor in breast" --report_to="wandb" --rank 4
By running below script, you can train DCGAN for the dataset. Replace <path_to_folder_each_class> with the path of the folder consists of images of each class. <path_to_model_weights> is the path when your model is going to be saved.
class_name_list=("benign" "malignant" "normal")
for class_name in "${class_name_list[@]}"; do
# Run the Python script with the appropriate arguments
python dcgan_pytorch.py \
--data_path <path_to_folder_each_class> \
--dest_path <path_to_model_weights>
doneIn this section, we aim to generate new datasets with different adjectives and expansion ratios.
python Inference_percent_breast.py --model_path <path_of_checkpoints>\
--adjective_list "","colorful","stylized","high-contrast","low-contrast","posterized","solarized","sheared","bright","dark" \
--modes "benign","malignant","normal"\
--percent_list 0.5,1.0,2.0\
--prompt_structure "an ultrasound photo of {class_name} tumor in breast"\
--existing_images_base_directory <path_of_existing_original_dataset_local>\
--save_dir <path_of_expanded_dataset>\
--copy_flag 0put the name of classes of your dataset in "mode" and the expansion ration in "percent_list". Remmember to not put any space between inputs seperating with cammas. "copy_flag" equal to 1 means if you want to copy the original images in new dataset as well, otherwise your new dataset contains just synthetic images.
In this section, we aim to generate new datasets with DCGAN with different expansion ratios. "save_path" shows the path where the synthetic images are saved and dest_path is the same as we used in training DCGAN. "base_dir" is the path to the folder of each classes of original dataset.
class_name_list=("benign" "malignant" "normal")
percentage=0.5
# Loop over each class name
for class_name in "${class_name_list[@]}"; do
# Define the save path and destination directory
save_path="DCGAN_{percentage}/${class_name}"
# Run the Python script with the appropriate arguments
python dcgan_inference.py \
--save_path "$save_path" \
--dest_path <path_to_model_weights> \
--base_dir <path_to_folder_each_class> \
--percentage $percentage
doneIn this section, we aim to train classifiers for both original and mixed (original + synthetic) ones. we used
'densenet121','resnet34','squeezenet1.1' as classifiers but we wrote a code for other classifiers in the code, you can choose them for your project but pay attention to change their last layer to work best with your problem. If you want to use it for other augmentation rather than SD, put "" for adjective_list and put 0 for adjective_flag. Moreover, you can put train to 0 if you just want to use the test procedure.
python evaluate_new.py --data_path <dir_of_dataset_local>\
--data_test_path <path_of_test_dataset>\
--data_valid_path <path_of_validation_set>\
--output_path <dir_of_output>\
--adjective_list "bright","colorful","dark","high-contrast","low-contrast","no_adjective","posterized","sheared","solarized","stylized" \
--baselines 'densenet121','resnet34','squeezenet1.1'\
--adjective_flag 1\
--batch_size 32\
--num_class 3\
--num_epochs 100\
--train 0\
--output_file_name <name_of_result_text_and_table>\
--size 224@article{medghalchi2025synthetic,
title={Synthetic vs. Classic Data Augmentation: Impacts on Breast Ultrasound Image Classification},
author={Medghalchi, Yasamin and Zakariaei, Niloufar and Rahmim, Arman and Hacihaliloglu, Ilker},
journal={IEEE Transactions on Ultrasonics, Ferroelectrics, and Frequency Control},
year={2025},
publisher={IEEE}
}