Skip to content

gagolucasm/quality_gap_ultrasound_segmentation

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

5 Commits
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Bridging the Quality Gap: Robust Colon Wall Segmentation in Noisy Transabdominal Ultrasound

Official implementation of the paper:

Bridging the quality gap: Robust colon wall segmentation in noisy transabdominal ultrasound
Lucas Gago, Miguel A. Fernández González, Justin Engelmann, Beatriz Remeseiro, Laura Igual
Computers in Biology and Medicine, Volume 197, Part B, 2025
DOI: 10.1016/j.compbiomed.2025.111077

Abstract

Colon wall segmentation in transabdominal ultrasound is challenging due to variations in image quality, speckle noise, and ambiguous boundaries. We present a novel quality-aware segmentation framework that simultaneously predicts image quality and adapts the segmentation process accordingly. Our approach uses a U-Net architecture with a ConvNeXt encoder backbone, enhanced with a parallel quality prediction branch that serves as a regularization mechanism.

Key Features

  • Quality-Aware U-Net (QA U-Net): Dual-branch architecture with ConvNeXt encoder backbone
  • Quality-Weighted Loss: Adaptive loss function based on image quality predictions
  • Comprehensive Data Augmentation: Ultrasound-specific augmentations including mixup and gradient clipping
  • State-of-the-Art Results: Dice scores of 0.7780, 0.7025, and 0.5970 for high, medium, and low-quality images

Repository Structure

quality_gap_ultrasound_segmentation/
├── configs/             # Configuration files for experiments
│   ├── qa_unet.yaml     # QA U-Net configuration
│   └── mask_rcnn.yaml   # Mask R-CNN configuration
├── src/                 # Source code
│   ├── models/          # Model architectures
│   ├── data/            # Dataset and data loaders
│   ├── losses/          # Loss functions
│   ├── training/        # Training utilities
│   ├── evaluation/      # Evaluation and metrics
│   └── utils/           # General utilities
├── scripts/             # Training and evaluation scripts
│   ├── train_qa_unet.py
│   ├── train_mask_rcnn.py
│   └── evaluate.py
├── analysis/            # Analysis and visualization scripts
│   ├── analyze_results.py
├── requirements.txt     # Python dependencies
└── README.md            # This file

Installation

Requirements

  • Python 3.8+
  • PyTorch 2.0+
  • CUDA 11.0+ (for GPU support)

Setup

# Clone the repository
git clone <repository-url>
cd github

# Create a virtual environment (recommended)
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install dependencies
pip install -r requirements.txt

Dataset

This work uses the C-TRUS dataset for colon wall segmentation in transabdominal ultrasound:

  • 827 ultrasound images (580 × 360 pixels)
  • 13 patients with ulcerative colitis
  • Expert annotations with quality ratings (high/medium/low)
  • 5-fold cross-validation splits provided

Dataset Structure:

data/
├── original/           # Original ultrasound images
├── labels/             # Segmentation masks
└── c-trus.csv  # Metadata with quality labels

For access to the C-TRUS dataset, please refer to:

https://github.com/wwu-mmll/c-trus

Usage

Training QA U-Net

# Train with default configuration
python scripts/train_qa_unet.py --config configs/qa_unet.yaml

# Train with custom parameters
python scripts/train_qa_unet.py \
    --config configs/qa_unet.yaml \
    --batch_size 4 \
    --learning_rate 5e-5 \
    --epochs 100

Training Mask R-CNN Baseline

python scripts/train_mask_rcnn.py --config configs/mask_rcnn.yaml

Evaluation

# Evaluate a trained model
python scripts/evaluate.py \
    --experiment_dir experiments/qa_unet_20250101_120000 \
    --csv_path data/c-trus-with-areas.csv \
    --use_tta

# Generate comprehensive analysis
python analysis/analyze_results.py \
    --experiment_dir experiments/qa_unet_20250101_120000 \
    --csv_path data/c-trus-with-areas.csv

Configuration

The training process is controlled via YAML configuration files in the configs/ directory. Key parameters include:

  • Model: Encoder backbone, decoder architecture
  • Training: Learning rate, batch size, epochs, early stopping
  • Data Augmentation: Transforms, mixup probability
  • Loss Functions: Quality weighting, BCE/Dice combination
  • Quality Assessment: Quality prediction branch settings

Example configuration:

# QA U-Net Configuration
model:
  encoder_name: "tu-convnext_base"
  encoder_weights: "imagenet"
  in_channels: 3
  classes: 1
  num_quality_classes: 3

training:
  learning_rate: 5e-5
  batch_size: 4
  epochs: 100
  patience: 15
  gradient_clip: 1.0
  
loss:
  quality_weight: 0.2
  quality_loss_weights:
    high: 1.0
    medium: 0.5
    low: 0.25
    
augmentation:
  mixup_alpha: 0.4
  mixup_prob: 0.2

Citation

If you use this code or methodology in your research, please cite:

@article{gago2025bridging,
  title={Bridging the quality gap: Robust colon wall segmentation in noisy transabdominal ultrasound},
  author={Gago, Lucas and Fern{\'a}ndez Gonz{\'a}lez, Miguel A. and Engelmann, Justin and Remeseiro, Beatriz and Igual, Laura},
  journal={Computers in Biology and Medicine},
  volume={197},
  pages={111077},
  year={2025},
  publisher={Elsevier},
  doi={10.1016/j.compbiomed.2025.111077}
}

License

This project is released under the Creative Commons license, consistent with the open access publication.

Acknowledgments

This work was partially funded by:

  • Spanish Ministry of Science and Innovation (MICINN) - Grant PID2022-136436NB-I00
  • Agency for Management of University and Research Grants of Catalonia (AGAUR) - Grant 2021-SGR-01104
  • Agency for Science, Business Competitiveness, and Innovation of the Principality of Asturias (SEKUENS) - Project GRU-GIC-24-018

References

Please refer to the paper for the complete list of references and related work.

About

Bridging the Quality Gap: Robust Colon Wall Segmentation in Noisy Transabdominal Ultrasound

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages