diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a3ed70e..de11514 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -26,6 +26,24 @@ jobs: - name: Test run: uv run pytest + wheel: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v10.2.0 + - name: Build and install the wheel outside the repository + run: | + uv build --wheel -o /tmp/dist + uv venv /tmp/venv + VIRTUAL_ENV=/tmp/venv uv pip install /tmp/dist/*.whl --extra-index-url https://download.pytorch.org/whl/cpu --index-strategy unsafe-best-match + - name: Run the commands from another folder + working-directory: /tmp + run: | + mkdir -p data/P0001 && echo '{}' > data/P0001/config.json + /tmp/venv/bin/deepaneseg-train --help + /tmp/venv/bin/deepaneseg-preprocess data_dir=data + test -f data/0Work/split_pats.json + docker: runs-on: ubuntu-latest steps: @@ -35,4 +53,4 @@ jobs: - name: Smoke test run: | docker run --rm deepaneseg:cpu python -c "import torch, deepaneseg.utils; print(torch.__version__)" - docker run --rm deepaneseg:cpu python scripts/train.py --help + docker run --rm deepaneseg:cpu deepaneseg-train --help diff --git a/Dockerfile b/Dockerfile index 5a05243..427424a 100644 --- a/Dockerfile +++ b/Dockerfile @@ -21,4 +21,4 @@ RUN uv sync --locked --no-install-project --no-dev --extra ${TORCH} && rm -rf /r COPY . . RUN uv sync --locked --no-dev --extra ${TORCH} && rm -rf /root/.cache/uv -CMD ["python", "scripts/train.py", "--help"] +CMD ["deepaneseg-train", "--help"] diff --git a/README.md b/README.md index 2e6a625..e30b9b5 100644 --- a/README.md +++ b/README.md @@ -1,138 +1,211 @@ > **Note:** This repository is a copy of the original code hosted at [gitlab.inria.fr/yassis/DeepAneSeg](https://gitlab.inria.fr/yassis/DeepAneSeg), developed during my PhD at Inria/LORIA. -# An Efficient Data Strategy for the Detection of Brain Aneurysms from MRA with Deep Learning -This repository contains the official software code related to the paper ["An Efficient Data Strategy for the Detection of Brain Aneurysms from MRA with Deep Learning"](https://link.springer.com/chapter/10.1007/978-3-030-88210-5_22). The paper was published in MICCAI Workshop on Data Augmentation, Labelling, and Imperfections 2021. - -# Abstract -The detection of intracranial aneurysms from Magnetic Resonance Angiography images is a problem of rapidly growing clinical importance. In the last 3 years, the raise of deep convolutional neural networks has instigated a streak of methods that have shown promising performance. The major issue to address is the very severe class imbalance. Previous authors have focused their efforts on the network architecture and loss function. This paper tackles the data. A rough but fast annotation is considered: each aneurysm is approximated by a sphere defined by two points. Second, a small patch approach is taken so as to increase the number of samples. Third, samples are generated by a combination of data selection (negative patches are centered half on blood vessels and half on parenchyma) and data synthesis (patches containing an aneurysm are duplicated and deformed by a 3D spline transform). This strategy is applied to train a 3D U-net model, with a binary cross entropy loss, on a data set of 111 patients (155 aneurysms, mean size 3.86 mm ± 2.39 mm, min 1.23 mm, max 19.63 mm). A 5-fold cross-validation evaluation provides state of the art results (sensitivity 0.72, false positive count 0.14, as per ADAM challenge criteria). The study also reports a comparison with the focal loss, and Cohen’s Kappa coefficient is shown to be a better metric than Dice for this highly unbalanced detection problem. - +# DeepAneSeg + +[![CI](https://github.com/youssefassis/DeepAneSeg/actions/workflows/ci.yml/badge.svg)](https://github.com/youssefassis/DeepAneSeg/actions/workflows/ci.yml) + +Detection of intracranial aneurysms in Time-of-Flight MRA with a 3D U-Net trained on small patches. This is the +code of the paper [An Efficient Data Strategy for the Detection of Brain Aneurysms from MRA with Deep +Learning](https://link.springer.com/chapter/10.1007/978-3-030-88210-5_22) (MICCAI 2021 DALI workshop), packaged as a +set of commands that go from raw scans to detections and evaluation metrics. + +## Method +The main difficulty is the extreme class imbalance: aneurysms occupy a tiny fraction of a scan. Rather than a new +architecture or loss, the paper proposes a data strategy: +- **Fast annotation:** each aneurysm is approximated by a sphere, defined by two points on a diameter. +- **Small patches**, to increase the number of training samples. +- **Sample selection and synthesis:** negative patches are centered half on blood vessels and half in the parenchyma, + and patches containing an aneurysm are duplicated and deformed by a random 3D spline transform. + +Trained with a binary cross-entropy loss on 111 patients (155 aneurysms), a 5-fold cross-validation reached a +sensitivity of 0.72 with 0.14 false positives per patient (ADAM challenge criteria). + +
+Abstract + +The detection of intracranial aneurysms from Magnetic Resonance Angiography images is a problem of rapidly growing +clinical importance. In the last 3 years, the raise of deep convolutional neural networks has instigated a streak of +methods that have shown promising performance. The major issue to address is the very severe class imbalance. +Previous authors have focused their efforts on the network architecture and loss function. This paper tackles the +data. A rough but fast annotation is considered: each aneurysm is approximated by a sphere defined by two points. +Second, a small patch approach is taken so as to increase the number of samples. Third, samples are generated by a +combination of data selection (negative patches are centered half on blood vessels and half on parenchyma) and data +synthesis (patches containing an aneurysm are duplicated and deformed by a 3D spline transform). This strategy is +applied to train a 3D U-net model, with a binary cross entropy loss, on a data set of 111 patients (155 aneurysms, +mean size 3.86 mm ± 2.39 mm, min 1.23 mm, max 19.63 mm). A 5-fold cross-validation evaluation provides state of the +art results (sensitivity 0.72, false positive count 0.14, as per ADAM challenge criteria). The study also reports a +comparison with the focal loss, and Cohen's Kappa coefficient is shown to be a better metric than Dice for this +highly unbalanced detection problem. +
+ +## Quick start +``` +git clone https://github.com/youssefassis/DeepAneSeg && cd DeepAneSeg +uv sync --extra cu126 # or --extra cpu, see Installation +uv run deepaneseg-prepare source_dir=ADAM data_dir=Data 'labels=[1]' # or lay out your data as in Data +uv run deepaneseg-remove-skull data_dir=Data +uv run deepaneseg-extract-points data_dir=Data +uv run deepaneseg-preprocess data_dir=Data +uv run deepaneseg-train data_dir=Data name=exp1 +uv run deepaneseg-predict train_dir=Data/0Work/exp1 +uv run deepaneseg-evaluate train_dir=Data/0Work/exp1 +``` +Each step is described in [Usage](#usage), and its results in [Outputs](#outputs). + +## Installation +The project uses [uv](https://docs.astral.sh/uv/) and Python 3.10+. From the repository folder, pick the PyTorch +build: +``` +uv sync --extra cu126 # NVIDIA GPU (CUDA 12.6) +uv sync --extra cpu # CPU only, also the build used on macOS +``` +This creates `.venv` with the exact versions of `uv.lock`. Commands then run with `uv run`, or directly once `.venv` +is activated. -# Background -The major issue to address in medical image segmentation is the very severe class imbalance. -To adress this problem, the most previous works were mainly focused on the network architecture and loss function. This paper tackles the data. -Our data stratery is composed of the following three stages: -* A rough but fast annotation is considered: each aneurysm is approximated by a sphere defined by two points. -* A small patch approach is taken so as to increase the number of samples. -* Samples are generated by a combination of **data selection** (negative patches are centered half on blood vessels and half on parenchyma) and **data synthesis** (patches containing an aneurysm are duplicated anddeformed by a 3D spline transform). +Without uv, `pip install git+https://github.com/youssefassis/DeepAneSeg` installs the package and its commands +(with PyTorch's default build from PyPI). -# Project description -The code is organised as a Python package, `deepaneseg`, plus command-line scripts: +### Docker +The image installs the same locked versions, with the GPU (default) or CPU PyTorch build: ``` -deepaneseg/ - data/ patient data I/O, train/valid/test splits, augmentation transforms - volume/ patch extraction, sphere burning, point selection, skull stripping - models/ 3D U-Net - training/ dataset, losses, metrics, trainer - inference/ patch-wise prediction and ADAM-style evaluation - experimental/ post-paper models (vessel-coupled U-Nets) and their trainer -configs/ Hydra configuration (see Configuration) -scripts/ pipeline steps (see Usage) - experimental/ analysis scripts from the PhD experiments +docker build -t deepaneseg --build-arg GIT_COMMIT=$(git rev-parse HEAD) . +docker build -t deepaneseg:cpu --build-arg TORCH=cpu . ``` +Mount the data folder and run any command; outputs are written to the mounted folder: +``` +docker run --rm --gpus all --shm-size=8g -v /path/to/Data:/data deepaneseg \ + deepaneseg-train data_dir=/data name=exp1 +``` +GPU runs need the [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/) on +the host. `--shm-size` gives the data-loading workers enough shared memory, and `--user $(id -u):$(id -g)` creates +the output files with your own user. -The data is expected in the following layout: +## Data +### Layout +Each patient is a folder named `P` followed by four digits: ``` -Data_dir/ +Data/ P0001/ - volume.nii.gz 3D MRA image - F.csv coordinates (x, y, z) of the two points defining each aneurysm + volume.nii.gz TOF-MRA scan + F.csv aneurysms: two points (x, y, z columns, in mm) on a diameter of each aneurysm config.json {"init volume": "volume.nii.gz", "pts aneurysm": "F.csv"} - noskull.nii.gz skull-stripped volume (created by remove_skull.py) - points.csv negative patch centers (created by extract_points.py) + noskull.nii.gz skull-stripped scan (written by deepaneseg-remove-skull) + points.csv negative patch centers (written by deepaneseg-extract-points) P0002/ ... - 0Work/ created by the scripts - split_pats.json train/valid/test split - / checkpoints, TensorBoard logs, predictions and resolved config (.hydra/) of one training - logs/ logs of the preprocessing scripts + 0Work/ splits, trainings and logs (see Outputs) ``` +Patients without aneurysm simply have no `pts aneurysm` entry. Coordinates are in the scanner (RAS) space of the +NIfTI file, so the points can be placed with any viewer, e.g. as markups in 3D Slicer. -# Configuration -Settings are [Hydra](https://hydra.cc) configs in `configs/`: one file per script (`train.yaml`, `predict.yaml`, -`preprocess.yaml`, `remove_skull.yaml`, `extract_points.yaml`) composed from groups (`data`, `model`, `optimizer`, -`scheduler`, `augmentation`). Any value can be overridden on the command line: +### Converting a dataset +`deepaneseg-prepare` converts a dataset of scans with aneurysm label masks to this layout. Its defaults match the +ADAM challenge (`/orig/TOF.nii.gz` and `/aneurysms.nii.gz`, where 2 marks treated aneurysms); other +datasets only need their own `scan` and `mask` patterns. Each aneurysm becomes a two-point annotation, and +`Data/cases.csv` maps the patients to the original cases. ``` -uv run python scripts/train.py data_dir=Data_dir name=my_training batch_size=8 optimizer.lr=3e-4 model.f_maps=32 -uv run python scripts/train.py --cfg job --resolve data_dir=Data_dir name=my_training # print the config +uv run deepaneseg-prepare source_dir=ADAM data_dir=Data 'labels=[1]' +uv run deepaneseg-prepare source_dir=my_data data_dir=Data 'scan=scans/{case}.nii.gz' 'mask=masks/{case}.nii.gz' ``` -Runs are reproducible: `seed` (0 by default) fixes the patient split and the training, and each training writes -`run_info.json` with the git commit and library versions next to its checkpoints. For identical GPU runs, also set -`deterministic=true` (slower). - -To add an alternative (for example another scheduler), add a file to the group, such as -`configs/scheduler/step.yaml` with a `_target_` class, and select it with `scheduler=step`. -# Usage -Run every command from the repository root. - -**1. Data preparation:** skull-strip the volumes, then select the negative patch centers. +## Usage +**1. Preprocessing:** skull-strip the scans, then select the negative patch centers. ``` -uv run python scripts/remove_skull.py data_dir=Data_dir -uv run python scripts/extract_points.py data_dir=Data_dir +uv run deepaneseg-remove-skull data_dir=Data +uv run deepaneseg-extract-points data_dir=Data ``` -**2. Split:** assign the patients to training, validation and testing. +**2. Split:** assign the patients to training, validation and testing (70/20/10 by default). ``` -uv run python scripts/preprocess_data.py data_dir=Data_dir # split=[0.7,0.2,0.1] overwrite=true +uv run deepaneseg-preprocess data_dir=Data # 'split=[0.8,0.1,0.1]' overwrite=true ``` -**3. Training:** outputs go to `Data_dir/0Work/my_training`; running the same command again resumes training. -The training and validation volumes are kept in memory as float32, about 4 bytes per voxel (150 MiB for a -512×512×150 scan). +**3. Training:** outputs go to `Data/0Work/exp1`; running the same command again resumes the training. ``` -uv run python scripts/train.py data_dir=Data_dir name=my_training +uv run deepaneseg-train data_dir=Data name=exp1 ``` -**4. Testing:** predictions of the test patients are written to `Data_dir/0Work/my_training/prediction/test`: a -probability map per patient, and its detections (center and radius in mm, size, probability) as CSV and as `.fcsv` -markups for 3D Slicer. +**4. Prediction:** writes a probability map and the detections of each test patient. +``` +uv run deepaneseg-predict train_dir=Data/0Work/exp1 +``` +Any other scan, or folder of scans, can be predicted the same way. Scans are skull-stripped first, like the +training data. ``` -uv run python scripts/predict.py train_dir=Data_dir/0Work/my_training +uv run deepaneseg-predict train_dir=Data/0Work/exp1 input=scan.nii.gz output_dir=results ``` -Any other scan, or folder of scans, can be predicted the same way (skull-stripped first, like the training data): + +**5. Evaluation:** detection metrics of the test predictions, with the ADAM challenge criteria (a detection counts +when its center lies within an aneurysm): sensitivity and false positives per patient. ``` -uv run python scripts/predict.py train_dir=Data_dir/0Work/my_training input=scan.nii.gz output_dir=results +uv run deepaneseg-evaluate train_dir=Data/0Work/exp1 # threshold=0.5 min_size=null ``` -**5. Evaluation:** detection metrics of the test predictions (ADAM challenge criteria: sensitivity and false -positives per patient), written to `Data_dir/0Work/my_training/evaluation`. +## Configuration +Settings are [Hydra](https://hydra.cc) configs in `deepaneseg/configs/`: one file per command, composed from groups +(`data`, `model`, `optimizer`, `scheduler`, `augmentation`). Any value can be overridden on the command line: ``` -uv run python scripts/evaluate.py train_dir=Data_dir/0Work/my_training # threshold=0.5 min_size=null +uv run deepaneseg-train data_dir=Data name=exp2 batch_size=8 optimizer.lr=3e-4 model.f_maps=32 +uv run deepaneseg-train --cfg job --resolve data_dir=Data name=exp2 # print the resulting config ``` +To add an alternative, e.g. another scheduler, add a file with a `_target_` class to the group, such as +`deepaneseg/configs/scheduler/step.yaml`, and select it with `scheduler=step`. -# Differences from the published code -The Kappa metric/loss follows Cohen's definition and the focal loss is computed per voxel (Lin et al.), so -trainings using them can give numbers that differ slightly from the published ones. +Runs are reproducible: `seed` (0 by default) fixes the split and the training, and each training records its git +commit and library versions in `run_info.json`. For identical GPU runs, also set `deterministic=true` (slower). -# Installation -The project uses [uv](https://docs.astral.sh/uv/) (Python 3.10+). From the repository root, pick the PyTorch build: +## Outputs +Everything is written under `Data/0Work`: ``` -uv sync --extra cu126 # NVIDIA GPU (CUDA 12.6) -uv sync --extra cpu # CPU only (also the build used on macOS) +split_pats.json train/valid/test split +logs/ logs of the preprocessing commands +exp1/ one training + .hydra/config.yaml resolved configuration, read back by deepaneseg-predict and deepaneseg-evaluate + run_info.json git commit and library versions + last_checkpoint.pytorch model after the last epoch, used for prediction by default + best_checkpoint.pytorch model with the best validation metric + logs/ TensorBoard logs: uv run tensorboard --logdir Data/0Work + train.log + prediction/test/ for each test patient (prediction/scans/ for input=...): + .nii.gz probability map, on the patient's voxel grid + _detections.csv one row per detection: x, y, z and radius (mm), voxels, probability + _detections.fcsv the same detections as 3D Slicer markups + evaluation/ per_patient.csv and summary.json ``` -This creates `.venv` with the exact versions of `uv.lock`, including pytest and ruff. Commands then run with -`uv run`, e.g. `uv run pytest`. Without uv, `pip install -e .` also works (PyTorch's default PyPI build). -# Docker -Build the image with the GPU (default) or CPU PyTorch build: +## Requirements +- **Memory:** training keeps the training and validation scans in memory, about 4 bytes per voxel: 150 MiB for a + 512×512×150 scan, around 16 GiB for 111 patients. +- **GPU:** the U-Net has 19.1 M parameters (73 MiB). Its memory use comes mostly from the batch of 20 patches of + 48³ voxels; if the GPU runs out of memory, lower `batch_size` and `validation_batch_size`. +- **Disk:** predictions are float32 probability maps, one per scan. + +## Development ``` -docker build -t deepaneseg --build-arg GIT_COMMIT=$(git rev-parse HEAD) . -docker build -t deepaneseg:cpu --build-arg TORCH=cpu . +uv run pytest # tests +uv run ruff check . # lint +uv run ruff format . # formatting ``` -Mount the data folder and run any script; outputs are written under `Data_dir/0Work` on the host: +The code is organised as follows: ``` -docker run --rm --gpus all --shm-size=8g -v /path/to/Data_dir:/data deepaneseg \ - python scripts/train.py data_dir=/data name=my_training -docker run --rm --gpus all --shm-size=8g -v /path/to/Data_dir:/data deepaneseg \ - python scripts/predict.py train_dir=/data/0Work/my_training +deepaneseg/ + cli/ commands + configs/ Hydra configuration + data/ patient data I/O, dataset conversion, splits, augmentation transforms + volume/ patch extraction, sphere burning, point selection, skull stripping + models/ 3D U-Net + training/ dataset, losses, metrics, trainer + inference/ patch-wise prediction, detections and evaluation + experimental/ post-paper models (vessel-coupled U-Nets) and their trainer +scripts/experimental/ analysis scripts from the PhD experiments ``` -GPU runs need the [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/) on the -host. `--shm-size` gives the data-loading workers enough shared memory, and `--user $(id -u):$(id -g)` creates the -output files with your own user. -# Citation -If you find this repository useful in your research, please consider citing the paper below. The same reference -is available from "Cite this repository" on GitHub (`CITATION.cff`). +## Differences from the published code +The Kappa metric/loss follows Cohen's definition and the focal loss is computed per voxel (Lin et al.), so trainings +using them can give numbers that differ slightly from the published ones. + +## Citation +If you find this work useful, please cite the paper (also available from "Cite this repository" on GitHub): ``` @inproceedings{assis2021efficient, title={An efficient data strategy for the detection of brain aneurysms from MRA with deep learning}, @@ -144,5 +217,5 @@ is available from "Cite this repository" on GitHub (`CITATION.cff`). } ``` -# Acknowledgements -This work was funded by the Grand-Est Region, and the University Hospital (CHRU) of Nancy, France +## Acknowledgements +This work was funded by the Grand-Est Region and the University Hospital (CHRU) of Nancy, France. diff --git a/deepaneseg/cli/__init__.py b/deepaneseg/cli/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/scripts/evaluate.py b/deepaneseg/cli/evaluate.py similarity index 94% rename from scripts/evaluate.py rename to deepaneseg/cli/evaluate.py index 1168e42..d845995 100644 --- a/scripts/evaluate.py +++ b/deepaneseg/cli/evaluate.py @@ -1,7 +1,6 @@ -#!/usr/bin/env python3 -"""Score a training's predictions on its test patients; configuration in configs/evaluate.yaml. +"""Score a training's predictions on its test patients; configuration in deepaneseg/configs/evaluate.yaml. -Example: python scripts/evaluate.py train_dir=/data/0Work/exp1 +Example: deepaneseg-evaluate train_dir=/data/0Work/exp1 A detection (connected component above the threshold) is a true positive when its center lies within an aneurysm's radius (ADAM challenge criteria). Writes per_patient.csv and summary.json to /evaluation. """ diff --git a/scripts/extract_points.py b/deepaneseg/cli/extract_points.py similarity index 81% rename from scripts/extract_points.py rename to deepaneseg/cli/extract_points.py index 571e5ba..b08eef4 100644 --- a/scripts/extract_points.py +++ b/deepaneseg/cli/extract_points.py @@ -1,7 +1,6 @@ -#!/usr/bin/env python3 """Select negative patch centers (on vessels and in the parenchyma) for every patient. -Configuration in configs/extract_points.yaml. Example: python scripts/extract_points.py data_dir=/data +Configuration in deepaneseg/configs/extract_points.yaml. Example: deepaneseg-extract-points data_dir=/data Writes (CSV) and its 3D Slicer .fcsv counterpart in each patient folder. Requires the 'noskull volume' produced by remove_skull.py. """ diff --git a/scripts/predict.py b/deepaneseg/cli/predict.py similarity index 91% rename from scripts/predict.py rename to deepaneseg/cli/predict.py index 30d952f..ed5bf67 100755 --- a/scripts/predict.py +++ b/deepaneseg/cli/predict.py @@ -1,9 +1,8 @@ -#!/usr/bin/env python3 -"""Predict with a trained model; configuration in configs/predict.yaml. +"""Predict with a trained model; configuration in deepaneseg/configs/predict.yaml. Examples: - python scripts/predict.py train_dir=/data/0Work/exp1 # test patients of the split - python scripts/predict.py train_dir=/data/0Work/exp1 input=scan.nii.gz # any scan, or a folder of scans + deepaneseg-predict train_dir=/data/0Work/exp1 # test patients of the split + deepaneseg-predict train_dir=/data/0Work/exp1 input=scan.nii.gz # any scan, or a folder of scans Data and model settings come from the training's resolved configuration (/.hydra/config.yaml). For each case, writes the probability map .nii.gz and the detections _detections.csv (and .fcsv, for 3D Slicer). diff --git a/deepaneseg/cli/prepare.py b/deepaneseg/cli/prepare.py new file mode 100644 index 0000000..60a6484 --- /dev/null +++ b/deepaneseg/cli/prepare.py @@ -0,0 +1,57 @@ +"""Convert a dataset of scans and aneurysm masks to the pipeline's layout; configuration in deepaneseg/configs/prepare.yaml. + +Example (ADAM challenge data): deepaneseg-prepare source_dir=/data/ADAM data_dir=/data/converted labels=[1] +Writes one P???? folder per case (volume, F.csv with the aneurysms, config.json) and cases.csv, which maps the +patients to the original case names. +""" + +import json +import os +import shutil +import sys + +import hydra +import pandas as pd + +import deepaneseg.data.io as dio +from deepaneseg.data.prepare import aneurysm_points, find_cases +from deepaneseg.utils import get_logger + + +@hydra.main(version_base="1.3", config_path="../configs", config_name="prepare") +def main(cfg): + logger = get_logger("Prepare") + cases = find_cases(cfg.source_dir, cfg.scan) + if not cases: + sys.exit(f"No scan matches '{cfg.scan}' in {cfg.source_dir}") + masks = {case: os.path.join(cfg.source_dir, cfg.mask.replace("{case}", case)) for case in cases} if cfg.mask else {} + missing = sorted(case for case, path in masks.items() if not os.path.isfile(path)) + if missing: + raise FileNotFoundError(f"No aneurysm mask '{cfg.mask}' for {', '.join(missing)}") + if dio.fetch_patient_dirs(cfg.data_dir) and not cfg.overwrite: + sys.exit(f"{cfg.data_dir} already contains patients; pass overwrite=true to replace them") + + rows = [] + for i, (case, scan) in enumerate(cases.items(), start=1): + patient = os.path.join(cfg.data_dir, f"P{i:04d}") + os.makedirs(patient, exist_ok=True) + volume = "volume.nii.gz" if scan.endswith(".gz") else "volume.nii" + shutil.copyfile(scan, os.path.join(patient, volume)) + config = {"init volume": volume} + points = [] + if case in masks: + mask, affine = dio.read_nii_from_file(masks[case]) + points = aneurysm_points(mask, affine, labels=cfg.labels) + if len(points): + pd.DataFrame(points, columns=["x", "y", "z"]).to_csv(os.path.join(patient, "F.csv"), index=False) + config["pts aneurysm"] = "F.csv" + with open(os.path.join(patient, "config.json"), "w") as f: + json.dump(config, f, indent=2) + rows.append({"patient": f"P{i:04d}", "case": case, "scan": scan, "aneurysms": len(points) // 2}) + logger.info(f"{case} -> P{i:04d}: {len(points) // 2} aneurysms") + + pd.DataFrame(rows).to_csv(os.path.join(cfg.data_dir, "cases.csv"), index=False) + + +if __name__ == "__main__": + main() diff --git a/scripts/preprocess_data.py b/deepaneseg/cli/preprocess.py similarity index 85% rename from scripts/preprocess_data.py rename to deepaneseg/cli/preprocess.py index 42d3347..d4c2a78 100644 --- a/scripts/preprocess_data.py +++ b/deepaneseg/cli/preprocess.py @@ -1,7 +1,6 @@ -#!/usr/bin/env python3 -"""Split the patients into training/validation/testing sets; configuration in configs/preprocess.yaml. +"""Split the patients into training/validation/testing sets; configuration in deepaneseg/configs/preprocess.yaml. -Example: python scripts/preprocess_data.py data_dir=/data split=[0.7,0.2,0.1] +Example: deepaneseg-preprocess data_dir=/data split=[0.7,0.2,0.1] Writes the split to data.split_file (/0Work/split_pats.json by default). """ diff --git a/scripts/remove_skull.py b/deepaneseg/cli/remove_skull.py similarity index 87% rename from scripts/remove_skull.py rename to deepaneseg/cli/remove_skull.py index 2a04a68..8753243 100644 --- a/scripts/remove_skull.py +++ b/deepaneseg/cli/remove_skull.py @@ -1,7 +1,6 @@ -#!/usr/bin/env python3 -"""Skull-strip every patient volume; configuration in configs/remove_skull.yaml. +"""Skull-strip every patient volume; configuration in deepaneseg/configs/remove_skull.yaml. -Example: python scripts/remove_skull.py data_dir=/data +Example: deepaneseg-remove-skull data_dir=/data The result is saved next to each volume and registered as 'noskull volume' in the patient's config.json. """ diff --git a/scripts/train.py b/deepaneseg/cli/train.py similarity index 94% rename from scripts/train.py rename to deepaneseg/cli/train.py index 7c08457..53c5cc1 100755 --- a/scripts/train.py +++ b/deepaneseg/cli/train.py @@ -1,7 +1,6 @@ -#!/usr/bin/env python3 -"""Train a model; configuration in configs/train.yaml. +"""Train a model; configuration in deepaneseg/configs/train.yaml. -Example: python scripts/train.py data_dir=/data name=exp1 batch_size=8 +Example: deepaneseg-train data_dir=/data name=exp1 batch_size=8 Checkpoints, TensorBoard logs and the resolved configuration (.hydra/) are written to /0Work/. Running the same command again resumes from the last checkpoint. """ diff --git a/deepaneseg/configs/__init__.py b/deepaneseg/configs/__init__.py new file mode 100644 index 0000000..bf4c8be --- /dev/null +++ b/deepaneseg/configs/__init__.py @@ -0,0 +1 @@ +"""Hydra configuration of the deepaneseg commands (a package so that installed commands find it).""" diff --git a/configs/augmentation/default.yaml b/deepaneseg/configs/augmentation/default.yaml similarity index 100% rename from configs/augmentation/default.yaml rename to deepaneseg/configs/augmentation/default.yaml diff --git a/configs/data/default.yaml b/deepaneseg/configs/data/default.yaml similarity index 77% rename from configs/data/default.yaml rename to deepaneseg/configs/data/default.yaml index 6dfcec3..5a33f7b 100644 --- a/configs/data/default.yaml +++ b/deepaneseg/configs/data/default.yaml @@ -1,9 +1,9 @@ # Patient data and patch geometry data_dir: ${data_dir} # directory containing the patient folders (P????) work_dir: ${.data_dir}/0Work # split file and trainings are written here -split_file: ${.work_dir}/split_pats.json # written by preprocess_data.py +split_file: ${.work_dir}/split_pats.json # written by deepaneseg-preprocess volume: init volume # key, in each patient's config.json, of the volume to train on normalize: Linear # Linear (rescale to [0, 1]), Normal (zero mean, unit variance) or null -negative_patch_centers: points.csv # written by extract_points.py in each patient folder +negative_patch_centers: points.csv # written by deepaneseg-extract-points in each patient folder patch_shape: [48, 48, 48] # patch dimension, in voxels patch_size: [38, 38, 38] # patch size, in mm (resolution = patch_size / patch_shape) diff --git a/configs/evaluate.yaml b/deepaneseg/configs/evaluate.yaml similarity index 87% rename from configs/evaluate.yaml rename to deepaneseg/configs/evaluate.yaml index 42d6a45..b5b32aa 100644 --- a/configs/evaluate.yaml +++ b/deepaneseg/configs/evaluate.yaml @@ -1,6 +1,6 @@ # Detection metrics (ADAM challenge criteria) of a training's predictions on its test patients -train_dir: ??? # training directory written by train.py -predictions: ${train_dir}/prediction/test # .nii.gz probability maps written by predict.py +train_dir: ??? # training directory written by deepaneseg-train +predictions: ${train_dir}/prediction/test # .nii.gz probability maps written by deepaneseg-predict threshold: 0.5 # probability above which a voxel belongs to a detection min_size: null # detections smaller than this, in voxels, are ignored; null keeps all diff --git a/configs/extract_points.yaml b/deepaneseg/configs/extract_points.yaml similarity index 100% rename from configs/extract_points.yaml rename to deepaneseg/configs/extract_points.yaml diff --git a/configs/model/unet3d.yaml b/deepaneseg/configs/model/unet3d.yaml similarity index 100% rename from configs/model/unet3d.yaml rename to deepaneseg/configs/model/unet3d.yaml diff --git a/configs/optimizer/adam.yaml b/deepaneseg/configs/optimizer/adam.yaml similarity index 100% rename from configs/optimizer/adam.yaml rename to deepaneseg/configs/optimizer/adam.yaml diff --git a/configs/predict.yaml b/deepaneseg/configs/predict.yaml similarity index 87% rename from configs/predict.yaml rename to deepaneseg/configs/predict.yaml index 422bc3a..26f0156 100644 --- a/configs/predict.yaml +++ b/deepaneseg/configs/predict.yaml @@ -1,10 +1,10 @@ # Predicts with a trained model; data and model settings are read from the training's config -train_dir: ??? # training directory written by train.py +train_dir: ??? # training directory written by deepaneseg-train input: null # NIfTI scan or folder of scans; null: the test patients of the training's split output_dir: null # null: /prediction/test, or /prediction/scans for input checkpoint: last_checkpoint.pytorch volume: noskull volume # key, in each patient's config.json, of the volume to predict on -skull_strip: true # skull-strip input scans as remove_skull.py does, to match volume above +skull_strip: true # skull-strip input scans as deepaneseg-remove-skull does, to match volume above threshold: 0.5 # detections: connected components above this probability min_size: null # detections smaller than this, in voxels, are ignored; null keeps all margin: 8 # voxels discarded at each patch border when stitching predictions diff --git a/deepaneseg/configs/prepare.yaml b/deepaneseg/configs/prepare.yaml new file mode 100644 index 0000000..709a7ea --- /dev/null +++ b/deepaneseg/configs/prepare.yaml @@ -0,0 +1,11 @@ +# Converts a dataset of scans and aneurysm label masks (defaults: ADAM challenge layout) to the pipeline's layout +source_dir: ??? # dataset to convert +data_dir: ??? # output folder, one P???? folder per case +scan: "{case}/orig/TOF.nii.gz" # scan of each case, relative to source_dir; {case} matches one name +mask: "{case}/aneurysms.nii.gz" # aneurysm label mask of a case, relative to source_dir; null if there are none +labels: null # mask values that are aneurysms, e.g. [1] for ADAM (2: treated); null: any non-zero +overwrite: false # replace existing patients in data_dir + +hydra: + run: + dir: ${data_dir}/0Work/logs/${hydra.job.name}/${now:%Y-%m-%d_%H-%M-%S} diff --git a/configs/preprocess.yaml b/deepaneseg/configs/preprocess.yaml similarity index 100% rename from configs/preprocess.yaml rename to deepaneseg/configs/preprocess.yaml diff --git a/configs/remove_skull.yaml b/deepaneseg/configs/remove_skull.yaml similarity index 100% rename from configs/remove_skull.yaml rename to deepaneseg/configs/remove_skull.yaml diff --git a/configs/scheduler/reduce_lr_on_plateau.yaml b/deepaneseg/configs/scheduler/reduce_lr_on_plateau.yaml similarity index 100% rename from configs/scheduler/reduce_lr_on_plateau.yaml rename to deepaneseg/configs/scheduler/reduce_lr_on_plateau.yaml diff --git a/configs/train.yaml b/deepaneseg/configs/train.yaml similarity index 100% rename from configs/train.yaml rename to deepaneseg/configs/train.yaml diff --git a/deepaneseg/data/prepare.py b/deepaneseg/data/prepare.py new file mode 100644 index 0000000..965a8f3 --- /dev/null +++ b/deepaneseg/data/prepare.py @@ -0,0 +1,42 @@ +"""Conversion of a dataset of scans and aneurysm masks to the layout used by the pipeline.""" + +import glob +import os +import re + +import numpy as np + +from deepaneseg.volume.selection import connected_components_to_spheres + + +def find_cases(source_dir, scan_pattern): + """ + Returns {case name: scan path}, sorted by case, for the files matching scan_pattern in source_dir. + scan_pattern is a path relative to source_dir where {case} stands for one folder or file name part, + e.g. "{case}/orig/TOF.nii.gz" (ADAM challenge) or "scans/{case}.nii.gz". + """ + if scan_pattern.count("{case}") != 1: + raise ValueError(f"scan pattern must contain {{case}} exactly once, got '{scan_pattern}'") + before, after = scan_pattern.split("{case}") + regex = re.compile(re.escape(before) + r"([^/]+)" + re.escape(after) + "$") + cases = {} + for path in glob.glob(os.path.join(source_dir, scan_pattern.replace("{case}", "*"))): + match = regex.search(os.path.relpath(path, source_dir)) + if match: + cases[match.group(1)] = path + return dict(sorted(cases.items())) + + +def aneurysm_points(mask, affine, labels=None): + """ + Annotation points of the aneurysms of a label mask: each connected component of the selected labels (any + non-zero value if labels is None) is approximated by a sphere, described by the two ends of a diameter along x, + as in the F.csv files. Returns a 2Nx3 array of points in mm. + """ + selected = mask > 0 if labels is None else np.isin(mask, list(labels)) + spheres = connected_components_to_spheres(selected, affine) + offset = np.zeros_like(spheres[:, :3]) + offset[:, 0] = spheres[:, 3] + points = np.empty((2 * len(spheres), 3)) + points[0::2], points[1::2] = spheres[:, :3] - offset, spheres[:, :3] + offset + return points diff --git a/pyproject.toml b/pyproject.toml index c921302..797d9ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,15 @@ dependencies = [ "hydra-core>=1.3", ] +[project.scripts] +deepaneseg-prepare = "deepaneseg.cli.prepare:main" +deepaneseg-remove-skull = "deepaneseg.cli.remove_skull:main" +deepaneseg-extract-points = "deepaneseg.cli.extract_points:main" +deepaneseg-preprocess = "deepaneseg.cli.preprocess:main" +deepaneseg-train = "deepaneseg.cli.train:main" +deepaneseg-predict = "deepaneseg.cli.predict:main" +deepaneseg-evaluate = "deepaneseg.cli.evaluate:main" + [project.optional-dependencies] # PyTorch build, for uv: `uv sync --extra cpu` or `uv sync --extra cu126` (see [tool.uv.sources]) cpu = ["torch"] @@ -54,6 +63,9 @@ explicit = true [tool.setuptools.packages.find] include = ["deepaneseg*"] +[tool.setuptools.package-data] +deepaneseg = ["configs/*.yaml", "configs/*/*.yaml"] + [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..c4f2916 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,37 @@ +"""The installed commands and the configs they read.""" + +import shutil +import subprocess +from importlib.metadata import entry_points +from importlib.resources import files + +import pytest + +COMMANDS = { + "deepaneseg-prepare": "prepare", + "deepaneseg-remove-skull": "remove_skull", + "deepaneseg-extract-points": "extract_points", + "deepaneseg-preprocess": "preprocess", + "deepaneseg-train": "train", + "deepaneseg-predict": "predict", + "deepaneseg-evaluate": "evaluate", +} + + +def test_every_command_is_installed(): + installed = {ep.name: ep for ep in entry_points(group="console_scripts") if ep.name.startswith("deepaneseg-")} + + assert set(installed) == set(COMMANDS) + assert all(callable(ep.load()) for ep in installed.values()) + + +@pytest.mark.parametrize("config", sorted(set(COMMANDS.values()))) +def test_every_command_has_its_config_in_the_package(config): + assert (files("deepaneseg.configs") / f"{config}.yaml").is_file() + + +@pytest.mark.parametrize("command", ["deepaneseg-train", "deepaneseg-predict"]) +def test_commands_run_from_any_folder(command, tmp_path): + result = subprocess.run([shutil.which(command), "--help"], cwd=tmp_path, capture_output=True, text=True) + + assert result.returncode == 0 and "Configuration groups" in result.stdout diff --git a/tests/test_config.py b/tests/test_config.py index 4408ed4..8d44417 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -11,7 +11,7 @@ from deepaneseg.training.trainer import create_trainer from deepaneseg.utils import get_model -CONFIG_DIR = str(Path(__file__).parents[1] / "configs") +CONFIG_DIR = str(Path(__file__).parents[1] / "deepaneseg" / "configs") def compose_config(name, overrides): diff --git a/tests/test_evaluate.py b/tests/test_evaluate.py index 41733f5..85a930e 100644 --- a/tests/test_evaluate.py +++ b/tests/test_evaluate.py @@ -1,7 +1,6 @@ import json import subprocess import sys -from pathlib import Path import nibabel as ni import numpy as np @@ -12,8 +11,6 @@ from deepaneseg.data.io import read_aneurysm_spheres from deepaneseg.inference.evaluation import summarize_detections -SCRIPTS = Path(__file__).parents[1] / "scripts" - def make_patient(data_dir, name, aneurysm_points=None): patient = data_dir / name @@ -66,7 +63,7 @@ def test_evaluate_script_scores_the_test_predictions(tmp_path): ni.save(ni.Nifti1Image(np.full((20, 20, 20), 0.1), np.eye(4)), predictions / "P0002.nii.gz") result = subprocess.run( - [sys.executable, str(SCRIPTS / "evaluate.py"), f"train_dir={train_dir}"], capture_output=True, text=True + [sys.executable, "-m", "deepaneseg.cli.evaluate", f"train_dir={train_dir}"], capture_output=True, text=True ) assert result.returncode == 0, result.stderr[-2000:] @@ -87,7 +84,7 @@ def test_evaluate_script_reports_missing_predictions(tmp_path): ) result = subprocess.run( - [sys.executable, str(SCRIPTS / "evaluate.py"), f"train_dir={train_dir}"], capture_output=True, text=True + [sys.executable, "-m", "deepaneseg.cli.evaluate", f"train_dir={train_dir}"], capture_output=True, text=True ) assert result.returncode != 0 and "P0001" in result.stderr diff --git a/tests/test_predict_scans.py b/tests/test_predict_scans.py index a6e7e5c..59983aa 100644 --- a/tests/test_predict_scans.py +++ b/tests/test_predict_scans.py @@ -47,7 +47,7 @@ def test_save_detections_writes_csv_and_slicer_markups(tmp_path): @pytest.fixture def trained(tmp_path): """A tiny untrained model with its saved training configuration.""" - with initialize_config_dir(config_dir=str(ROOT / "configs"), version_base="1.3"): + with initialize_config_dir(config_dir=str(ROOT / "deepaneseg" / "configs"), version_base="1.3"): cfg = compose( "train", [ @@ -67,7 +67,7 @@ def trained(tmp_path): def predict(*args): return subprocess.run( - [sys.executable, str(ROOT / "scripts" / "predict.py"), *map(str, args)], capture_output=True, text=True + [sys.executable, "-m", "deepaneseg.cli.predict", *map(str, args)], capture_output=True, text=True ) diff --git a/tests/test_prepare.py b/tests/test_prepare.py new file mode 100644 index 0000000..fef6f0f --- /dev/null +++ b/tests/test_prepare.py @@ -0,0 +1,104 @@ +import subprocess +import sys + +import nibabel as ni +import numpy as np +import pandas as pd +import pytest + +from deepaneseg.data.io import fetch_patient_dirs, read_aneurysm_spheres, read_patient_data_base +from deepaneseg.data.prepare import aneurysm_points, find_cases +from deepaneseg.volume.selection import connected_components_to_spheres + +AFFINE = np.diag([0.4, 0.4, 0.6, 1.0]) +AFFINE[:3, 3] = [-50, -60, 10] + + +def write_adam_case(source, name, blobs=()): + """ADAM challenge layout: /orig/TOF.nii.gz and /aneurysms.nii.gz (1: untreated, 2: treated).""" + (source / name / "orig").mkdir(parents=True) + ni.save( + ni.Nifti1Image(np.random.default_rng(0).random((40, 40, 30)).astype(np.float32), AFFINE), + source / name / "orig" / "TOF.nii.gz", + ) + mask = np.zeros((40, 40, 30), dtype=np.uint8) + for label, (x, y, z) in blobs: + mask[x - 2 : x + 3, y - 2 : y + 3, z - 1 : z + 2] = label + ni.save(ni.Nifti1Image(mask, AFFINE), source / name / "aneurysms.nii.gz") + return mask + + +def test_find_cases_matches_the_pattern(tmp_path): + for name in ("10078F", "10001", "10042B"): + write_adam_case(tmp_path, name) + (tmp_path / "notes").mkdir() # no scan inside: ignored + + cases = find_cases(str(tmp_path), "{case}/orig/TOF.nii.gz") + + assert list(cases) == ["10001", "10042B", "10078F"] + assert cases["10001"] == str(tmp_path / "10001" / "orig" / "TOF.nii.gz") + + +def test_find_cases_needs_a_case_placeholder(tmp_path): + with pytest.raises(ValueError, match="{case}"): + find_cases(str(tmp_path), "*/orig/TOF.nii.gz") + + +def test_aneurysm_points_describe_each_selected_component(tmp_path): + mask = write_adam_case(tmp_path, "c", [(1, (10, 10, 10)), (1, (30, 30, 20)), (2, (10, 30, 15))]) + + points = aneurysm_points(mask, AFFINE, labels=[1]) + + assert points.shape == (4, 3) # two points per untreated aneurysm + spheres = np.hstack( + [(points[::2] + points[1::2]) / 2, np.linalg.norm(points[::2] - points[1::2], axis=1, keepdims=True) / 2] + ) + np.testing.assert_allclose(spheres, connected_components_to_spheres(mask == 1, AFFINE)) + assert len(aneurysm_points(mask, AFFINE, labels=None)) == 6 # any non-zero value + + +def run_prepare(source, data_dir, *args): + return subprocess.run( + [sys.executable, "-m", "deepaneseg.cli.prepare", f"source_dir={source}", f"data_dir={data_dir}", *args], + capture_output=True, + text=True, + ) + + +def test_prepare_converts_an_adam_style_dataset(tmp_path): + source, data_dir = tmp_path / "adam", tmp_path / "data" + mask = write_adam_case(source, "10002", [(1, (15, 15, 12)), (2, (30, 30, 20))]) + write_adam_case(source, "10001") # no aneurysm + + result = run_prepare(source, data_dir, "labels=[1]") + + assert result.returncode == 0, result.stderr[-2000:] + patients = fetch_patient_dirs(str(data_dir)) + assert [p.split("/")[-1] for p in patients] == ["P0001", "P0002"] + cases = pd.read_csv(data_dir / "cases.csv") + assert cases[["patient", "case", "aneurysms"]].values.tolist() == [["P0001", 10001, 0], ["P0002", 10002, 1]] + assert read_aneurysm_spheres(patients[0]).shape == (0, 4) + np.testing.assert_allclose(read_aneurysm_spheres(patients[1]), connected_components_to_spheres(mask == 1, AFFINE)) + patient = read_patient_data_base([patients[1]])[0] + np.testing.assert_allclose(patient["affine"], AFFINE) + assert patient["data"].shape == (40, 40, 30) + + +def test_prepare_refuses_to_overwrite_patients(tmp_path): + source, data_dir = tmp_path / "adam", tmp_path / "data" + write_adam_case(source, "10001") + run_prepare(source, data_dir) + + result = run_prepare(source, data_dir) + + assert result.returncode != 0 and "overwrite=true" in result.stderr + + +def test_prepare_reports_missing_masks(tmp_path): + source = tmp_path / "adam" + write_adam_case(source, "10001") + (source / "10001" / "aneurysms.nii.gz").unlink() + + result = run_prepare(source, tmp_path / "data") + + assert result.returncode != 0 and "10001" in result.stderr diff --git a/tests/test_reproducibility.py b/tests/test_reproducibility.py index 61e2ebc..06ba285 100644 --- a/tests/test_reproducibility.py +++ b/tests/test_reproducibility.py @@ -14,8 +14,6 @@ from deepaneseg.data.io import fetch_patient_dirs from deepaneseg.utils import run_info, seed_everything -SCRIPTS = Path(__file__).parents[1] / "scripts" - def test_split_depends_only_on_the_seed(): patients = [f"P{i:04d}" for i in range(20)] @@ -89,7 +87,7 @@ def train(name): "data.patch_size=[16,16,16]", "augmentation.positive.duplicates=2", ] - result = subprocess.run([sys.executable, str(SCRIPTS / "train.py"), *args], capture_output=True, text=True) + result = subprocess.run([sys.executable, "-m", "deepaneseg.cli.train", *args], capture_output=True, text=True) assert result.returncode == 0, result.stderr[-2000:] return torch.load(tmp_path / "0Work" / name / "last_checkpoint.pytorch")["model_state_dict"] diff --git a/tests/test_scripts.py b/tests/test_scripts.py index c701ed6..b243e26 100644 --- a/tests/test_scripts.py +++ b/tests/test_scripts.py @@ -1,13 +1,12 @@ import json import subprocess import sys -from pathlib import Path -SCRIPTS = Path(__file__).parents[1] / "scripts" - -def run(script, *args): - return subprocess.run([sys.executable, str(SCRIPTS / script), *map(str, args)], capture_output=True, text=True) +def run(command, *args): + return subprocess.run( + [sys.executable, "-m", f"deepaneseg.cli.{command}", *map(str, args)], capture_output=True, text=True + ) def make_patients(data_dir, n): @@ -19,8 +18,8 @@ def make_patients(data_dir, n): def test_preprocess_data_splits_patients_and_refuses_to_overwrite(tmp_path): make_patients(tmp_path, 10) - first = run("preprocess_data.py", f"data_dir={tmp_path}") - second = run("preprocess_data.py", f"data_dir={tmp_path}") + first = run("preprocess", f"data_dir={tmp_path}") + second = run("preprocess", f"data_dir={tmp_path}") assert first.returncode == 0, first.stderr split = json.loads((tmp_path / "0Work" / "split_pats.json").read_text()) @@ -29,6 +28,6 @@ def test_preprocess_data_splits_patients_and_refuses_to_overwrite(tmp_path): def test_train_requires_data_dir_and_name(): - result = run("train.py") + result = run("train") assert result.returncode != 0 and "Missing mandatory value" in result.stderr