Add Docker

Author: saundersresearch

.dockerignore

@@ -0,0 +1 @@
+.venv/

Dockerfile

@@ -0,0 +1,55 @@
+FROM python:3.12-slim-trixie
+
+RUN apt-get update && apt-get install -y git
+
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin
+
+WORKDIR /app
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,source=uv.lock,target=uv.lock \
+    --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
+    uv sync --locked --no-install-project
+
+ADD . /app
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv sync --locked
+
+# Need to add hsd as a package, but original GitHub does not have pyproject.toml
+# Use this as a workaround
+RUN mkdir -p /opt/hsd/src/hsd/
+RUN git clone https://github.com/AxelElaldi/fast-equivariant-deconv /opt/hsd/src/hsd/
+RUN touch /opt/hsd/src/hsd/__init__.py /opt/hsd/src/hsd/utils/__init__.py /opt/hsd/src/hsd/model/__init__.py
+RUN cat > /opt/hsd/pyproject.toml <<'PYPROJECT'
+[project]
+name = "hsd"
+version = "0.1.0"
+description = "https://github.com/AxelElaldi/fast-equivariant-deconv"
+readme = "README.md"
+dependencies = [
+    "h5py==3.11.0",
+    "healpy==1.16.6",
+    "joblib==1.4.2",
+    "matplotlib==3.9.0",
+    "nibabel==5.2.1",
+    "numpy<2",
+    "pandas==2.2.2",
+    "pygsp==0.5.1",
+    "pyyaml==6.0.1",
+    "scipy==1.13.0",
+    "tensorboard==2.16.2",
+    "torch==2.3.0",
+    "torchaudio==2.3.0",
+    "torchvision==0.18.0",
+]
+
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+PYPROJECT
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install /opt/hsd
+
+ENV PATH="/app/.venv/bin:${PATH}"

README.md

@@ -1,18 +1,59 @@
-# DeepFixel
-Deep learning-based identification of crossing fiber bundle elements
+# DeepFixel: Crossing white matter fiber identification through spherical convolutional neural networks
+[![arXiv](https://img.shields.io/badge/arXiv-2511.03893-b31b1b.svg)](https://arxiv.org/abs/2511.03893)
 
-## Training and testing the model
+DeepFixel is a deep learning method for identification of crossing fiber bundle elements from diffusion MRI.
+
+> Adam M. Saunders, Lucas W. Remedios, Elyssa M. McMaster, Jongyeon Yoon, Gaurav Rudravaram, Adam Sadriddinov, Praitayini Kanakaraj, Bennett A, Landman, and Adam W. Anderson. DeepFixel: Crossing white matter fiber identification through spherical convolutional neural networks. SPIE Medical Imaging: Clinical and Biomedical Imaging, 2026. [https://arxiv.org/abs/2511.03893](https://arxiv.org/abs/2511.03893)
+
+## Installation
 You can set up an environment using [`uv`](https://github.com/astral-sh/uv) by running the following command:
 ```bash
 uv sync
 ```
 
-To run the model, download the weights and testing dataset from the following link: [https://zenodo.org/records/16587458](https://zenodo.org/records/16587458).
+Alternatively, you can use Docker or Apptainer (see instructions below).
 
+## Usage
+To run the model, download the weights and testing dataset from the following link: [https://zenodo.org/records/16587458](https://zenodo.org/records/16587458). 
 - Unzip and copy the testing data to `./test_data`
 - Put the weights in `./models/pretrained`
 
-See `run_pretrained_deep_fixel.py` to test the pretrained model and `run_deep_fixel.py` to train and test a new model.
 
-## Using the model
-If you wish to apply the model to your own dataset, you can use `fissile.test_mesh_model()` as a basis for your code. You can also use `fissile.dataset.GeneratedMeshNIFTIDataset()` if your data is stored as spherical harmonic coefficients in a NIFTI file.
+To train the model:
+```bash
+python train_deep_fixel.py --config config/example_scnn.yaml
+```
+
+To test the model on the provided testing dataset:
+```bash
+python test_deep_fixel.py --config config/example_scnn.yaml
+```
+
+
+## Usage (Docker)
+To build the Docker image, clone the repository and run the following command in the root directory:
+```bash
+sudo docker build -t spherical_deep_fixel:v1.0.0 .
+```
+
+Then run the Docker container with the following command (note you will likely need to bind in local directories with `-v`):
+```bash
+sudo docker run --rm -it --gpus all -v $(pwd):$(pwd) $spherical_deep_fixel:v1.0.0 python train_deep_fixel.py --config /path/to/config/example_scnn.yaml
+sudo docker run --rm -it --gpus all -v $(pwd):$(pwd) $spherical_deep_fixel:v1.0.0 python test_deep_fixel.py --config /path/to/config/example_scnn.yaml
+```
+
+## Usage (Apptainer)
+A pre-built Apptainer image is available on Zenodo: 
+
+```bash
+apptainer run -C -B $(pwd):$(pwd) --nv https://zenodo.org/records/16587458/files/spherical_deep_fixel_v1.0.0.sif python /app/train_deep_fixel.py --config /path/to/config/example_scnn.yaml
+apptainer run -C -B $(pwd):$(pwd) --nv https://zenodo.org/records/16587458/files/spherical_deep_fixel_v1.0.0.sif python /app/test_deep_fixel.py --config /path/to/config/example_scnn.yaml
+```
+
+## Applying the model to your own data
+If you wish to apply the model to your own dataset, you can use `fissile.test_mesh_model()` as a basis for your code. You can also use `fissile.dataset.GeneratedMeshNIFTIDataset()` if your data is stored as spherical harmonic coefficients in a NIFTI file.
+
+## Citation
+If you use this code in your research, please cite the following paper:
+
+> Adam M. Saunders, Lucas W. Remedios, Elyssa M. McMaster, Jongyeon Yoon, Gaurav Rudravaram, Adam Sadriddinov, Praitayini Kanakaraj, Bennett A, Landman, and Adam W. Anderson. DeepFixel: Crossing white matter fiber identification through spherical convolutional neural networks. SPIE Medical Imaging: Clinical and Biomedical Imaging, 2026. [https://arxiv.org/abs/2511.03893](https://arxiv.org/abs/2511.03893)

config/example_mlp.yaml

@@ -0,0 +1,47 @@
+# Example config for MLP model training and testing
+
+# wandb project name
+project_name: "deepfixel"
+
+# Learning rate
+lr: 1e-3 
+
+# Batch size for training and testing
+batch_size: 512
+
+# Loss function, can be "MSE" or "L1"
+loss: "MSE"
+
+# Model type, can be "mesh_mlp" or "mesh_scnn"
+model: "mesh_mlp"
+
+# GPU for training
+gpu_id: 0
+
+# Random seed for reproducibility
+seed: 42
+
+# Keep this at 1 for correct sampling of sphere
+mesh_subdivide: 1
+
+# Keep this true for correct sampling of sphere
+healpix: true
+
+# Spread parameter for Bingham distribution of simulated data
+kappa: 100
+
+# Minimum separation angle between fibers in degrees for simulated data
+min_separation_angle: 0
+
+# Number of fibers per voxel, either 1, 2 or 'both'
+n_fibers: 'both'
+
+# Directories for testing data and testing outputs
+test_dir: "./test_data"
+output_dir: './outputs/pretrained_mlp'
+
+# Threshold for amplitude of FOD peaks (below this, peaks are ignored)
+amp_threshold: 0.1
+
+# Path to pretrained model
+pretrained_model_path: "./models/pretrained/best_model_mlp.pth"

config/example_scnn.yaml

@@ -0,0 +1,49 @@
+# Example config for spherical CNN model training and testing
+
+# wandb project name
+project_name: "deepfixel"
+
+# Learning rate
+lr: 1e-3 
+
+# Batch size for training and testing
+batch_size: 512
+
+# Loss function, can be "MSE" or "L1"
+loss: "MSE"
+
+# Model type, can be "mesh_mlp" or "mesh_scnn"
+model: "mesh_scnn"
+
+# GPU for training
+gpu_id: 0
+
+# Random seed for reproducibility
+seed: 42
+
+# Keep this at 1 for correct sampling of sphere
+mesh_subdivide: 1
+
+# Keep this true for correct sampling of sphere
+healpix: true
+
+# Spread parameter for Bingham distribution of simulated data
+kappa: 100
+
+# Minimum separation angle between fibers in degrees for simulated data
+min_separation_angle: 0
+
+# Number of fibers per voxel, either 1, 2 or 'both'
+n_fibers: 'both'
+
+### Testing parameters ###
+
+# Directories for testing data and testing outputs
+test_dir: "./test_data"
+output_dir: './outputs/pretrained_scnn'
+
+# Threshold for amplitude of FOD peaks (below this, peaks are ignored)
+amp_threshold: 0.1
+
+# Path to pretrained model
+pretrained_model_path: "./models/pretrained/best_model_scnn.pth"

run_deep_fixel.py

@@ -126,7 +126,7 @@ def generate_odf(self, seed=None):
         if self.n_fibers == "both":
             n_fibers = self.rng.choice([2, 3])
         else:
-            n_fibers = self.n_fibers
+            n_fibers = int(self.n_fibers)
 
         # Generate random volume fraction using Dirichlet distribution
         vol = self.rng.dirichlet(np.ones(n_fibers))
@@ -317,7 +317,7 @@ def generate_odf(self, seed=None):
         if self.n_fibers == "both":
             n_fibers = self.rng.choice([2, 3])
         else:
-            n_fibers = self.n_fibers
+            n_fibers = int(self.n_fibers)
 
         # Generate random volume fraction using Dirichlet distribution
         vol = self.rng.dirichlet(np.ones(n_fibers))

run_pretrained_deep_fixel.py

@@ -13,6 +13,7 @@
 
 def train_mesh_model(
     run_name,
+    project_name="deepfixel",
     lr=1e-3,
     batch_size=512,
     n_fibers=2,
@@ -56,7 +57,7 @@ def train_mesh_model(
 
     # Set up Weights and Biases
     wandb.login()
-    run = wandb.init(project="deepfixel", name=run_name, config=config)
+    run = wandb.init(project=project_name, name=run_name, config=config)
 
     # Set up datasets
     train_dataset = RandomMeshDataset(n_fibers=n_fibers, l_max=6, seed=seed, subdivide=mesh_subdivide, kappa=kappa, healpix=healpix, csd=csd, snr=snr)