upload training_code

.gitignore

@@ -0,0 +1,216 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+#poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+#pdm.lock
+#pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+#pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# Redis
+*.rdb
+*.aof
+*.pid
+
+# RabbitMQ
+mnesia/
+rabbitmq/
+rabbitmq-data/
+
+# ActiveMQ
+activemq-data/
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Abstra
+# Abstra is an AI-powered process automation framework.
+# Ignore directories containing user credentials, local state, and settings.
+# Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
+#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#  and can be added to the global gitignore or merged into this file. However, if you prefer,
+#  you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/
+
+# Streamlit
+.streamlit/secrets.toml

training_code/model/resnet.py

@@ -0,0 +1,54 @@
+import torch
+import torch.nn as nn
+from monai.networks.nets import ResNet
+from monai.networks.nets.resnet import ResNetBlock, ResNetBottleneck
+
+
+class Resnet(nn.Module):
+    
+    def __init__(self, model_name: str, num_classes: int, norm: str = "batch"):
+        super().__init__()
+        
+        self.model_name = model_name
+        valid_model_names = ["resnet10", "resnet18", "resnet34", "resnet50", "resnet101", "resnet152", "resnet200"]
+        assert self.model_name in valid_model_names, f"Model name must be one of {valid_model_names}"
+
+        # Get default resnet parameters
+        # >> model_name: (block, layers, shortcut_type, bias_downsample)
+        resnet_params = {
+            "resnet10": ("basic", [1, 1, 1, 1], "B", False),
+            "resnet18": ("basic", [2, 2, 2, 2], "A", True),
+            "resnet34": ("basic", [3, 4, 6, 3], "A", True),
+            "resnet50": ("bottleneck", [3, 4, 6, 3], "B", False),
+            "resnet101": ("bottleneck", [3, 4, 23, 3], "B", False),
+            "resnet152": ("bottleneck", [3, 8, 36, 3], "B", False),
+            "resnet200": ("bottleneck", [3, 24, 36, 3], "B", False),
+        }
+        block, layers, shortcut_type, bias_downsample = resnet_params[self.model_name]
+        block = ResNetBlock if block == "basic" else ResNetBottleneck
+        
+        # Get default norm parameters
+        norm_mapper = {
+            "group": ("group", {"num_groups": 8}),
+            "instance": ("instance", {"affine": True}),
+            "batch": "batch"
+        }
+        _norm = norm_mapper[norm]
+        
+        # Create resnet model
+        self.net = ResNet(                         
+            block=block,                  
+            layers=layers,                
+            block_inplanes=[64, 128, 256, 512],  
+            spatial_dims=3,
+            n_input_channels=1,
+            num_classes=num_classes,
+            bias_downsample=bias_downsample,
+            shortcut_type=shortcut_type,
+            norm=_norm
+        )
+        
+        return
+        
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)

training_code/train.py

@@ -0,0 +1,261 @@
+import argparse
+import json
+import pandas as pd
+import numpy as np
+from tqdm import tqdm
+from pathlib import Path
+import torch
+import torch.nn as nn
+import matplotlib
+matplotlib.use("Agg")
+import random
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning, module="torchio.data.image")
+import os
+import gc
+
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from torch.amp import GradScaler, autocast
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+torch.multiprocessing.set_sharing_strategy("file_system")
+
+from model.resnet import Resnet
+from utils.dataloader import prepare_loaders, prepare_batch_single_scan
+from utils.metrics import compute_metrics_mc
+from utils.plots import plot_training_progress_classification, plot_pred_summary_mc
+from utils.utils import compute_class_weights
+
+
+def collect_arguments():
+    """
+    Collects arguments from the command line.
+    """
+    
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--df", type=Path, required=True)
+    parser.add_argument("--savedir", type=Path, required=True)
+    parser.add_argument("--model_name", type=str, default="resnet18", required=True)
+    parser.add_argument("--optimizer", type=str, choices=["adamw", "sgd"], default="adamw", required=True)
+    parser.add_argument("--lr", type=float, default=1e-5, required=True)
+    parser.add_argument("--augment", action="store_true", default=False)
+    parser.add_argument("--accum_steps", type=int, default=4, required=True)
+    parser.add_argument("--batch_size", type=int, default=4, required=True)
+    parser.add_argument("--cnn_wd", type=float, default=1e-3, required=True)
+    parser.add_argument("--debug_subset", action="store_true", default=False)
+    parser.add_argument("--balanced_sampling", action="store_true", default=False)
+    parser.add_argument("--val_fold", type=int, default=4, required=True)
+    args = parser.parse_args()
+    
+    assert args.df.exists(), "Dataset file does not exist"
+    
+    return args.df, args.savedir, args.model_name, args.optimizer, args.lr, args.augment, args.accum_steps, args.batch_size, args.cnn_wd, args.debug_subset, args.balanced_sampling, args.val_fold
+
+
+def main():
+    """
+    Baseline method that uses a simple resnet with a single scan as input.
+    """
+    
+    df, savedir, model_name, optimizer_name, lr, do_augmentation, accum_steps, batch_size, cnn_wd, debug_subset, balanced_sampling, val_fold = collect_arguments()
+    
+    run_id = ''.join([random.choice('0123456789abcdef') for _ in range(6)])
+    save_dir = Path(str(savedir).replace("run_id", run_id))
+    save_dir.mkdir(parents=True, exist_ok=True)
+    save_dir.joinpath("preds").mkdir(parents=True, exist_ok=True)
+    save_dir.joinpath("model_weights").mkdir(parents=True, exist_ok=True)
+    print(f"Saving results to {save_dir}")
+    
+    # Load label file, only use odelia for now
+    df = pd.read_excel(df)
+    df = df[(df["scan_number"] == "mip") & (df["breast_label"] != "n.a.")]
+    df.to_excel(save_dir.joinpath("train_data.xlsx"), index=False)
+
+    # Get tio datasets and loaders   
+    train_loader, val_loader, _, _ = prepare_loaders(
+        df=df, 
+        do_augmentation=do_augmentation, 
+        mode="training", 
+        batch_size=batch_size, 
+        all_scans=False,
+        max_num_scans=1, 
+        finetune_label="breast_label",
+        debug_subset=debug_subset,
+        balanced_sampling=balanced_sampling,
+        num_classes=3,
+        val_fold=val_fold
+    )
+
+    # Loss function and optimizer.
+    n_epochs = 200
+    if not balanced_sampling:
+        class_weights = compute_class_weights(df=df, label="breast_label", dataset="odelia", all_scans=False)
+        criterion = nn.CrossEntropyLoss(weight=class_weights.half())    
+    else:
+        criterion = nn.CrossEntropyLoss()
+    
+    model = Resnet(model_name=model_name, num_classes=3, norm="batch")
+    model.to("cuda")
+    
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=cnn_wd)
+    scheduler = CosineAnnealingLR(optimizer, T_max=n_epochs, eta_min=lr/10)
+    
+    all_train_loss, all_val_loss = [], []
+    all_train_acc, all_val_acc = [], []
+    all_train_auc, all_val_auc = [], []
+    
+    # Same some results for later debugging
+    training_config = {
+        "model_name": model_name,
+        "optimizer_name": optimizer_name,
+        "scheduler": "cosine",
+        "start_lr": lr,
+        "end_lr": lr/10,
+        "n_epochs": n_epochs,
+        "do_augmentation": do_augmentation,
+        "batch_size": batch_size,
+        "run_id": run_id,
+        "balanced_sampling": balanced_sampling,
+        "accum_steps": accum_steps,
+        "cnn_wd": cnn_wd,
+        "finetune_label": "breast_label",
+        "use_amp": True,
+        "mode": "training",
+        "norm": "batch",
+        "val_fold": val_fold
+    }
+    with open(save_dir.joinpath("training_config.json"), "w") as f:
+        json.dump(training_config, f, indent=4, sort_keys=True)
+    
+    scaler = GradScaler("cuda")
+    for epoch in range(n_epochs):
+        
+        model.train()
+        
+        train_epoch_loss = 0
+        train_epoch_preds, train_epoch_probs, train_epoch_gt = [], [], []
+        
+        for batch_idx, batch in enumerate(tqdm(train_loader, desc=f"Epoch {epoch+1}/{n_epochs}")):
+                        
+            # Forward pass          
+            x, y = prepare_batch_single_scan(batch)
+            
+            with autocast(device_type='cuda'):
+                y_pred = model(x)
+                train_loss = criterion(y_pred, y)
+                train_loss /= accum_steps
+            
+            scaler.scale(train_loss).backward()
+            
+            # Gradient accumulation step
+            if (batch_idx + 1) % accum_steps == 0:
+                scaler.step(optimizer)
+                scaler.update()
+                optimizer.zero_grad()
+            
+            probs = torch.softmax(y_pred, dim=1).detach().cpu().numpy()
+            preds = np.argmax(probs, axis=1).tolist()
+            gts = y.detach().cpu().numpy().tolist()
+            
+            train_epoch_probs.extend(probs.tolist())
+            train_epoch_preds.extend(preds)
+            train_epoch_gt.extend(gts)
+            
+            train_epoch_loss += train_loss.item() * accum_steps
+            
+        # Handle remaining gradients
+        if len(train_loader) % accum_steps != 0:
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer.zero_grad()
+        
+        all_train_loss.append(train_epoch_loss / len(train_loader.dataset))
+        train_metrics = compute_metrics_mc(labels=train_epoch_gt, preds=train_epoch_preds, probs=train_epoch_probs, num_classes=3)
+        all_train_acc.append(train_metrics["balanced_accuracy"])
+        all_train_auc.append(train_metrics["auc"])
+        
+        # Validation
+        model.eval()
+        val_epoch_loss = 0
+        val_epoch_preds, val_epoch_probs, val_epoch_gt = [], [], []
+        
+        with torch.no_grad():
+            for batch in tqdm(val_loader, desc=f"Epoch {epoch+1}/{n_epochs} - Val"):
+                
+                x, y = prepare_batch_single_scan(batch)
+                
+                with autocast(device_type='cuda'):  
+                    y_pred = model(x)
+                    val_loss = criterion(y_pred, y)
+                
+                probs = torch.softmax(y_pred, dim=1).detach().cpu().numpy()
+                preds = np.argmax(probs, axis=1).tolist()
+                gts = y.detach().cpu().numpy().tolist()
+                
+                val_epoch_loss += val_loss.item()
+                val_epoch_probs.extend(probs.tolist())
+                val_epoch_preds.extend(preds)
+                val_epoch_gt.extend(gts)
+                
+        all_val_loss.append(val_epoch_loss / len(val_loader.dataset))
+        val_metrics = compute_metrics_mc(labels=val_epoch_gt, preds=val_epoch_preds, probs=val_epoch_probs, num_classes=3)
+        all_val_acc.append(val_metrics["balanced_accuracy"])
+        all_val_auc.append(val_metrics["auc"])
+        
+        scheduler.step()
+        
+        # Update training progress
+        save_path = save_dir.joinpath(f"progress.png")
+        plot_training_progress_classification(
+            all_train_loss=all_train_loss, 
+            all_val_loss=all_val_loss, 
+            all_train_acc=all_train_acc, 
+            all_val_acc=all_val_acc, 
+            all_train_auc=all_train_auc, 
+            all_val_auc=all_val_auc, 
+            save_path=save_path
+        )
+            
+        if all_val_auc[-1] == max(all_val_auc):
+            torch.save(model.state_dict(), save_dir.joinpath("model_weights", "best_model.pth"))
+    
+        if epoch % 5 == 0:
+            torch.save(model.state_dict(), save_dir.joinpath("model_weights", f"epoch_{str(epoch).zfill(3)}_model.pth"))
+            plot_pred_summary_mc(
+                preds=train_epoch_preds, 
+                probs=train_epoch_probs, 
+                gts=train_epoch_gt, 
+                n_classes=3,
+                save_path=save_dir.joinpath("preds", f"epoch_{str(epoch).zfill(3)}_train_preds.png")
+            )
+            plot_pred_summary_mc(
+                preds=val_epoch_preds, 
+                probs=val_epoch_probs, 
+                gts=val_epoch_gt, 
+                n_classes=3, 
+                save_path=save_dir.joinpath("preds", f"epoch_{str(epoch).zfill(3)}_val_preds.png")
+            )
+        
+        metric_df = pd.DataFrame({
+            "epoch": list(range(1, epoch+2)),
+            "train_loss": all_train_loss,
+            "val_loss": all_val_loss,
+            "train_acc": all_train_acc,
+            "val_acc": all_val_acc,
+            "train_auc": all_train_auc,
+            "val_auc": all_val_auc
+        })
+        metric_df.to_excel(save_dir.joinpath("metrics.xlsx"), index=False)
+        
+        print(f"Epoch {epoch+1}/{n_epochs}: > Train Loss: {all_train_loss[-1]:.4f} - Val Loss: {all_val_loss[-1]:.4f}") 
+        print(f"Epoch {epoch+1}/{n_epochs}: > Train Acc: {all_train_acc[-1]:.4f} - Val Acc: {all_val_acc[-1]:.4f}")
+        print(f"Epoch {epoch+1}/{n_epochs}: > Train AUC: {all_train_auc[-1]:.4f} - Val AUC: {all_val_auc[-1]:.4f}\n")
+        
+        # Clean up
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    return
+
+
+if __name__ == "__main__":
+    main()

training_code/utils/dataloader.py

@@ -0,0 +1,165 @@
+import pandas as pd
+import numpy as np
+from tqdm import tqdm
+from pathlib import Path
+import torch
+import torchio as tio
+import SimpleITK as sitk
+from typing import List, Tuple
+
+from utils.transforms import get_image_transforms
+
+
+def create_subjects_list(df: pd.DataFrame, all_scans: bool = False, finetune_label: str = None, dataset: str = "odelia", debug_subset: bool = False) -> List[tio.Subject]:
+    """
+    Create a list of subjects from a dataframe of file paths and labels.
+    """
+    
+    subjects = []
+    patient_ids = df["patient_id"].unique().tolist()
+    
+    # Sample 5 patients from each class in debug mode
+    if debug_subset:
+        if finetune_label is not None:
+            patient_ids = []
+            for cls in df[finetune_label].dropna().unique().tolist():
+                cls_patients = df[df[finetune_label] == cls]["patient_id"].unique()
+                n_select = min(5, len(cls_patients))
+                patient_ids.extend(np.random.choice(cls_patients, size=n_select, replace=False).tolist())
+        else:
+            patient_ids = patient_ids[:10]
+    
+    if dataset == "odelia":
+        odelia_mapper = {"normal": 0, "benign": 1, "malignant": 2}
+    
+    for patient_id in tqdm(patient_ids, total=len(patient_ids)):
+        
+        pt_df = df[df["patient_id"] == patient_id]
+        
+        # Get all breast volumes per patient
+        for side in pt_df["side"].unique().tolist():
+            side_df = pt_df[pt_df["side"] == side]
+            side_df = side_df.sort_values(by="filepath")
+            
+            if finetune_label == "breast_label":
+                side_label = torch.tensor(odelia_mapper[side_df["breast_label"].values[0]], dtype=torch.float32)
+            else:
+                side_label = "n.a."
+            
+            # Use a single scan per patient
+            if not all_scans:
+                side_images = [tio.ScalarImage(side_df["filepath"].values[0])]
+                side_image_keys = ["image_1"]
+            
+            side_subject_dict = tio.Subject(name=f"{patient_id}_{side}", **{key: value for key, value in zip(side_image_keys, side_images)}, label=side_label)
+            subjects.append(side_subject_dict)
+            
+    return subjects
+
+
+def run_sanity_check_single_image_dataloader(dataloader: tio.data.SubjectsLoader, num_samples: int, save_dir: Path):
+    """
+    Sanity check the dataloader.
+    """
+    
+    save_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Extract samples
+    examples = []
+    
+    for batch in dataloader:
+        examples.extend(batch)
+        if len(examples) == num_samples:
+            break
+    
+    # Save them to visually inspect train transforms
+    for example in examples:
+        image = sitk.GetImageFromArray(example["image_1"][tio.DATA].squeeze(0).numpy())
+        label = int(example["label"].numpy())
+        sitk.WriteImage(image, save_dir.joinpath(f"{example['name']}_(gt={label}).nii.gz"))
+        
+    return
+
+
+def create_balanced_sampler(subjects, num_classes: int = 3):
+    """
+    Create a balanced sampler to handle class imbalance.
+    """
+    
+    labels = [subject["label"].item() for subject in subjects]
+    unique, counts = np.unique(labels, return_counts=True)
+    
+    if len(unique) != num_classes:
+        print(f"WARNING: Number of classes ({len(unique)}) does not match expected number of classes ({num_classes}) in balanced sampler")
+    
+    # Calculate weights for each sample
+    class_weights = {label: 1.0 / count for label, count in zip(unique, counts)}
+    sample_weights = [class_weights[label] for label in labels]
+    
+    sampler = torch.utils.data.WeightedRandomSampler(
+        weights=sample_weights,
+        num_samples=len(sample_weights),
+        replacement=True
+    )
+    
+    return sampler
+
+
+def prepare_batch_single_scan(batch: List[tio.Subject]) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Helper function to prepare a batch of subjects for training.
+    """
+    
+    x = batch["image_1"][tio.DATA].to(device='cuda', dtype=torch.float16, non_blocking=True)
+    y = batch["label"].to(device='cuda', dtype=torch.long, non_blocking=True)
+    
+    return x, y
+
+
+def prepare_loaders(
+        df: pd.DataFrame, 
+        do_augmentation: bool, 
+        mode: str, 
+        batch_size: int, 
+        all_scans: bool,
+        max_num_scans: int, 
+        finetune_label: str = None, 
+        debug_subset: bool = False, 
+        balanced_sampling: bool = False,
+        num_classes: int = None,
+        val_fold: int = 0
+    ):
+    """
+    Helper function to prepare data loaders.
+    """
+    
+    # Get transforms
+    train_transforms, val_transforms = get_image_transforms(do_augmentation=do_augmentation)
+    
+    # Get data loaders
+    print(f"Preparing data loaders...")
+    
+    all_folds = [0, 1, 2, 3, 4]
+    train_folds = [i for i in all_folds if i != val_fold]
+    
+    train_subjects = create_subjects_list(df[df["fold"].isin(train_folds)], all_scans=all_scans, finetune_label=finetune_label, debug_subset=debug_subset)
+    train_subjects_names = [i["name"] for i in train_subjects]
+    train_dataset = tio.data.SubjectsDataset(train_subjects, transform=train_transforms, load_getitem=False)
+    
+    if balanced_sampling and num_classes is not None:
+        train_sampler = create_balanced_sampler(train_subjects, num_classes=num_classes)
+        train_loader = tio.data.SubjectsLoader(train_dataset, batch_size=batch_size, sampler=train_sampler, num_workers=6, prefetch_factor=4, persistent_workers=True)
+    else:
+        train_loader = tio.data.SubjectsLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=6, prefetch_factor=4, persistent_workers=True)
+    
+    val_subjects = create_subjects_list(df[df["fold"].isin([val_fold])], all_scans=all_scans, finetune_label=finetune_label, debug_subset=debug_subset)
+    val_dataset = tio.data.SubjectsDataset(val_subjects, transform=val_transforms, load_getitem=False)
+    val_loader = tio.data.SubjectsLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=6, prefetch_factor=4, persistent_workers=True)
+    
+    # Sanity loader to fetch attention weights from some training and val cases
+    sanity_subjects = np.random.choice(train_subjects, size=5, replace=False).tolist() + np.random.choice(val_subjects, size=5, replace=False).tolist()
+    sanity_dataset = tio.data.SubjectsDataset(sanity_subjects, transform=train_transforms, load_getitem=False)
+    sanity_loader = tio.data.SubjectsLoader(sanity_dataset, batch_size=batch_size, shuffle=False, num_workers=6, prefetch_factor=4, persistent_workers=True)
+    
+    return train_loader, val_loader, sanity_loader, train_subjects_names
+    

training_code/utils/metrics.py

@@ -0,0 +1,104 @@
+import numpy as np
+from typing import Dict, Any
+from sklearn.metrics import balanced_accuracy_score, roc_auc_score, roc_curve, confusion_matrix
+from sklearn.preprocessing import label_binarize 
+
+
+def compute_metrics_bc(labels: np.ndarray, preds: np.ndarray, probs: np.ndarray) -> Dict[str, Any]:
+    """
+    Compute all relevant metrics for binary classification.
+    
+    Args:
+        labels: Ground truth binary labels (0 or 1)
+        preds: Binary predictions (0 or 1)
+        probs: Prediction probabilities (0-1)
+        
+    Returns:
+        Dictionary containing all calculated metrics
+    """
+    
+    metrics = {}
+    
+    # Balanced accuracy
+    metrics['balanced_accuracy'] = balanced_accuracy_score(labels, preds)
+    
+    # AUC
+    metrics['auc'] = roc_auc_score(labels, probs)
+    
+    # Confusion matrix for sensitivity and specificity
+    cm = confusion_matrix(labels, preds)
+    metrics['confusion_matrix'] = cm
+    tn, fp, fn, tp = cm.ravel()
+    
+    # Sensitivity and specificity
+    metrics['sensitivity'] = tp / (tp + fn) if (tp + fn) > 0 else 0
+    metrics['specificity'] = tn / (tn + fp) if (tn + fp) > 0 else 0
+    
+    # ROC curve data
+    metrics['roc_curve'] = roc_curve(labels, probs)
+    
+    return metrics
+
+
+def compute_metrics_mc(labels: np.ndarray, preds: np.ndarray, probs: np.ndarray, num_classes: int = 3) -> Dict[str, Any]:
+    """
+    Compute relevant metrics for multiclass classification (default: 3 classes).
+
+    Args:
+        labels: Ground-truth class indices with shape (N,), values in \{0, ..., num_classes-1\}
+        preds: Predicted class indices with shape (N,)
+        probs: Predicted class probabilities with shape (N, num_classes)
+        num_classes: Number of classes present in the problem (default 3)
+
+    Returns:
+        Dictionary mapping metric names to their computed values.
+    """
+
+    metrics = {}
+
+    # Normalize probabilities to handle floating point precision errors from mixed precision training
+    probs = np.array(probs)
+    probs = probs / probs.sum(axis=1, keepdims=True)
+    labels = np.array(labels)
+
+    # Overall balanced accuracy (macro-averaged recall)
+    metrics["balanced_accuracy"] = balanced_accuracy_score(labels, preds)
+
+    # Macro-averaged ROC-AUC using the one-vs-rest strategy
+    labels_auc = label_binarize(labels, classes=list(range(num_classes)))
+    metrics["auc"] = roc_auc_score(labels_auc, probs, average="micro")
+
+    # Confusion matrix (shape: num_classes × num_classes)
+    cm = confusion_matrix(labels, preds, labels=list(range(num_classes)))
+    metrics["confusion_matrix"] = cm
+
+    # Per-class sensitivity (recall) and specificity
+    sensitivities = []
+    specificities = []
+    total = cm.sum()
+    for c in range(num_classes):
+        tp = cm[c, c]
+        fn = cm[c, :].sum() - tp
+        fp = cm[:, c].sum() - tp
+        tn = total - (tp + fp + fn)
+
+        sens = tp / (tp + fn) if (tp + fn) > 0 else 0.0
+        spec = tn / (tn + fp) if (tn + fp) > 0 else 0.0
+
+        sensitivities.append(sens)
+        specificities.append(spec)
+
+    metrics["sensitivity"] = np.array(sensitivities)
+    metrics["specificity"] = np.array(specificities)
+
+    # ROC curve per class (fpr, tpr, thresholds for each class)
+    roc_curves = []
+    for c in range(num_classes):
+        try:
+            fpr, tpr, thresholds = roc_curve((np.array(labels) == c).astype(int), np.array(probs)[:, c])
+            roc_curves.append((fpr, tpr, thresholds))
+        except ValueError:
+            roc_curves.append((np.array([]), np.array([]), np.array([])))
+    metrics["roc_curve"] = roc_curves
+
+    return metrics

training_code/utils/plots.py

@@ -0,0 +1,282 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from sklearn.metrics import balanced_accuracy_score, roc_auc_score, roc_curve, confusion_matrix, auc
+
+    
+def moving_average(data, window_size=5):
+    """
+    Compute moving average over specified window size.
+    """
+    
+    if len(data) < window_size:
+        return data
+    
+    return np.convolve(data, np.ones(window_size)/window_size, mode='valid')
+
+
+def plot_training_progress_classification(all_train_loss, all_val_loss, all_train_acc, all_val_acc, 
+                          all_train_auc, all_val_auc, save_path, window_size=5):
+    """
+    Plot training progress with both raw epoch-by-epoch data and moving averages.
+    
+    Args:
+        all_train_loss: List of training losses per epoch
+        all_val_loss: List of validation losses per epoch  
+        all_train_acc: List of training accuracies per epoch
+        all_val_acc: List of validation accuracies per epoch
+        all_train_auc: List of training AUCs per epoch
+        all_val_auc: List of validation AUCs per epoch
+        save_path: Path to save the plot
+        window_size: Window size for moving average (default: 5)
+    """
+    
+    plt.figure(figsize=(10, 5))
+    epochs = np.arange(1, len(all_train_loss) + 1)
+    
+    # Loss subplot
+    plt.subplot(1, 3, 1)
+    train_loss_line = plt.plot(epochs, all_train_loss, '--', alpha=0.6, label="Train Loss (raw)")[0]
+    val_loss_line = plt.plot(epochs, all_val_loss, '--', alpha=0.6, label="Val Loss (raw)")[0]
+    if len(all_train_loss) >= window_size:
+        ma_epochs = np.arange(window_size, len(all_train_loss) + 1)
+        plt.plot(ma_epochs, moving_average(all_train_loss, window_size), '-', 
+                color=train_loss_line.get_color(), label=f"Train Loss (MA-{window_size})")
+        plt.plot(ma_epochs, moving_average(all_val_loss, window_size), '-', 
+                color=val_loss_line.get_color(), label=f"Val Loss (MA-{window_size})")
+    plt.legend()
+    plt.xlabel("Epoch")
+    plt.ylabel("Loss")
+    
+    # Accuracy subplot
+    plt.subplot(1, 3, 2)
+    train_acc_line = plt.plot(epochs, all_train_acc, '--', alpha=0.6, label="Train Acc (raw)")[0]
+    val_acc_line = plt.plot(epochs, all_val_acc, '--', alpha=0.6, label="Val Acc (raw)")[0]
+    if len(all_train_acc) >= window_size:
+        ma_epochs = np.arange(window_size, len(all_train_acc) + 1)
+        plt.plot(ma_epochs, moving_average(all_train_acc, window_size), '-', 
+                color=train_acc_line.get_color(), label=f"Train Acc (MA-{window_size})")
+        plt.plot(ma_epochs, moving_average(all_val_acc, window_size), '-', 
+                color=val_acc_line.get_color(), label=f"Val Acc (MA-{window_size})")
+    plt.legend()
+    plt.xlabel("Epoch")
+    plt.ylabel("Accuracy")
+    
+    # AUC subplot
+    plt.subplot(1, 3, 3)
+    train_auc_line = plt.plot(epochs, all_train_auc, '--', alpha=0.6, label="Train AUC (raw)")[0]
+    val_auc_line = plt.plot(epochs, all_val_auc, '--', alpha=0.6, label="Val AUC (raw)")[0]
+    if len(all_train_auc) >= window_size:
+        ma_epochs = np.arange(window_size, len(all_train_auc) + 1)
+        plt.plot(ma_epochs, moving_average(all_train_auc, window_size), '-', 
+                color=train_auc_line.get_color(), label=f"Train AUC (MA-{window_size})")
+        plt.plot(ma_epochs, moving_average(all_val_auc, window_size), '-', 
+                color=val_auc_line.get_color(), label=f"Val AUC (MA-{window_size})")
+    plt.legend()
+    plt.xlabel("Epoch")
+    plt.ylabel("AUC")
+    
+    plt.tight_layout()
+    plt.savefig(save_path)
+    plt.close() 
+    
+    return
+
+
+def plot_pred_summary_bc(preds, probs, gts, save_path):
+    """
+    Plot predictions for binary classification with confusion matrix, performance metrics, and ROC curve.
+    """
+    
+    probs = np.array(probs)
+    preds = np.array(preds)
+    gts = np.array(gts)
+    
+    # Compute metrics
+    balanced_accuracy = balanced_accuracy_score(gts, preds)
+    auc = roc_auc_score(gts, probs)
+    
+    # Confusion matrix
+    cm = confusion_matrix(gts, preds)
+    tn, fp, fn, tp = cm.ravel()
+    
+    # Additional metrics
+    sensitivity = tp / (tp + fn) if (tp + fn) > 0 else 0
+    specificity = tn / (tn + fp) if (tn + fp) > 0 else 0
+    ppv = tp / (tp + fp) if (tp + fp) > 0 else 0
+    npv = tn / (tn + fn) if (tn + fn) > 0 else 0
+    
+    # ROC curve
+    fpr, tpr, _ = roc_curve(gts, probs[:, 1] if probs.ndim == 2 else probs)
+    
+    # Create figure with 3 subplots
+    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(15, 5))
+    
+    # Confusion matrix
+    im = ax1.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)
+    ax1.set_title('Confusion Matrix')
+    ax1.set_xlabel('Predicted')
+    ax1.set_ylabel('Actual')
+    ax1.set_xticks([0, 1])
+    ax1.set_yticks([0, 1])
+    ax1.set_xticklabels(['Negative', 'Positive'])
+    ax1.set_yticklabels(['Negative', 'Positive'])
+    
+    # Add text annotations
+    thresh = cm.max() / 2
+    for i in range(2):
+        for j in range(2):
+            ax1.text(j, i, format(cm[i, j], 'd'),
+                    ha="center", va="center",
+                    color="white" if cm[i, j] > thresh else "black")
+    
+    # Performance metrics bar chart
+    metrics = ['Sensitivity', 'Specificity', 'PPV', 'NPV', 'Balanced Accuracy']
+    values = [sensitivity, specificity, ppv, npv, balanced_accuracy]
+    bars = ax2.bar(metrics, values, color=['skyblue', 'lightcoral', 'lightgreen', 'gold', 'lightpink'])
+    ax2.set_title('Performance Metrics')
+    ax2.set_ylabel('Value')
+    ax2.set_ylim(0, 1)
+    
+    # Add value labels on bars
+    for bar, value in zip(bars, values):
+        height = bar.get_height()
+        ax2.text(bar.get_x() + bar.get_width()/2., height + 0.01,
+                f'{value:.3f}', ha='center', va='bottom')
+    
+    # ROC curve
+    ax3.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (AUC = {auc:.3f})')
+    ax3.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
+    ax3.set_xlim([0.0, 1.0])
+    ax3.set_ylim([0.0, 1.05])
+    ax3.set_xlabel('False Positive Rate')
+    ax3.set_ylabel('True Positive Rate')
+    ax3.set_title('ROC Curve')
+    ax3.legend(loc="lower right")
+    ax3.grid(True)
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=150, bbox_inches='tight')
+    plt.close()
+    
+    return
+
+
+def plot_pred_summary_mc(preds, probs, gts, n_classes, save_path):
+    """
+    Plot predictions for multiclass classification with confusion matrix and performance metrics.
+    """
+    
+    # Normalize probabilities to handle floating point precision errors from mixed precision training
+    probs = np.array(probs)
+    probs = probs / probs.sum(axis=1, keepdims=True)
+    
+    preds = np.array(preds)
+    gts = np.array(gts)
+    
+    # Compute metrics
+    balanced_accuracy = balanced_accuracy_score(gts, preds)
+    auc = roc_auc_score(gts, probs, multi_class='ovo', average='macro', labels=list(range(n_classes)))
+    
+    # Confusion matrix
+    cm = confusion_matrix(gts, preds, labels=list(range(n_classes)))
+    
+    # Create figure with 2 subplots for multiclass
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
+    
+    # Confusion matrix
+    im = ax1.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)
+    ax1.set_title('Confusion Matrix')
+    ax1.set_xlabel('Predicted Class')
+    ax1.set_ylabel('True Class')
+    ax1.set_xticks(range(n_classes))
+    ax1.set_yticks(range(n_classes))
+    ax1.set_xticklabels([f'Class {i}' for i in range(n_classes)])
+    ax1.set_yticklabels([f'Class {i}' for i in range(n_classes)])
+    
+    # Add text annotations
+    thresh = cm.max() / 2
+    for i in range(n_classes):
+        for j in range(n_classes):
+            ax1.text(j, i, format(cm[i, j], 'd'),
+                    ha="center", va="center",
+                    color="white" if cm[i, j] > thresh else "black")
+    
+    # Performance metrics bar chart
+    metrics = ['Balanced Accuracy', 'Macro AUC']
+    values = [balanced_accuracy, auc]
+    bars = ax2.bar(metrics, values, color=['lightpink', 'lightblue'])
+    ax2.set_title('Performance Metrics')
+    ax2.set_ylabel('Value')
+    ax2.set_ylim(0, 1)
+    
+    # Add value labels on bars
+    for bar, value in zip(bars, values):
+        height = bar.get_height()
+        ax2.text(bar.get_x() + bar.get_width()/2., height + 0.01,
+                f'{value:.3f}', ha='center', va='bottom')
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=150, bbox_inches='tight')
+    plt.close()
+    
+    return
+
+
+def plot_multiclass_roc_curve(roc_curve_data, save_path=None):
+    """
+    Plot the ROC curve for multiclass classification.
+    """
+
+    n_classes = len(roc_curve_data)
+    plt.figure(figsize=(n_classes * 5, 5))
+
+    for c, (fpr, tpr, _) in enumerate(roc_curve_data):
+        plt.subplot(1, n_classes, c + 1)
+        plt.plot(fpr, tpr, label=f'ROC curve of class {c}')
+        plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
+        plt.xlim([0.0, 1.0])
+        plt.ylim([0.0, 1.0])
+        auc_val = auc(fpr, tpr)
+        plt.legend([f"AUC = {auc_val:.2f}"])
+        plt.xlabel("False Positive Rate")
+        plt.ylabel("True Positive Rate")
+        plt.title(f"ROC Curve of Class {c}")
+        plt.grid(True)
+
+    plt.savefig(save_path, dpi=150, bbox_inches="tight")
+    plt.close()
+
+    return
+    
+def plot_multiclass_confusion_matrix(cm, save_path=None):
+    """
+    Plot the confusion matrix.
+    """
+    
+    n_classes = cm.shape[0]
+    fig, ax1 = plt.subplots(1, 1, figsize=(n_classes * 3, n_classes * 3))
+    
+    # Confusion matrix
+    im = ax1.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)
+    ax1.set_title('Confusion Matrix')
+    ax1.set_xlabel('Predicted Class')
+    ax1.set_ylabel('True Class')
+    ax1.set_xticks(range(n_classes))
+    ax1.set_yticks(range(n_classes))
+    ax1.set_xticklabels([f'Class {i}' for i in range(n_classes)])
+    ax1.set_yticklabels([f'Class {i}' for i in range(n_classes)])
+    
+    # Add text annotations
+    thresh = cm.max() / 2
+    for i in range(n_classes):
+        for j in range(n_classes):
+            ax1.text(j, i, format(cm[i, j], 'd'),
+                    ha="center", va="center",
+                    color="white" if cm[i, j] > thresh else "black")
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=150, bbox_inches='tight')
+    plt.close()
+    
+    return
+    

training_code/utils/transforms.py

@@ -0,0 +1,87 @@
+import torchio as tio
+import numpy as np
+import torch
+
+
+# Helper function for multiprocessing-safe casting
+def _to_float32(tensor: torch.Tensor) -> torch.Tensor:
+    """Cast TorchIO tensor to float32 (multiprocessing-safe)."""
+    return tensor.to(torch.float32)
+
+
+class PadToMultiple(tio.Transform):
+    """
+    Pads a 3D image so that each spatial dimension is a multiple of a given number.
+    """
+    def __init__(self, multiple: int = 16, **kwargs):
+        super().__init__(**kwargs)
+        self.multiple = multiple
+
+    def apply_transform(self, subject: tio.Subject) -> tio.Subject:
+        image_keys = sorted([i for i in subject.keys() if i.startswith("image_")])
+        for image_key in image_keys:
+            image = subject[image_key]
+            shape = image.spatial_shape
+            
+            target_shape = [int(np.ceil(dim / self.multiple)) * self.multiple for dim in shape]
+
+            pad_transform = tio.CropOrPad(
+                target_shape=tuple(target_shape), padding_mode=0
+            )
+            subject = pad_transform(subject)
+            
+        return subject 
+
+
+def pad_fixed_size_collate_fn(batch):
+    """
+    Pads each image in the batch to the same size and returns a tensor of images and a tensor of labels.
+    Assumes that the subject only contains one image and one label.
+    """
+    
+    # Find the maximum spatial dimensions in the batch
+    max_shape = np.max([item['image_1'][tio.DATA].shape[1:] for item in batch], axis=0)
+
+    # Make it a multiple of 16
+    target_shape = tuple([int(np.ceil(dim / 16)) * 16 for dim in max_shape])
+    pad_transform = tio.CropOrPad(
+        target_shape=target_shape, padding_mode=0
+    )
+    
+    # Pad each image to the max shape
+    for c, item in enumerate(batch):
+        image_keys = sorted([i for i in item.keys() if i.startswith("image_")])
+        for image_key in image_keys:
+            image = item[image_key]
+            padded_image = pad_transform(image)
+            batch[c][image_key] = padded_image
+    
+    return batch
+
+
+
+def get_image_transforms(do_augmentation: bool = True):
+    """
+    Get the image transforms for the training and validation sets.
+    """
+    
+    if do_augmentation:
+        train_transforms = tio.Compose([
+            tio.RandomFlip(axes=('LR', 'AP', 'IS'), flip_probability=0.5),
+            tio.RandomAffine(scales=(0.8, 1.2), degrees=45, translation=(15, 15, 15), p=0.3),
+            tio.RandomNoise(std=(0.01, 0.10), p=0.2),
+            # tio.RandomBiasField(coefficients=0.5, p=0.2),
+            # tio.RandomBlur(std=(0, 2), p=0.2),
+            tio.RandomGamma(log_gamma=(-0.3, 0.3), p=0.2),
+            tio.RandomSwap(patch_size=10, num_iterations=30, p=0.2)
+        ])
+    else:
+        train_transforms = tio.Compose([
+            tio.RandomFlip(axes=('LR', 'AP', 'IS'), flip_probability=0.5),
+            tio.RandomAffine(scales=(0.8, 1.2), degrees=45, translation=(15, 15, 15), p=0.3),
+        ])
+    
+    # Add type transform to ensure image data loading for validation
+    val_transforms = tio.Compose([tio.Lambda(_to_float32)])
+    
+    return train_transforms, val_transforms

training_code/utils/utils.py

@@ -0,0 +1,60 @@
+import torch
+import torch.nn as nn
+import pandas as pd
+
+
+def get_optimizer(model, cnn_wd, transformer_wd, cnn_lr, transformer_lr, general_lr):
+    """
+    Helper function to set different weight decay for CNN and transformers.
+    """
+    
+    # Figure out which parameters need decay
+    decay, no_decay = set(), set()
+    for mn, m in model.named_modules():
+        for pn, p in m.named_parameters(recurse=False):
+            fqn = f"{mn}.{pn}" if mn else pn
+            if isinstance(m, (nn.LayerNorm, nn.BatchNorm1d, nn.BatchNorm2d)) or pn.endswith("bias"):
+                no_decay.add(fqn)
+            else:
+                decay.add(fqn)
+                
+    # Set different weight decay for CNN and transformers
+    param_dict = {pn: p for pn, p in model.named_parameters()}
+    optimizer = torch.optim.AdamW(
+        [
+            {"params": [param_dict[n] for n in sorted(decay) if "transformer" not in n], "weight_decay": cnn_wd, "lr": cnn_lr},
+            {"params": [param_dict[n] for n in sorted(decay) if "transformer" in n], "weight_decay": transformer_wd, "lr": transformer_lr},
+            {"params": [param_dict[n] for n in sorted(no_decay)], "weight_decay": 0.0, "lr": general_lr},
+        ]
+    )
+    
+    return optimizer
+
+
+def compute_class_weights(df: pd.DataFrame, label: str, dataset: str, all_scans: bool = False):
+    """
+    Compute the class weights for the given label.
+    """
+    
+    # Make a copy to avoid SettingWithCopyWarning
+    df = df.copy()
+    
+    # Map values to range [0, 2]
+    if dataset == "odelia":
+        odelia_mapper = {"normal": 0, "benign": 1, "malignant": 2}
+        df["breast_label"] = df["breast_label"].map(odelia_mapper)
+
+    # Only select a single scan per patient
+    if all_scans:
+        df = df[df["scan_number"] == 1]
+
+    # Get counts for how often each class occurs
+    class_counts = df[label].value_counts().sort_index()
+    
+    # Account for multiclass and binary classification
+    if dataset == "odelia":
+        class_weights = class_counts.max() / class_counts
+        return torch.tensor(class_weights.values, dtype=torch.float32, device="cuda")
+    else:
+        class_weights = class_counts[0] / class_counts[1]
+        return torch.tensor(class_weights, dtype=torch.float32, device="cuda")