Model Training & SlideLab Pipeline

The ColabBio training workflow combines SlideLab preprocessing and PyTorch Multiple Instance Learning (ABMIL) within a distributed Nextflow DSL2 pipeline.

graph TD
    WSI["Raw Gigapixel Slide (.svs, .tiff)"] --> SL1["SlideLab: Tissue Masking (Otsu/HSV)"]
    SL1 --> SL2["SlideLab: Color Normalization (Macenko / Vahadane)"]
    SL2 --> SL3["SlideLab: Tile Extraction (512x512)"]
    SL3 --> SL4["Feature Extraction: ResNet50 / UNI / CTransPath"]
    SL4 --> H5["HDF5 Tile Embeddings (.h5)"]
    
    H5 --> TRAIN["ABMIL Attention Training (PyTorch)"]
    CLINICAL["OMOP / FHIR Tabular Covariates"] --> TRAIN
    
    TRAIN --> METRICS["Validation (AUC, Precision, F1, Loss)"]
    TRAIN --> WEIGHTS["Model Weights (.pt / .onnx)"]
    
    METRICS --> MLFLOW["MLflow Tracking Server"]
    WEIGHTS --> MLFLOW_REGISTRY["MLflow Model Registry"]

---

1. SlideLab Preprocessing Pipeline

SlideLab transforms heterogeneous gigapixel pathology slides into compact, normalized feature embeddings.

Steps Executed by SlideLab:

1. Tissue Masking: Identifies valid tissue regions while discarding background glass, marker ink, and air bubbles using adaptive HSV and Otsu thresholding.
2. Color Normalization: Applies Macenko or Vahadane optical density matrix decomposition to eliminate staining variations across hospital laboratories.
3. Tiling: Divides high-magnification tissue regions into non-overlapping 512x512 patches at 20x (0.50 µm/px).
4. Embedding Generation: Passes tiles through pre-trained foundation models (UNI, CONCH, CTransPath) to generate dense feature matrices saved in HDF5 (.h5) format.

---

2. Distributed Training (MODEL_TRAINING)

The training process reads preprocessed .h5 embeddings and clinical tables through MedVFS paths:

process MODEL_TRAINING {
    tag "Train: ${experiment_name} (${training_type})"
    container 'ghcr.io/colabbio/slidelab:latest'

    input:
    val experiment_name
    val training_type
    val target_label
    val epochs
    val learning_rate
    path dataset_manifest

    output:
    path "trained_model_weights.pt", emit: weights
    path "training_metrics.json",   emit: metrics

    script:
    """
    python3 /opt/colabbio/train/train_abmil.py \
        --experiment-name "${experiment_name}" \
        --manifest "${dataset_manifest}" \
        --target-label "${target_label}" \
        --epochs ${epochs} \
        --lr ${learning_rate} \
        --mlflow-tracking-uri "${params.mlflow_tracking_uri}" \
        --output-weights "trained_model_weights.pt" \
        --output-metrics "training_metrics.json"
    """
}

---

3. MLflow Tracking & Validation

During training, the script logs parameters, epoch metrics, and attention heatmaps directly to the centralized MLflow Tracking Server:

• Logged Hyperparameters: learning_rate, batch_size, attention_heads, weight_decay.
• Logged Metrics: train_loss, val_loss, val_auc_roc, val_f1_score, cohen_kappa.
• Artifacts: Checkpoint weights, confusion matrices, and ROC curves.