Skip to content
WSAttention-Prostate Logo

WSAttention-Prostate

Weakly Supervised Attention-Based Deep Learning for Prostate Cancer Characterization from Bi-Parametric Prostate MRI.

WSAttention-Prostate is a two-stage deep learning pipeline that predicts clinically significant prostate cancer (csPCa) risk and PI-RADS score (2 to 5) from T2W, DWI, and ADC bpMRI sequences. The backbone is a patch based 3D Multiple-Instance Learning (MIL) model pre-trained to classify PI-RADS scores and fine-tuned to predict csPCa risk — all without requiring lesion-level annotations.

💡 GUI for real-time inference available at Hugging Face Spaces

Key Features

  • Weakly-supervised attention — Heatmap-guided patch sampling and cosine-similarity attention loss replace the need for voxel-level labels
  • 3D Multiple Instance Learning — Extracts volumetric patches from MRI scans and aggregates them via transformer + attention pooling
  • Two-stage pipeline — Stage 1 trains a 4-class PI-RADS classifier; Stage 2 freezes its backbone and trains a binary csPCa head
  • Preprocessing — Preprocessing to minimize inter-center MRI acquisiton variability.
  • End-to-end pipeline — Registration, segmentation, histogram matching, and heatmap generation, and inferencing in a single configurable pipeline

Pipeline Overview

%%{init: {'themeVariables': { 'fontSize': '20px' }}}%%
flowchart LR
    A[Raw bpMRI</br>T2 + DWI + ADC] --> B[Preprocessing]
    B --> C[Stage 1:</br>PI-RADS Classification]
    C --> D[Stage 2:</br>csPCa Prediction]
    D --> E[Risk Score + Top-5 Salient Patches]