Multi-SWAV: Self-Supervised Learning for Spatio-Temporal Enhanced Ultrasound-Based Prostate Cancer Detection

This repository contains the implementation of a novel multi-modal framework for prostate cancer detection in ultrasound data, published at ISBI 2025

This repository is the official implementation of the ISBI 2025 paper:
Self-supervised Prototype Learning for Spatio-Temporal Enhanced Ultrasound-based Prostate Cancer Detection.

Abstract

Prostate cancer detection in ultrasound data presents significant challenges due to the highly heterogeneous nature of the cancer and its appearance in ultrasound images. In this work, we introduce a novel multi-modal framework that integrates spatial information in ultrasound images with temporal time-series signals across a sequence of ultrasound images for robust prostate cancer detection. Our framework employs a clustering-based self-supervised learning approach, which allows the model to learn shared, complementary representations from combined spatio-temporal data. This approach not only captures the unique properties of spatial and temporal information, but also enhances the model's ability to handle noisy labels due to imperfect alignment between imaging data and histopathology reports, resulting in improved classification performance. Our method achieves a 7% improvement in AUC and a 10% increase in balanced accuracy compared to models relying on spatial or temporal ultrasound data alone.

Key Features

• Multi-modal Architecture: Integrates spatial ultrasound images with temporal time-series signals
• Self-supervised Learning: Clustering-based approach for learning complementary representations
• Robust to Noisy Labels: Handles imperfect alignment between imaging data and histopathology reports
• Multiple Model Support: Includes InceptionTime, TimesNet, Transformer, and ResNet architectures
• Comprehensive Evaluation: Supports various evaluation metrics and cross-validation

Installation

Prerequisites

• Python 3.8+
• PyTorch 1.9+
• CUDA (optional, for GPU acceleration)

Setup

1. Clone the repository:

git clone https://github.com/DeepRCL/Multi-SWAV.git
cd Multi-SWAV

2. Install dependencies:

pip install -r requirements.txt

3. Install the package in development mode:

pip install -e .

Data Preparation

Important Note

The code includes a default file_resolver function in main.py that maps dataset metadata to file paths. You may need to modify this function based on your specific data structure and file naming conventions.

Dataset Structure

The expected dataset structure should be organized as follows:

data/
├── BK_CENTER_CORES/
│   ├── patXX_corX_rf.npy      # RF ultrasound data
│   ├── patXX_corX_needle.npy   # Needle segmentation masks
│   ├── patXX_corX_prostate.npy # Prostate segmentation masks
│   ├── patXX_corX_bmode.npy    # B-mode ultrasound images
│   └── patXX_corX_roi.npy      # Region of interest masks
└── metadata.csv                # Dataset metadata with splits

Metadata Format

The metadata.csv file should contain the following columns:

• id: Unique identifier for each core
• PatientId: Patient identifier
• center: Data collection center
• filetemplate: File naming template
• fold_idx: Cross-validation fold index (0-4)
• set_name: Dataset split (train/val/test)
• TrueLabel: Ground truth label (0: benign, 1: cancer)
• Involvement: Cancer involvement percentage
• GleasonScore: Gleason score

Usage

Basic Training

Train a model with default settings:

python main.py --dataset_name patch_1d --model_name inception --epochs 100

Advanced Configuration

Use configuration files for complex setups:

python main.py --config configs/baseline_16_16.yaml

Available Models

• InceptionTime: --model_name inception
• TimesNet: --model_name timesnet
• Transformer: --model_name transformer
• ResNet: --model_name resnet

Available Datasets

• 1D Patch: --dataset_name patch_1d
• 2D Patch: --dataset_name patch_2d
• 3D Patch: --dataset_name patch_3d
• TeUS Patch: --dataset_name patch_teus
• Needle TeUS: --dataset_name needle_teus

Self-supervised Pre-training

Enable unsupervised pre-training:

python main.py --dataset_name patch_1d --model_name inception \
    --unsupervised_epochs 10 --unsupervised_lr 0.001

Co-teaching Training

Use co-teaching for robust training with noisy labels:

python main.py --dataset_name patch_1d --model_name inception \
    --use_coteaching --forget_rate 0.2

Configuration

Key Parameters

• --epochs: Number of training epochs
• --batch_size: Batch size for training
• --learning_rate: Learning rate
• --fold_idx: Cross-validation fold (0-4)
• --seq_len: Sequence length for temporal models
• --patch_dim: Patch dimensions (height, vertical_stride, horizontal_stride)
• --min_inv: Minimum cancer involvement threshold
• --min_gs: Minimum Gleason score threshold

Loss Functions

• --criterion ce: Cross-entropy loss
• --criterion isomax: IsoMax loss
• --criterion focal: Focal loss

Optimizers

• --optim adam: Adam optimizer
• --optim sgd: SGD optimizer
• --optim novograd: NovoGrad optimizer
• --optim lion: Lion optimizer

Evaluation

Metrics

The framework evaluates models using:

• AUC: Area Under the ROC Curve
• Balanced Accuracy: Balanced classification accuracy
• Sensitivity: True positive rate
• Specificity: True negative rate

Cross-validation

The framework supports 5-fold cross-validation:

for fold in {0..4}; do
    python main.py --dataset_name patch_1d --model_name inception --fold_idx $fold
done

Results

Our Multi-SWAV framework achieves:

• 7% improvement in AUC compared to spatial-only models
• 10% increase in balanced accuracy compared to temporal-only models
• Robust performance across different data splits and noise levels

File Structure

Multi-SWAV/
├── main.py                 # Main training script
├── build.py               # Model and data loader setup
├── configs.py             # Configuration classes
├── src/
│   ├── data/              # Data loading and preprocessing
│   ├── networks/          # Model architectures
│   └── training/          # Training loops and utilities
├── utils/                 # Utility functions and loss functions
├── configs/               # Configuration files
└── notebooks/             # Jupyter notebooks for analysis

Citation

If you use this code in your research, please cite our paper:

@INPROCEEDINGS{10981217,
  author={Elghareb, Tarek and Jamzad, Amoon and Nhat To, Minh Nguyen and Fooladgar, Fahimeh and Wilson, Paul F.R. and Sojoudi, Samira and Reznik, Gabrielle and Leveridge, Michael and Siemens, Robert and Chang, Silvia and Black, Peter and Mousavi, Parvin and Abolmaesumi, Purang},
  booktitle={2025 IEEE 22nd International Symposium on Biomedical Imaging (ISBI)}, 
  title={Self-supervised Prototype Learning for Spatio-Temporal Enhanced Ultrasound-based Prostate Cancer Detection}, 
  year={2025},
  volume={},
  number={},
  pages={1-5},
  keywords={Ultrasonic imaging;Accuracy;Biological system modeling;Prototypes;Self-supervised learning;Data models;Robustness;Noise measurement;Prostate cancer;Principal component analysis;Prostate Cancer;Ultrasound Imaging;Multi-Modal SSL learning;Robustness},
  doi={10.1109/ISBI60581.2025.10981217}}