Add documentation site

Add full project documentation and tooling.

.github/workflows/docs.yml

@@ -0,0 +1,26 @@
+name: Deploy Documentation
+
+on:
+  push:
+    branches: [main]
+    paths:
+      - 'docs/**'
+      - 'mkdocs.yml'
+  workflow_dispatch:
+
+permissions:
+  contents: write
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+
+      - run: pip install mkdocs-material
+
+      - run: mkdocs gh-deploy --force

.gitignore

@@ -6,4 +6,5 @@ temp.ipynb
 __pycache__/
 **/__pycache__/
 *.pyc
-.ruff_cache
+.ruff_cache
+site/

README.md

@@ -1 +1,122 @@
-# WSAttention-Prostate
+<p align="center">
+  <img src="docs/assets/logo.svg" alt="WSAttention-Prostate Logo" width="240">
+</p>
+
+<p align="center">
+  <img src="https://img.shields.io/badge/python-3.11-blue?logo=python&logoColor=white" alt="Python 3.11">
+  <img src="https://img.shields.io/badge/pytorch-2.5-ee4c2c?logo=pytorch&logoColor=white" alt="PyTorch 2.5">
+  <img src="https://img.shields.io/badge/MONAI-1.4-3ddc84" alt="MONAI 1.4">
+  <img src="https://img.shields.io/badge/license-Apache%202.0-green" alt="License">
+  <a href="https://ai-assisted-healthcare.github.io/WSAttention-Prostate/"><img src="https://img.shields.io/badge/docs-mkdocs-blue" alt="Docs"></a>
+</p>
+
+# WSAttention-Prostate
+
+**Weakly-supervised attention-based 3D Multiple Instance Learning for prostate cancer risk prediction on multiparametric MRI.**
+
+WSAttention-Prostate is a two-stage deep learning pipeline that predicts clinically significant prostate cancer (csPCa) risk from T2-weighted, DWI, and ADC MRI sequences. It uses 3D patch-based Multiple Instance Learning with transformer attention to first classify PI-RADS scores, then predict csPCa risk — all without requiring lesion-level annotations.
+
+## 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
+- **Multi-seed confidence intervals** — Runs 20 random seeds and reports 95% CI on AUC, sensitivity, and specificity
+- **End-to-end preprocessing** — Registration, segmentation, histogram matching, and heatmap generation in a single configurable pipeline
+
+## Pipeline Overview
+
+```mermaid
+flowchart LR
+    A[Raw MRI\nT2 + DWI + ADC] --> B[Preprocessing]
+    B --> C[Stage 1:\nPI-RADS Classification]
+    C --> D[Stage 2:\ncsPCa Prediction]
+    D --> E[Risk Score\n+ Top-5 Patches]
+```
+
+## Quick Start
+
+```bash
+git clone https://github.com/ai-assisted-healthcare/WSAttention-Prostate.git
+cd WSAttention-Prostate
+pip install -r requirements.txt
+pytest tests/
+```
+
+## Usage
+
+### Preprocessing
+
+```bash
+python preprocess_main.py --config config/config_preprocess.yaml \
+    --steps register_and_crop get_segmentation_mask histogram_match get_heatmap
+```
+
+### PI-RADS Training
+
+```bash
+python run_pirads.py --mode train --config config/config_pirads_train.yaml
+```
+
+### csPCa Training
+
+```bash
+python run_cspca.py --mode train --config config/config_cspca_train.yaml
+```
+
+### Inference
+
+```bash
+python run_pirads.py --mode test --config config/config_pirads_test.yaml --checkpoint <path>
+python run_cspca.py --mode test --config config/config_cspca_test.yaml --checkpoint_cspca <path>
+python run_inference.py --config config/config_preprocess.yaml
+```
+
+See the [full documentation](https://ai-assisted-healthcare.github.io/WSAttention-Prostate/) for detailed configuration options and data format requirements.
+
+## Project Structure
+
+```
+WSAttention-Prostate/
+├── run_pirads.py              # PI-RADS training/testing entry point
+├── run_cspca.py               # csPCa training/testing entry point
+├── run_inference.py           # Full inference pipeline
+├── preprocess_main.py         # Preprocessing entry point
+├── config/                    # YAML configuration files
+├── src/
+│   ├── model/
+│   │   ├── MIL.py             # MILModel_3D — core MIL architecture
+│   │   └── csPCa_model.py     # csPCa_Model + SimpleNN head
+│   ├── data/
+│   │   ├── data_loader.py     # MONAI data pipeline
+│   │   └── custom_transforms.py
+│   ├── train/
+│   │   ├── train_pirads.py    # PI-RADS training loop
+│   │   └── train_cspca.py     # csPCa training loop
+│   ├── preprocessing/         # Registration, segmentation, heatmaps
+│   └── utils.py               # Shared utilities and step validation
+├── tests/
+├── dataset/                   # Reference images for histogram matching
+└── models/                    # Downloaded checkpoints (not in repo)
+```
+
+## Architecture
+
+Input MRI patches are processed independently through a 3D ResNet18 backbone, then aggregated via a transformer encoder and attention pooling:
+
+```mermaid
+flowchart TD
+    A["Input [B, N, C, D, H, W]"] --> B["Reshape to [B*N, C, D, H, W]"]
+    B --> C[ResNet18-3D Backbone]
+    C --> D["Reshape to [B, N, 512]"]
+    D --> E[Transformer Encoder\n4 layers, 8 heads]
+    E --> F[Attention Pooling\n512 → 2048 → 1]
+    F --> G["Weighted Sum [B, 512]"]
+    G --> H["FC Head [B, num_classes]"]
+```
+
+For csPCa prediction, the backbone is frozen and a 3-layer MLP (`512 → 256 → 128 → 1`) replaces the classification head.
+
+## License
+
+Apache-2.0 — see [LICENSE](LICENSE).

docs/api/data.md

@@ -0,0 +1,140 @@
+# Data Loading Reference
+
+## get_dataloader
+
+```python
+def get_dataloader(args, split: Literal["train", "test"]) -> DataLoader
+```
+
+Creates a PyTorch DataLoader with MONAI transforms and persistent caching.
+
+**Parameters:**
+
+| Parameter | Description |
+|-----------|-------------|
+| `args` | Namespace with `dataset_json`, `data_root`, `tile_size`, `tile_count`, `depth`, `use_heatmap`, `batch_size`, `workers`, `dry_run`, `logdir` |
+| `split` | `"train"` or `"test"` |
+
+**Behavior:**
+
+- Loads data lists from a MONAI decathlon-format JSON
+- In `dry_run` mode, limits to 8 samples
+- Uses `PersistentDataset` with cache stored at `<logdir>/cache/<split>/`
+- Training split is shuffled; test split is not
+- Uses `list_data_collate` to stack patches into `[B, N, C, D, H, W]`
+
+## Transform Pipeline
+
+Two variants depending on `args.use_heatmap`:
+
+### With Heatmaps (default)
+
+| Step | Transform | Description |
+|------|-----------|-------------|
+| 1 | `LoadImaged` | Load T2, mask, DWI, ADC, heatmap (ITKReader, channel-first) |
+| 2 | `ClipMaskIntensityPercentilesd` | Clip T2 intensity to [0, 99.5] percentiles within mask |
+| 3 | `ConcatItemsd` | Stack T2 + DWI + ADC → 3-channel image |
+| 4 | `NormalizeIntensity_customd` | Z-score normalize per channel using mask-only statistics |
+| 5 | `ElementwiseProductd` | Multiply mask * heatmap → `final_heatmap` |
+| 6 | `RandWeightedCropd` | Extract N patches weighted by `final_heatmap` |
+| 7 | `EnsureTyped` | Cast labels to float32 |
+| 8 | `Transposed` | Reorder image dims for 3D convolution |
+| 9 | `DeleteItemsd` | Remove intermediate keys (mask, dwi, adc, heatmap) |
+| 10 | `ToTensord` | Convert to PyTorch tensors |
+
+### Without Heatmaps
+
+| Step | Transform | Description |
+|------|-----------|-------------|
+| 1 | `LoadImaged` | Load T2, mask, DWI, ADC |
+| 2 | `ClipMaskIntensityPercentilesd` | Clip T2 intensity to [0, 99.5] percentiles within mask |
+| 3 | `ConcatItemsd` | Stack T2 + DWI + ADC → 3-channel image |
+| 4 | `NormalizeIntensityd` | Standard channel-wise normalization (MONAI built-in) |
+| 5 | `RandCropByPosNegLabeld` | Extract N patches from positive (mask) regions |
+| 6 | `EnsureTyped` | Cast labels to float32 |
+| 7 | `Transposed` | Reorder image dims |
+| 8 | `DeleteItemsd` | Remove intermediate keys |
+| 9 | `ToTensord` | Convert to tensors |
+
+## list_data_collate
+
+```python
+def list_data_collate(batch: Sequence) -> dict
+```
+
+Custom collation function that stacks per-patient patch lists into batch tensors.
+
+Each sample from the dataset is a list of N patch dictionaries. This function:
+
+1. Stacks `image` across patches: `[N, C, D, H, W]` per sample
+2. Stacks `final_heatmap` if present
+3. Applies PyTorch's `default_collate` to form the batch dimension
+
+Result: `{"image": [B, N, C, D, H, W], "label": [B], ...}`
+
+## Custom Transforms
+
+### ClipMaskIntensityPercentilesd
+
+```python
+ClipMaskIntensityPercentilesd(
+    keys: KeysCollection,
+    mask_key: str,
+    lower: float | None,
+    upper: float | None,
+    sharpness_factor: float | None = None,
+    channel_wise: bool = False,
+    dtype: DtypeLike = np.float32,
+)
+```
+
+Clips image intensity to percentiles computed only from the **masked region**. Supports both hard clipping (default) and soft clipping (via `sharpness_factor`).
+
+### NormalizeIntensity_customd
+
+```python
+NormalizeIntensity_customd(
+    keys: KeysCollection,
+    mask_key: str,
+    subtrahend: NdarrayOrTensor | None = None,
+    divisor: NdarrayOrTensor | None = None,
+    nonzero: bool = False,
+    channel_wise: bool = False,
+    dtype: DtypeLike = np.float32,
+)
+```
+
+Z-score normalization where mean and standard deviation are computed only from **masked voxels**. Supports channel-wise normalization.
+
+### ElementwiseProductd
+
+```python
+ElementwiseProductd(
+    keys: KeysCollection,
+    output_key: str,
+)
+```
+
+Computes the element-wise product of two arrays from the data dictionary and stores the result in `output_key`. Used to combine the prostate mask with the attention heatmap.
+
+## Dataset JSON Format
+
+The pipeline expects a MONAI decathlon-format JSON file:
+
+```json
+{
+    "train": [
+        {
+            "image": "relative/path/to/t2.nrrd",
+            "dwi": "relative/path/to/dwi.nrrd",
+            "adc": "relative/path/to/adc.nrrd",
+            "mask": "relative/path/to/mask.nrrd",
+            "heatmap": "relative/path/to/heatmap.nrrd",
+            "label": 2
+        }
+    ],
+    "test": [...]
+}
+```
+
+Paths are relative to `data_root`. The `heatmap` key is only required when `use_heatmap=True`.

docs/api/models.md

@@ -0,0 +1,104 @@
+# Models Reference
+
+## MILModel_3D
+
+```python
+class MILModel_3D(nn.Module):
+    def __init__(
+        self,
+        num_classes: int,
+        mil_mode: str = "att",
+        pretrained: bool = True,
+        backbone: str | nn.Module | None = None,
+        backbone_num_features: int | None = None,
+        trans_blocks: int = 4,
+        trans_dropout: float = 0.0,
+    )
+```
+
+**Constructor arguments:**
+
+| Argument | Type | Default | Description |
+|----------|------|---------|-------------|
+| `num_classes` | `int` | — | Number of output classes |
+| `mil_mode` | `str` | `"att"` | MIL aggregation mode |
+| `pretrained` | `bool` | `True` | Use pretrained backbone weights |
+| `backbone` | `str \| nn.Module \| None` | `None` | Backbone CNN (None = ResNet18-3D) |
+| `backbone_num_features` | `int \| None` | `None` | Output features of custom backbone |
+| `trans_blocks` | `int` | `4` | Number of transformer encoder layers |
+| `trans_dropout` | `float` | `0.0` | Transformer dropout rate |
+
+**MIL modes:**
+
+| Mode | Description |
+|------|-------------|
+| `mean` | Average logits across all patches — equivalent to pure CNN |
+| `max` | Keep only the max-probability instance for loss |
+| `att` | Attention-based MIL ([Ilse et al., 2018](https://arxiv.org/abs/1802.04712)) |
+| `att_trans` | Transformer + attention MIL ([Shao et al., 2021](https://arxiv.org/abs/2111.01556)) |
+| `att_trans_pyramid` | Pyramid transformer using intermediate ResNet layers |
+
+**Key methods:**
+
+- `forward(x, no_head=False)` — Full forward pass. If `no_head=True`, returns patch-level features `[B, N, 512]` before transformer and attention pooling (used during attention loss computation).
+- `calc_head(x)` — Applies the MIL aggregation and classification head to patch features.
+
+**Example:**
+
+```python
+import torch
+from src.model.MIL import MILModel_3D
+
+model = MILModel_3D(num_classes=4, mil_mode="att_trans")
+# Input: [batch, patches, channels, depth, height, width]
+x = torch.randn(2, 24, 3, 3, 64, 64)
+logits = model(x)  # [2, 4]
+```
+
+## csPCa_Model
+
+```python
+class csPCa_Model(nn.Module):
+    def __init__(self, backbone: nn.Module)
+```
+
+Wraps a pre-trained `MILModel_3D` backbone for binary csPCa prediction. The backbone's feature extractor, transformer, and attention mechanism are reused. The original classification head (`myfc`) is replaced by a `SimpleNN`.
+
+**Attributes:**
+
+| Attribute | Type | Description |
+|-----------|------|-------------|
+| `backbone` | `MILModel_3D` | Frozen PI-RADS backbone |
+| `fc_cspca` | `SimpleNN` | Binary classification head |
+| `fc_dim` | `int` | Feature dimension (512 for ResNet18) |
+
+**Example:**
+
+```python
+import torch
+from src.model.MIL import MILModel_3D
+from src.model.csPCa_model import csPCa_Model
+
+backbone = MILModel_3D(num_classes=4, mil_mode="att_trans")
+model = csPCa_Model(backbone=backbone)
+
+x = torch.randn(2, 24, 3, 3, 64, 64)
+prob = model(x)  # [2, 1] — sigmoid probabilities
+```
+
+## SimpleNN
+
+```python
+class SimpleNN(nn.Module):
+    def __init__(self, input_dim: int)
+```
+
+A lightweight MLP for binary classification:
+
+```
+Linear(input_dim, 256) → ReLU
+Linear(256, 128) → ReLU → Dropout(0.3)
+Linear(128, 1) → Sigmoid
+```
+
+Input: `[B, input_dim]` — Output: `[B, 1]` (probability).

docs/api/preprocessing.md

@@ -0,0 +1,103 @@
+# Preprocessing Reference
+
+## Overview
+
+Preprocessing is orchestrated by `preprocess_main.py`, which runs steps in sequence. Each step receives and returns the `args` namespace, updating directory paths as it goes.
+
+## Step Dependencies
+
+| Step | Requires |
+|------|----------|
+| `register_and_crop` | — |
+| `get_segmentation_mask` | `register_and_crop` |
+| `histogram_match` | `register_and_crop`, `get_segmentation_mask` |
+| `get_heatmap` | `register_and_crop`, `get_segmentation_mask`, `histogram_match` |
+
+Dependencies are validated at runtime — the pipeline will exit with an error if steps are out of order.
+
+## register_files
+
+```python
+def register_files(args) -> args
+```
+
+Registers and crops T2, DWI, and ADC images to a standardized spacing and size.
+
+**Process:**
+
+1. Reads images from `args.t2_dir`, `args.dwi_dir`, `args.adc_dir`
+2. Resamples to spacing `(0.4, 0.4, 3.0)` mm using `picai_prep.Sample`
+3. Center-crops with `args.margin` (default 0.2) in x/y dimensions
+4. Saves to `<output_dir>/t2_registered/`, `DWI_registered/`, `ADC_registered/`
+
+**Updates `args`:** `t2_dir`, `dwi_dir`, `adc_dir` → registered directories.
+
+## get_segmask
+
+```python
+def get_segmask(args) -> args
+```
+
+Generates prostate segmentation masks from T2W images using a pre-trained model.
+
+**Process:**
+
+1. Loads model config from `<project_dir>/config/inference.json`
+2. Loads checkpoint from `<project_dir>/models/prostate_segmentation_model.pt`
+3. Applies MONAI transforms: orientation (RAS), spacing (0.5 mm isotropic), intensity normalization
+4. Runs inference and inverts transforms to original space
+5. Post-processes: retains only top 10 slices by non-zero voxel count
+6. Saves NRRD masks to `<output_dir>/prostate_mask/`
+
+**Updates `args`:** adds `seg_dir`.
+
+## histmatch
+
+```python
+def histmatch(args) -> args
+```
+
+Matches the intensity histogram of each modality to a reference image.
+
+**Process:**
+
+1. Reads reference images from `<project_dir>/dataset/` (`t2_reference.nrrd`, `dwi_reference.nrrd`, `adc_reference.nrrd`, `prostate_segmentation_reference.nrrd`)
+2. For each patient, matches histograms within the prostate mask using `skimage.exposure.match_histograms`
+3. Saves to `<output_dir>/t2_histmatched/`, `DWI_histmatched/`, `ADC_histmatched/`
+
+**Updates `args`:** `t2_dir`, `dwi_dir`, `adc_dir` → histogram-matched directories.
+
+### get_histmatched
+
+```python
+def get_histmatched(
+    data: np.ndarray,
+    ref_data: np.ndarray,
+    mask: np.ndarray,
+    ref_mask: np.ndarray,
+) -> np.ndarray
+```
+
+Low-level function that performs histogram matching on masked regions only. Unmasked pixels remain unchanged.
+
+## get_heatmap
+
+```python
+def get_heatmap(args) -> args
+```
+
+Generates combined DWI/ADC attention heatmaps.
+
+**Process:**
+
+1. For each file, reads DWI, ADC, and prostate mask
+2. Computes DWI heatmap: `(dwi - min) / (max - min)` within mask
+3. Computes ADC heatmap: `(max - adc) / (max - min)` within mask (inverted — low ADC = high attention)
+4. Combines via element-wise multiplication
+5. Re-normalizes to [0, 1]
+6. Saves to `<output_dir>/heatmaps/`
+
+**Updates `args`:** adds `heatmapdir`.
+
+!!! info "Edge cases"
+    If all values within the mask are identical for a modality (DWI or ADC), that modality's heatmap is skipped. If both are constant, the heatmap defaults to all ones.

docs/architecture.md

@@ -0,0 +1,97 @@
+# Architecture
+
+## Tensor Shape Convention
+
+Throughout the pipeline, tensors follow the shape `[B, N, C, D, H, W]`:
+
+| Dim | Meaning | Typical Value |
+|-----|---------|---------------|
+| B | Batch size | 4–8 |
+| N | Number of patches (instances) | 24 |
+| C | Channels (T2 + DWI + ADC) | 3 |
+| D | Depth (slices per patch) | 3 |
+| H | Patch height | 64 |
+| W | Patch width | 64 |
+
+## MILModel_3D
+
+The core model processes each patch independently through a CNN backbone, then aggregates patch-level features via a transformer encoder and attention pooling.
+
+```mermaid
+flowchart TD
+    A["Input [B, N, C, D, H, W]"] --> B["Reshape to [B*N, C, D, H, W]"]
+    B --> C[ResNet18-3D Backbone]
+    C --> D["Reshape to [B, N, 512]"]
+    D --> E[Transformer Encoder\n4 layers, 8 heads]
+    E --> F[Attention Pooling\n512 → 2048 → 1]
+    F --> G["Weighted Sum [B, 512]"]
+    G --> H["FC Head [B, num_classes]"]
+```
+
+### Forward Pass
+
+1. **Backbone**: Input is reshaped from `[B, N, C, D, H, W]` to `[B*N, C, D, H, W]` and passed through a 3D ResNet18 (with 3 input channels). The final FC layer is removed, yielding 512-dimensional features per patch.
+
+2. **Transformer**: Features are reshaped to `[B, N, 512]`, permuted to `[N, B, 512]` for the transformer encoder (4 layers, 8 attention heads), then permuted back.
+
+3. **Attention**: A two-layer attention network (`512 → 2048 → 1` with Tanh) computes a scalar weight per patch, normalized via softmax.
+
+4. **Classification**: The attention-weighted sum of patch features produces a single `[B, 512]` vector per scan, which is projected to class logits by a linear layer.
+
+### MIL Modes
+
+| Mode | Aggregation Strategy |
+|------|---------------------|
+| `mean` | Average logits across patches |
+| `max` | Max logits across patches |
+| `att` | Attention-weighted feature pooling |
+| `att_trans` | Transformer encoder + attention pooling (primary mode) |
+| `att_trans_pyramid` | Pyramid transformer on intermediate ResNet layers + attention |
+
+The default and primary mode is `att_trans`.
+
+## csPCa_Model
+
+Wraps a frozen `MILModel_3D` backbone and replaces the classification head:
+
+```mermaid
+flowchart TD
+    A["Input [B, N, C, D, H, W]"] --> B["Frozen Backbone\n(ResNet18 + Transformer)"]
+    B --> C["Pooled Features [B, 512]"]
+    C --> D["SimpleNN Head\n512 → 256 → 128 → 1"]
+    D --> E["Sigmoid → csPCa Probability"]
+```
+
+### SimpleNN
+
+```
+Linear(512, 256) → ReLU
+Linear(256, 128) → ReLU → Dropout(0.3)
+Linear(128, 1) → Sigmoid
+```
+
+During csPCa training, the backbone's `net` (ResNet18), `transformer`, and `myfc` parameters are frozen. The `attention` module and `SimpleNN` head remain trainable.
+
+## Attention Loss
+
+During PI-RADS training with heatmaps enabled, the model uses a dual-loss objective:
+
+```
+total_loss = class_loss + lambda_att * attention_loss
+```
+
+- **Classification loss**: Standard CrossEntropy on PI-RADS labels
+- **Attention loss**: `1 - cosine_similarity(predicted_attention, heatmap_attention)`
+    - Heatmap-derived attention labels are computed by summing spatial heatmap values per patch, squaring for sharpness, and normalizing
+    - PI-RADS 2 samples get uniform attention (no expected lesion)
+    - `lambda_att` warms up linearly from 0 to 2.0 over the first 25 epochs
+    - The attention predictions are computed with detached transformer outputs to avoid gradient interference with classification
+
+## Patch Extraction
+
+Patches are extracted using MONAI's `RandWeightedCropd` (when heatmaps are available) or `RandCropByPosNegLabeld` (without heatmaps):
+
+- **With heatmaps**: The combined DWI/ADC heatmap multiplied by the prostate mask serves as the sampling weight map — regions with high DWI and low ADC are sampled more frequently
+- **Without heatmaps**: Crops are sampled from positive (prostate) regions based on the binary mask
+
+Each scan yields `N` patches (default 24) of size `tile_size x tile_size x depth` (default 64x64x3).

docs/configuration.md

@@ -0,0 +1,147 @@
+# Configuration
+
+## Config System
+
+Configuration follows a three-level hierarchy:
+
+1. **CLI defaults** — Argparse defaults in `run_pirads.py`, `run_cspca.py`, etc.
+2. **YAML overrides** — Values from `--config <file>.yaml` override CLI defaults
+3. **SLURM job name** — If `SLURM_JOB_NAME` is set, it overrides `run_name`
+
+```bash
+# CLI defaults are overridden by YAML config
+python run_pirads.py --mode train --config config/config_pirads_train.yaml
+```
+
+!!! note
+    YAML values **always override** CLI defaults for any key present in the YAML file (`args.__dict__.update(config)`). To override a YAML value, edit the YAML file or omit the key from YAML so the CLI default is used.
+
+## PI-RADS Training Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `mode` | — | `train` or `test` (required) |
+| `config` | — | Path to YAML config file |
+| `data_root` | — | Root folder of images |
+| `dataset_json` | — | Path to dataset JSON file |
+| `num_classes` | `4` | Number of output classes (PI-RADS 2–5) |
+| `mil_mode` | `att_trans` | MIL algorithm (`mean`, `max`, `att`, `att_trans`, `att_trans_pyramid`) |
+| `tile_count` | `24` | Number of patches per scan |
+| `tile_size` | `64` | Patch spatial size in pixels |
+| `depth` | `3` | Number of slices per patch |
+| `use_heatmap` | `True` | Enable heatmap-guided patch sampling |
+| `workers` | `2` | DataLoader workers |
+| `checkpoint` | `None` | Path to resume from checkpoint |
+| `epochs` | `50` | Max training epochs |
+| `early_stop` | `40` | Epochs without improvement before stopping |
+| `batch_size` | `4` | Scans per batch |
+| `optim_lr` | `3e-5` | Base learning rate |
+| `weight_decay` | `0` | Optimizer weight decay |
+| `amp` | `False` | Enable automatic mixed precision |
+| `val_every` | `1` | Validation frequency (epochs) |
+| `wandb` | `False` | Enable Weights & Biases logging |
+| `project_name` | `Classification_prostate` | W&B project name |
+| `run_name` | `train_pirads` | Run name for logging |
+| `dry_run` | `False` | Quick test mode |
+
+## csPCa Training Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `mode` | — | `train` or `test` (required) |
+| `config` | — | Path to YAML config file |
+| `data_root` | — | Root folder of images |
+| `dataset_json` | — | Path to dataset JSON file |
+| `num_classes` | `4` | PI-RADS classes (for backbone initialization) |
+| `mil_mode` | `att_trans` | MIL algorithm for backbone |
+| `tile_count` | `24` | Number of patches per scan |
+| `tile_size` | `64` | Patch spatial size |
+| `depth` | `3` | Slices per patch |
+| `use_heatmap` | `True` | Enable heatmap-guided patch sampling |
+| `workers` | `2` | DataLoader workers |
+| `checkpoint_pirads` | — | Path to pre-trained PI-RADS model (required for train) |
+| `checkpoint_cspca` | — | Path to csPCa checkpoint (required for test) |
+| `epochs` | `30` | Max training epochs |
+| `batch_size` | `32` | Scans per batch |
+| `optim_lr` | `2e-4` | Learning rate |
+| `num_seeds` | `20` | Number of random seeds for CI |
+| `val_every` | `1` | Validation frequency |
+| `dry_run` | `False` | Quick test mode |
+
+## Preprocessing Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `config` | — | Path to YAML config file |
+| `steps` | — | Steps to execute (required, one or more) |
+| `t2_dir` | — | Directory of T2W images |
+| `dwi_dir` | — | Directory of DWI images |
+| `adc_dir` | — | Directory of ADC images |
+| `seg_dir` | — | Directory of segmentation masks |
+| `output_dir` | — | Output directory |
+| `margin` | `0.2` | Center-crop margin fraction |
+| `project_dir` | — | Project root (for reference images and models) |
+
+## Example YAML
+
+=== "PI-RADS Training"
+
+    ```yaml
+    data_root: /path/to/registered/t2_hist_matched
+    dataset_json: /path/to/PI-RADS_data.json
+    num_classes: 4
+    mil_mode: att_trans
+    tile_count: 24
+    tile_size: 64
+    depth: 3
+    use_heatmap: true
+    workers: 4
+    epochs: 100
+    batch_size: 8
+    optim_lr: 2e-4
+    weight_decay: 1e-5
+    amp: true
+    wandb: true
+    ```
+
+=== "csPCa Training"
+
+    ```yaml
+    data_root: /path/to/registered/t2_hist_matched
+    dataset_json: /path/to/csPCa_data.json
+    num_classes: 4
+    mil_mode: att_trans
+    tile_count: 24
+    tile_size: 64
+    depth: 3
+    use_heatmap: true
+    workers: 6
+    checkpoint_pirads: /path/to/models/pirads.pt
+    epochs: 80
+    batch_size: 8
+    optim_lr: 2e-4
+    ```
+
+=== "Preprocessing"
+
+    ```yaml
+    t2_dir: /path/to/raw/t2
+    dwi_dir: /path/to/raw/dwi
+    adc_dir: /path/to/raw/adc
+    output_dir: /path/to/processed
+    project_dir: /path/to/WSAttention-Prostate
+    ```
+
+## Dry-Run Mode
+
+The `--dry_run` flag configures a minimal run for quick testing:
+
+- Epochs: 2
+- Batch size: 2
+- Workers: 0
+- Seeds: 2
+- W&B: disabled
+
+```bash
+python run_pirads.py --mode train --config config/config_pirads_train.yaml --dry_run
+```

docs/contributing.md

@@ -0,0 +1,62 @@
+# Contributing
+
+## Running Tests
+
+```bash
+# Full test suite
+pytest tests/
+
+# Single test
+pytest tests/test_run.py::test_run_pirads_training
+```
+
+Tests run in `--dry_run` mode (2 epochs, batch_size=2, no W&B logging).
+
+## Linting
+
+This project uses [Ruff](https://docs.astral.sh/ruff/) for linting and formatting:
+
+```bash
+# Check for lint errors
+ruff check .
+
+# Auto-format code
+ruff format .
+```
+
+**Ruff configuration** (from `pyproject.toml`):
+
+| Setting | Value |
+|---------|-------|
+| Line length | 100 |
+| Quote style | Double quotes |
+| Rules | E (errors), W (warnings) |
+| Ignored | E501 (line too long) |
+
+## SLURM Job Scripts
+
+Job scripts are in `job_scripts/` and are configured for GPU partitions:
+
+```bash
+sbatch job_scripts/train_pirads.sh
+sbatch job_scripts/train_cspca.sh
+```
+
+Key SLURM settings used:
+
+| Setting | Value |
+|---------|-------|
+| Partition | `gpu` |
+| Memory | 128 GB |
+| GPUs | 1 |
+| Time limit | 48 hours |
+
+!!! tip
+    The SLURM job name (`--job-name`) automatically becomes the `run_name`, which determines the log directory at `logs/<run_name>/`.
+
+## Project Conventions
+
+- **Configs** are stored in `config/` as YAML files
+- **Logs** are written to `logs/<run_name>/` including TensorBoard events and training logs
+- **Models** are saved to `logs/<run_name>/` during training; best models are saved to `models/` for deployment
+- **Cache** is stored at `logs/<run_name>/cache/` and cleaned up automatically after training

docs/getting-started.md

@@ -0,0 +1,104 @@
+# Getting Started
+
+## Prerequisites
+
+- Python 3.11+
+- NVIDIA GPU recommended (CUDA-compatible)
+- ~128 GB RAM for training (configurable via batch size)
+
+## Installation
+
+```bash
+git clone https://github.com/ai-assisted-healthcare/WSAttention-Prostate.git
+cd WSAttention-Prostate
+pip install -r requirements.txt
+```
+
+### External Git Dependencies
+
+Two packages are installed directly from GitHub repositories:
+
+| Package | Source | Purpose |
+|---------|--------|---------|
+| `AIAH_utility` | `ai-assisted-healthcare/AIAH_utility` | Healthcare imaging utilities |
+| `grad-cam` | `jacobgil/pytorch-grad-cam` | Gradient-weighted class activation maps |
+
+These are included in `requirements.txt` and install automatically.
+
+## Verify Installation
+
+Run the test suite in dry-run mode:
+
+```bash
+pytest tests/
+```
+
+Tests use `--dry_run` mode internally (2 epochs, batch_size=2, no W&B).
+
+## Data Format
+
+Input MRI scans should be in **NRRD** or **NIfTI** format with three modalities per patient:
+
+- T2-weighted (T2W)
+- Diffusion-weighted imaging (DWI)
+- Apparent diffusion coefficient (ADC)
+
+### Dataset JSON Structure
+
+The data pipeline uses MONAI's decathlon-format JSON:
+
+```json
+{
+    "train": [
+        {
+            "image": "path/to/t2.nrrd",
+            "dwi": "path/to/dwi.nrrd",
+            "adc": "path/to/adc.nrrd",
+            "mask": "path/to/prostate_mask.nrrd",
+            "heatmap": "path/to/heatmap.nrrd",
+            "label": 0
+        }
+    ],
+    "test": [
+        ...
+    ]
+}
+```
+
+The `image` key points to the T2W image, which serves as the reference modality. Labels for PI-RADS are 0-indexed: label `0` = PI-RADS 2, label `3` = PI-RADS 5. For csPCa, labels are binary (0 or 1).
+
+## Project Structure
+
+```
+WSAttention-Prostate/
+├── run_pirads.py              # PI-RADS training/testing entry point
+├── run_cspca.py               # csPCa training/testing entry point
+├── run_inference.py           # Full inference pipeline
+├── preprocess_main.py         # Preprocessing entry point
+├── config/                    # YAML configuration files
+│   ├── config_pirads_train.yaml
+│   ├── config_pirads_test.yaml
+│   ├── config_cspca_train.yaml
+│   ├── config_cspca_test.yaml
+│   └── config_preprocess.yaml
+├── src/
+│   ├── model/
+│   │   ├── MIL.py             # MILModel_3D — core MIL architecture
+│   │   └── csPCa_model.py     # csPCa_Model + SimpleNN head
+│   ├── data/
+│   │   ├── data_loader.py     # MONAI data pipeline
+│   │   └── custom_transforms.py
+│   ├── train/
+│   │   ├── train_pirads.py    # PI-RADS training loop
+│   │   └── train_cspca.py     # csPCa training loop
+│   ├── preprocessing/
+│   │   ├── register_and_crop.py
+│   │   ├── prostate_mask.py
+│   │   ├── histogram_match.py
+│   │   └── generate_heatmap.py
+│   └── utils.py
+├── job_scripts/               # SLURM job templates
+├── tests/
+├── dataset/                   # Reference images for histogram matching
+└── models/                    # Pre-trained model checkpoints
+```

docs/index.md

@@ -0,0 +1,35 @@
+<div style="text-align: center; margin-bottom: 2em;">
+  <img src="assets/logo.svg" alt="WSAttention-Prostate Logo" width="240">
+</div>
+
+# WSAttention-Prostate
+
+**Weakly-supervised attention-based 3D Multiple Instance Learning for prostate cancer risk prediction on multiparametric MRI.**
+
+WSAttention-Prostate is a two-stage deep learning pipeline that predicts clinically significant prostate cancer (csPCa) risk from T2-weighted, DWI, and ADC MRI sequences. It uses 3D patch-based Multiple Instance Learning with transformer attention to first classify PI-RADS scores, then predict csPCa risk — all without requiring lesion-level annotations.
+
+## 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
+- **Multi-seed confidence intervals** — Runs 20 random seeds and reports 95% CI on AUC, sensitivity, and specificity
+- **End-to-end preprocessing** — Registration, segmentation, histogram matching, and heatmap generation in a single configurable pipeline
+
+## Pipeline Overview
+
+```mermaid
+flowchart LR
+    A[Raw MRI\nT2 + DWI + ADC] --> B[Preprocessing]
+    B --> C[Stage 1:\nPI-RADS Classification]
+    C --> D[Stage 2:\ncsPCa Prediction]
+    D --> E[Risk Score\n+ Top-5 Patches]
+```
+
+## Quick Links
+
+- [Getting Started](getting-started.md) — Installation and first run
+- [Pipeline](pipeline.md) — Full walkthrough of preprocessing, training, and evaluation
+- [Architecture](architecture.md) — Model design and tensor shapes
+- [Configuration](configuration.md) — YAML config reference
+- [Inference](inference.md) — Running predictions on new data

docs/inference.md

@@ -0,0 +1,81 @@
+# Inference
+
+## Full Pipeline
+
+`run_inference.py` runs the complete pipeline: preprocessing followed by PI-RADS classification and csPCa risk prediction.
+
+```bash
+python run_inference.py --config config/config_preprocess.yaml
+```
+
+This script:
+
+1. Runs all four preprocessing steps (register, segment, histogram match, heatmap)
+2. Loads the PI-RADS model from `models/pirads.pt`
+3. Loads the csPCa model from `models/cspca_model.pth`
+4. For each scan: predicts PI-RADS score, csPCa risk probability, and identifies the top-5 most-attended patches
+
+### Required Model Files
+
+Place these in the `models/` directory:
+
+| File | Description |
+|------|-------------|
+| `pirads.pt` | Trained PI-RADS MIL model checkpoint |
+| `cspca_model.pth` | Trained csPCa model checkpoint |
+| `prostate_segmentation_model.pt` | Pre-trained prostate segmentation model |
+
+### Output Format
+
+Results are saved to `<output_dir>/results.json`:
+
+```json
+{
+    "patient_001.nrrd": {
+        "Predicted PIRAD Score": 4.0,
+        "csPCa risk": 0.8234,
+        "Top left coordinate of top 5 patches(x,y,z)": [
+            [32, 45, 7],
+            [28, 50, 7],
+            [35, 42, 8],
+            [30, 48, 6],
+            [33, 44, 8]
+        ]
+    }
+}
+```
+
+### Label Mapping
+
+PI-RADS predictions are 0-indexed internally and shifted by +2 for display:
+
+| Internal Label | PI-RADS Score |
+|---------------|---------------|
+| 0 | PI-RADS 2 |
+| 1 | PI-RADS 3 |
+| 2 | PI-RADS 4 |
+| 3 | PI-RADS 5 |
+
+csPCa risk is a continuous probability in [0, 1].
+
+## Testing Individual Models
+
+### PI-RADS Testing
+
+```bash
+python run_pirads.py --mode test \
+    --config config/config_pirads_test.yaml \
+    --checkpoint models/pirads.pt
+```
+
+Reports Quadratic Weighted Kappa (QWK) across multiple seeds.
+
+### csPCa Testing
+
+```bash
+python run_cspca.py --mode test \
+    --config config/config_cspca_test.yaml \
+    --checkpoint_cspca models/cspca_model.pth
+```
+
+Reports AUC, sensitivity, and specificity with 95% confidence intervals across 20 seeds (default).

docs/pipeline.md

@@ -0,0 +1,99 @@
+# Pipeline
+
+The full pipeline has three phases: preprocessing, PI-RADS training (Stage 1), and csPCa training (Stage 2).
+
+```mermaid
+flowchart TD
+    subgraph Preprocessing
+        R[register_and_crop] --> S[get_segmentation_mask]
+        S --> H[histogram_match]
+        H --> G[get_heatmap]
+    end
+
+    subgraph Stage 1
+        P[PI-RADS Training\nCrossEntropy + Attention Loss]
+    end
+
+    subgraph Stage 2
+        C[csPCa Training\nFrozen Backbone + BCE Loss]
+    end
+
+    G --> P
+    P -->|frozen backbone| C
+```
+
+## Preprocessing
+
+Run all four steps in sequence:
+
+```bash
+python preprocess_main.py \
+    --config config/config_preprocess.yaml \
+    --steps register_and_crop get_segmentation_mask histogram_match get_heatmap
+```
+
+### Step 1: Register and Crop
+
+Resamples T2, DWI, and ADC to a common spacing of `(0.4, 0.4, 3.0)` mm using `picai_prep`, then center-crops with a configurable margin (default 20%).
+
+### Step 2: Prostate Segmentation
+
+Runs a pre-trained segmentation model on T2W images to generate binary prostate masks. Post-processing retains only the top 10 slices by non-zero voxel count.
+
+### Step 3: Histogram Matching
+
+Matches the intensity histogram of each modality to a reference image within masked (prostate) regions using `skimage.exposure.match_histograms`.
+
+### Step 4: Heatmap Generation
+
+Creates attention heatmaps from DWI and ADC:
+
+- **DWI heatmap**: `(dwi - min) / (max - min)` — higher DWI signal = higher attention
+- **ADC heatmap**: `(max - adc) / (max - min)` — lower ADC = higher attention (inverted)
+- **Combined**: element-wise product, re-normalized to [0, 1]
+
+!!! note "Step Dependencies"
+    Steps must run in the order shown above. The pipeline validates dependencies automatically — for example, `get_heatmap` requires `get_segmentation_mask` and `histogram_match` to have run first.
+
+## Stage 1: PI-RADS Classification
+
+Trains a 4-class PI-RADS classifier (grades 2–5, mapped to labels 0–3).
+
+```bash
+python run_pirads.py --mode train --config config/config_pirads_train.yaml
+```
+
+**Training details:**
+
+| Component | Value |
+|-----------|-------|
+| Loss | CrossEntropy + cosine-similarity attention loss |
+| Attention loss weight | Linear warmup over 25 epochs to `lambda=2.0` |
+| Optimizer | AdamW (base LR `3e-5`, transformer LR `6e-5`) |
+| Scheduler | CosineAnnealingLR |
+| Metric | Quadratic Weighted Kappa (QWK) |
+| Early stopping | After 40 epochs without validation loss improvement |
+| AMP | Disabled by default (enabled in example YAML config) |
+
+**Attention loss**: For each batch, the model's learned attention weights are compared against heatmap-derived attention labels via cosine similarity. PI-RADS 2 samples receive uniform attention (no lesion expected). The loss is weighted by `lambda_att`, which warms up linearly over the first 25 epochs.
+
+## Stage 2: csPCa Risk Prediction
+
+Builds on a frozen PI-RADS backbone to predict binary csPCa risk.
+
+```bash
+python run_cspca.py --mode train --config config/config_cspca_train.yaml
+```
+
+**Training details:**
+
+| Component | Value |
+|-----------|-------|
+| Loss | Binary Cross-Entropy (BCE) |
+| Backbone | Frozen PI-RADS model (ResNet18 + Transformer); attention module is trainable |
+| Head | SimpleNN: `512 → 256 → 128 → 1` with ReLU + Dropout(0.3) + Sigmoid |
+| Optimizer | AdamW (LR `2e-4`) |
+| Seeds | 20 random seeds (default) for 95% CI |
+| Metrics | AUC, Sensitivity, Specificity |
+
+The backbone's feature extractor (`net`), transformer, and `myfc` are frozen. The attention module and `SimpleNN` classification head are trained. After training across all seeds, the framework reports mean and 95% confidence intervals for AUC, sensitivity, and specificity.

mkdocs.yml

@@ -0,0 +1,67 @@
+site_name: WSAttention-Prostate
+site_description: Weakly-supervised attention-based 3D MIL for prostate cancer risk prediction on multiparametric MRI
+repo_url: https://github.com/ai-assisted-healthcare/WSAttention-Prostate
+repo_name: WSAttention-Prostate
+
+theme:
+  name: material
+  logo: assets/logo.svg
+  favicon: assets/logo.svg
+  palette:
+    - scheme: default
+      primary: indigo
+      accent: teal
+      toggle:
+        icon: material/brightness-7
+        name: Switch to dark mode
+    - scheme: slate
+      primary: indigo
+      accent: teal
+      toggle:
+        icon: material/brightness-4
+        name: Switch to light mode
+  font:
+    text: Inter
+    code: JetBrains Mono
+  features:
+    - navigation.instant
+    - navigation.sections
+    - navigation.top
+    - search.suggest
+    - search.highlight
+    - content.code.copy
+    - content.tabs.link
+
+plugins:
+  - search
+  # TODO: add mkdocstrings[python] plugin for autodoc support
+
+markdown_extensions:
+  - admonition
+  - pymdownx.details
+  - pymdownx.superfences:
+      custom_fences:
+        - name: mermaid
+          class: mermaid
+          format: !!python/name:pymdownx.superfences.fence_code_format
+  - pymdownx.highlight:
+      anchor_linenums: true
+  - pymdownx.inlinehilite
+  - pymdownx.tabbed:
+      alternate_style: true
+  - tables
+  - attr_list
+  - md_in_html
+
+nav:
+  - Home: index.md
+  - Getting Started: getting-started.md
+  - Pipeline: pipeline.md
+  - Architecture: architecture.md
+  - Configuration: configuration.md
+  - Inference: inference.md
+  - API Reference:
+      - Models: api/models.md
+      - Preprocessing: api/preprocessing.md
+      - Data Loading: api/data.md
+  - Contributing: contributing.md