On-Device Fine-Tuning for Privacy-Preserving Mammography Classification

This is the repository for my final project for my BSc in Computer Science at University of London.

The aim of this project is demonstrate that we can build effective Breast Cancer Mammography classification models on Apple Silicon hardware, to support offline fine-tuning on hospital specific datasets.

The project comprises of two major parts:

1. Model Training Scripts - A series of scripts for training models and demonstrating the problem of domain shift.
2. Web-based Interface - A web-based interface to support easy and accessible fine-tuning.

This repository uses Git LFS to store large files like model weights.

Install Git LFS before cloning:

# macOS
brew install git-lfs

# Ubuntu/Debian
sudo apt install git-lfs

# Then initialise
git lfs install

If you've already cloned without LFS, pull the actual files with:

git lfs pull

Project utilises uv as it provides one of the best tools for managing Python dependenices.

Uv can be installed following these instructions.

Then, the project dependencies can be installed with the sync command, as follows:

uv sync

To run the web-based interface, it can be run as follows:

uv run streamlit run src/app.py

The web-interface is comprised of 3 sections, separated by tabs:

1. The Project Overview tab, where users select the model and configure training/test data folders.
2. The Inference tab, which provides users with the ability to classify Mammography images into benign or malignant, including batch evaluation on test datasets.
3. The Fine-Tune tab, which can be used to fine-tune on clinical data.

See docs/training_log.md for a complete record of experiments with metrics and W&B links.

The model training is broken down into a few different categories:

• Downloading datasets.
• Dataset EDA via Juypter notebooks.
• Training and inference scripts.

I utilised 3 separate Mammography datasets to demonstrate base model training, and then to highlight the problem of domain shift, and show how fine-tuning can recover performance. The datasets are:

• CBIS-DDSM.
• InBreast.
• VinDr-Mammo.

I choose to use the official CBIS-DDSM dataset.

See 01-cbis-ddsm for details of how it was downloaded and preprocessed.

Download from https://drive.google.com/file/d/19n-p9p9C0eCQA1ybm6wkMo-bbeccT_62/view?usp=sharing

Run Jupyter Lab as follows:

uv run jupyter lab

EDA notebook for the CBIS-DDSM dataset.

EDA notebook for the INBreast dataset.

Preprocessing pipeline demonstrations.

Patch extraction experiments.

EDA notebook for the VinDr Mammogram dataset.

The preparation scripts extract ROI crops, resize images to a fixed resolution (256×256), and split the data at patient level to avoid data leakage.

Run the CBIS-DDSM preparation script with default settings (70/10/20 train/val/test split):

uv run prepare-cbis

Fine-tune EfficientNet-B0 on the CBIS-DDSM dataset:

uv run train

Create small VinDr datasets for UI testing:

uv run python scripts/create_dataset.py --preset large
uv run python scripts/create_dataset.py --preset tiny
uv run python scripts/create_dataset.py --preset small

Output: datasets/prep/vindr-ui-${present}/

Run unit tests:

uv run python -m pytest tests/

Run the accessibility audit (requires the app's model weights):

uv run pytest tests/test_accessibility.py -v --no-cov

This spins up the Streamlit app and checks each tab for WCAG violations using axe-core.

Note: there are currently many skipped violations that come from Streamlit's component library. See issue.

This project uses ruff for linting and code formatting.

Commands are as follows:

# Check for lint errors
uv run ruff check .

# Auto-fix fixable issues
uv run ruff check --fix .

# Format code
uv run ruff format .

This project uses pyright for static type checking.

Run as follows:

uv run pyright src/