Skip to main content

Preparing the data

Our goal in this step: turn the UCSF-BMSR MRI volumes into something a 2D model can train on — pairs of (image slice, mask slice) — and wrap them in a PyTorch Dataset.

What the dataset gives you​

The UCSF-BMSR dataset provides, per subject, a set of 3D MRI volumes and an expert segmentation mask marking the brain-metastasis lesions. We use a single input sequence:

  • T1post — the T1-weighted post-contrast scan. Metastases enhance brightly with contrast, which is why this sequence is the workhorse for detecting them.

The volumes are 3D (a stack of 2D slices). The mask is a matching 3D volume where lesion voxels are labelled 1 and everything else 0.

The layout on disk​

After downloading and unzipping, the training set is a folder of per-subject directories:

UCSF_BrainMetastases_v1.3/
└── UCSF_BrainMetastases_TRAIN/
    ├── 100101A/
    │   ├── 100101A_T1post.nii.gz       ← T1 post-contrast  (our input)
    │   ├── 100101A_seg.nii.gz          ← metastasis mask   (our target)
    │   ├── 100101A_T1pre.nii.gz        ← other sequences, unused here
    │   ├── 100101A_FLAIR.nii.gz
    │   ├── 100101A_T2Synth.nii.gz
    │   ├── 100101A_subtraction.nii.gz
    │   └── 100101A_BraTS-seg.nii.gz    ← BraTS-format labels (alternative mask)
    ├── 100102A/
    └── ...                             (several hundred studies; some patients have A/B timepoints)

Every file is NIfTI (.nii.gz). We use exactly two per subject: *_T1post.nii.gz as the input and *_seg.nii.gz as the target mask. (Note the filename is T1post — lower-case "post".)

note

Whether volumes are skull-stripped, and their exact intensity ranges, depend on the release; the benchmarks repo also provides skull-stripping models. Treat the dataset page as authoritative for the data itself.

From 3D volumes to 2D slices​

A 2D model looks at one slice at a time, so we walk the axial slices of each subject and pair each image slice with its mask slice. Two preprocessing choices matter:

  • Normalisation. MRI intensities aren't calibrated between scanners, so we z-score normalise each volume using its brain (non-zero) voxels.
  • Class imbalance. Brain-mets slices are mostly background; training on every slice teaches the model to predict "nothing". Keeping only lesion-bearing slices (lesion_only=True) is the simplest fix.
Consistency across sites matters

In a federation, every site must preprocess the same way — same sequence, same normalisation, same slice size. The models are averaged together, so they must be learning on comparable inputs. Agree the recipe with your collaborators up front.

The Dataset​

We pre-load the wanted slices into memory so training is fast. Each item is an (image, mask) pair of shape (1, H, W) — one channel in, one channel out (lesion vs background).

import glob, os
import numpy as np
import nibabel as nib
import torch
from torch.utils.data import Dataset, DataLoader

IMG_SIZE = 256  # UCSF-BMSR volumes are already 256x256 in-plane


def zscore(vol):
    """Z-score normalise using non-zero (brain) voxels."""
    brain = vol[vol > 0]
    if brain.size == 0:
        return vol
    return (vol - brain.mean()) / (brain.std() + 1e-8)


def fit_size(arr, size=IMG_SIZE):
    """Centre-crop or pad a 2D array to (size, size)."""
    h, w = arr.shape
    out = np.zeros((size, size), dtype=arr.dtype)
    hh, ww = min(h, size), min(w, size)
    sy, sx = (h - hh) // 2, (w - ww) // 2
    dy, dx = (size - hh) // 2, (size - ww) // 2
    out[dy:dy + hh, dx:dx + ww] = arr[sy:sy + hh, sx:sx + ww]
    return out


class BrainMetsSlices(Dataset):
    """2D axial T1post slices with their (binary) lesion masks, pre-loaded."""

    def __init__(self, subject_dirs, lesion_only=True, size=IMG_SIZE):
        self.images, self.masks = [], []
        for subj in subject_dirs:
            t1 = glob.glob(os.path.join(subj, "*_T1post.nii.gz"))
            sg = glob.glob(os.path.join(subj, "*_seg.nii.gz"))
            if not (t1 and sg):
                continue
            vol = zscore(nib.load(t1[0]).get_fdata().astype("float32"))
            seg = (nib.load(sg[0]).get_fdata() > 0).astype("float32")
            for z in range(seg.shape[2]):
                if lesion_only and seg[:, :, z].sum() == 0:
                    continue
                self.images.append(fit_size(vol[:, :, z], size))
                self.masks.append(fit_size(seg[:, :, z], size))

    def __len__(self):
        return len(self.images)

    def __getitem__(self, i):
        x = torch.from_numpy(self.images[i]).unsqueeze(0).float()   # (1, H, W)
        y = torch.from_numpy(self.masks[i]).unsqueeze(0).float()    # (1, H, W)
        return x, y

Building the loaders​

Point DATA_DIR at your UCSF_BrainMetastases_TRAIN folder, then split subjects into train/val. We use a small subset so it trains on a laptop; raise it for real work.

DATA_DIR = os.environ.get("UCSF_BMSR_DIR", "UCSF_BrainMetastases_v1.3/UCSF_BrainMetastases_TRAIN")
subjects = sorted(d for d in glob.glob(os.path.join(DATA_DIR, "*")) if os.path.isdir(d))

subset = subjects[:16]                      # increase for real training
split = max(1, int(0.8 * len(subset)))
train_subj, val_subj = subset[:split], subset[split:]

train_ds = BrainMetsSlices(train_subj, lesion_only=True)
val_ds = BrainMetsSlices(val_subj, lesion_only=True)
print("train slices:", len(train_ds), "| val slices:", len(val_ds))

train_loader = DataLoader(train_ds, batch_size=8, shuffle=True, num_workers=0)
val_loader = DataLoader(val_ds, batch_size=8, shuffle=False, num_workers=0)

A couple of deliberate simplifications, so you can improve on them: binary masks (any lesion label → 1), lesion_only=True, and a modest subset. None of them affect how the federation works.

Next​

You now have training data as clean 2D tensors. Next, the model that consumes them:

➡️ Building and training the model