Skip to main content

Brain metastasis segmentation (UCSF-BMSR)

A complete, worked walkthrough: take a real brain-MRI dataset and a segmentation model, and run it as a federated learning project on BranchKey — step by step, from the ground up.

Read it in order, adapt the code to your paths and model, and by the end you will have a working federated segmentation client.

Runnable companion repository

All the code on these pages lives in a ready-to-run repository with a Makefile and a Jupyter notebook: branchkey/demo-applications/brain-metastases-stereotactic-radiosurgery. Clone it, make run, and follow along in the browser.

The task​

  • Task: segment brain metastases on MRI.
  • Dataset: the UCSF Brain Metastases Stereotactic Radiosurgery (UCSF-BMSR) MRI dataset.
  • Input sequence: the T1-weighted post-contrast scan (referred to as T1Post in the dataset and benchmarks). Metastases enhance brightly with contrast, so this is the natural sequence for detecting them.
  • Model: a small 2D U-Net that we build ourselves, in PyTorch.

Why build a 2D model from scratch first?​

The published UCSF-BMSR benchmarks use nnU-Net — a powerful, self-configuring segmentation framework. It is the destination, not the starting point, and for one practical reason:

Federated learning requires you to extract your model's parameters, average them across sites, and load them back. That is three lines each when you built the model and can call named_parameters() on it (the technique from the Concepts section). It is considerably harder inside a framework like nnU-Net, which wraps the network, its preprocessing, and its training loop behind its own machinery — and which you drive from the command line rather than own as a model. Getting at "just the weights" to federate them each round is not the easy path.

Note this is only about where the weights live — the dataset itself needs no nnU-Net. It is ordinary MRI (NIfTI volumes + masks) that our own loader reads directly. So the plan is:

  1. Learn the whole federated loop on a model you fully control — a hand-written 2D U-Net.
  2. Get a real federation running end to end on the UCSF-BMSR T1Post data.
  3. Then graduate to nnU-Net if you need its accuracy, knowing exactly what "extract, average, reload" must achieve.

Working in 2D (one axial slice at a time) keeps the model small, fast to train on a laptop or a single GPU, and easy to reason about — ideal for learning. You can move to 3D later.

What you'll need​

  • A BranchKey account, and a leaf created for your site — see Getting Started. Keep its credentials JSON handy.
  • Python 3.9+ with: branchkey, torch, numpy, nibabel (for reading MRI volumes). 
    pip install branchkey torch numpy nibabel
  • Access to the UCSF-BMSR dataset (request/download via the dataset page; follow its licence and usage terms).
This is a guide, not a drop-in repo

The code here is written to be read and understood, then adapted to your exact file layout. Paths, image dimensions, and hyperparameters are illustrative — treat the dataset page and the benchmarks repo as authoritative for the data itself.

The walkthrough​

  1. Preparing the data — from 3D MRI volumes to 2D training slices.
  2. Building and training the model — a 2D U-Net and a local training loop, before any federation.
  3. Federating it with BranchKey — wire in the client so multiple sites train together without sharing data.

Before you start, it helps to have read the Concepts section — especially Extracting & reloading weights and Federated averaging, by hand. This example puts those ideas to work on real data.