Update Dockerfile for inference

Author: saundersresearch

Dockerfile

@@ -52,4 +52,8 @@ PYPROJECT
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install /opt/hsd
 
+RUN mkdir -p /app/models
+ADD https://zenodo.org/records/17859792/files/best_model_mlp.pth?download=1 /app/models/best_model_mlp.pth
+ADD https://zenodo.org/records/17859792/files/best_model_scnn.pth?download=1 /app/models/best_model_scnn.pth
+
 ENV PATH="/app/.venv/bin:${PATH}"

README.md

@@ -1,5 +1,5 @@
 # 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) [![Zenodo](https://zenodo.org/badge/DOI/10.5281/zenodo.13121149.svg)](https://doi.org/10.5281/zenodo.13121149)
+[![arXiv](https://img.shields.io/badge/arXiv-2511.03893-b31b1b.svg)](https://arxiv.org/abs/2511.03893) [![Zenodo](https://zenodo.org/badge/DOI/10.5281/zenodo.17834289.svg)](https://doi.org/10.5281/zenodo.17834289)
 
 DeepFixel is a deep learning method for identification of crossing fiber bundle elements from diffusion MRI.
 
@@ -14,11 +14,28 @@ uv sync
 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://doi.org/10.5281/zenodo.13121149](https://doi.org/10.5281/zenodo.13121149). 
+You can use DeepFixel to split multi-fiber ODFs into the underlying single-fiber ODFs. The pretrained weights are available on Zenodo: [https://doi.org/10.5281/zenodo.17834289](https://doi.org/10.5281/zenodo.17834289). 
+
+```bash
+# For pretrained models: --lmax 6, --subdivide 1
+deepfixel /path/to/input/fod.nii.gz \
+    /path/to/output_dir \
+    /path/to/best_model_scnn.pth \
+    --mask /path/to/mask.nii.gz \
+    --maxnum 2 \
+    --lmax 6 \
+    --subdivide 1 \
+    --amp_threshold 0.1 \
+    --model mesh_scnn \
+    --batch_size 512 \
+    --gpu_id 1
+```
+
+## Training and testing the model
+To run the model, download the weights and testing dataset from the following link: [https://doi.org/10.5281/zenodo.17834289](https://doi.org/10.5281/zenodo.17834289). 
 - Unzip and copy the testing data to `./test_data`
 - Put the weights in `./models/pretrained`
 
-
 To train the model:
 ```bash
 python train_deep_fixel.py --config config/example_scnn.yaml
@@ -30,28 +47,58 @@ python test_deep_fixel.py --config config/example_scnn.yaml
 ```
 
 
-## Usage (Docker)
+### 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 .
+sudo docker build -t spherical_deep_fixel:v1.2.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
+sudo docker run --rm -it --gpus all \
+    /path/to/input/fod.nii.gz \
+    /path/to/output_dir \
+    /app/models/best_model_scnn.pth \
+    --mask /path/to/mask.nii.gz \
+    --maxnum 2 \
+    --lmax 6 \
+    --subdivide 1 \
+    --amp_threshold 0.1 \
+    --model mesh_scnn \
+    --batch_size 512 \
+    --gpu_id 0
 ```
 
-## Usage (Apptainer)
-A pre-built Apptainer image is available on Zenodo ([https://doi.org/10.5281/zenodo.13121149](https://doi.org/10.5281/zenodo.13121149)):
+For training and testing:
+```bash
+sudo docker run --rm -it --gpus all $spherical_deep_fixel:v1.2.0 python train_deep_fixel.py --config /path/to/config/example_scnn.yaml
+sudo docker run --rm -it --gpus all $spherical_deep_fixel:v1.2.0 python test_deep_fixel.py --config /path/to/config/example_scnn.yaml
+```
+
+### Apptainer
+A pre-built Apptainer image is available on Zenodo ([https://doi.org/10.5281/zenodo.17834289](https://doi.org/10.5281/zenodo.17834289)). (Note you will likely need to bind in local directories with `-B`):
 
 ```bash
-apptainer run -C -B $(pwd):$(pwd) --nv https://zenodo.org/records/17834290/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/17834290/files/spherical_deep_fixel_v1.0.0.sif python /app/test_deep_fixel.py --config /path/to/config/example_scnn.yaml
+apptainer run -C --nv spherical_deep_fixel_v1.2.0.sif \
+    /path/to/input/fod.nii.gz \
+    /path/to/output_dir \
+    /app/models/best_model_scnn.pth \
+    --mask /path/to/mask.nii.gz \
+    --maxnum 2 \
+    --lmax 6 \
+    --subdivide 1 \
+    --amp_threshold 0.1 \
+    --model mesh_scnn \
+    --batch_size 512 \
+    --gpu_id 0
 ```
 
-## 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.
+For training and testing:
+```bash
+apptainer run -C --nv spherical_deep_fixel_v1.2.0.sif python /app/train_deep_fixel.py --config /path/to/config/example_scnn.yaml
+apptainer run -C --nv spherical_deep_fixel_v1.2.0.sif python /app/test_deep_fixel.py --config /path/to/config/example_scnn.yaml
+```
 
 ## Citation
 If you use this code in your research, please cite the following paper:

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "deep_fixel"
-version = "1.0.0"
+version = "1.2.0"
 authors = [
     {name = "Adam Saunders", email = "adam.m.saunders@vanderbilt.edu"}
 ]
@@ -41,5 +41,8 @@ dependencies = [
     "wandb>=0.19.7",
 ]
 
+[project.scripts]
+deepfixel = "deep_fixel.scripts.inference:main"
+
 [project.urls]
-Homepage = "https://github.com/MASILab/deep_fixel"
+Homepage = "https://github.com/MASILab/spherical_deep_fixel"

src/deep_fixel/dataset.py

@@ -26,6 +26,7 @@
 
 from .utils import load_fissile_mat, fiber_response
 
+
 class RandomODFDataset(IterableDataset):
     def __init__(self, n_fibers, l_max=6, seed=None, size=None, deterministic=False):
         """Generate ODFs at random angles and volume fractions (using Tournier07/mrtrix convention)
@@ -182,7 +183,8 @@ def __iter__(self):
         else:
             while True:
                 yield self.generate_odf()
-                
+
+
 class RandomMeshDataset(IterableDataset):
     def __init__(
         self,
@@ -382,7 +384,7 @@ def generate_odf(self, seed=None):
             sh_order_max=self.l_max,
             basis_type="tournier07",
         )
-        
+
         pdf = [
             v * vonmises_fisher(mu, self.kappa).pdf(self.icosphere.vertices)
             for v, mu in zip(vol, xyz.T)
@@ -614,37 +616,36 @@ def __getitem__(self, idx):
 class GeneratedMeshNIFTIDataset(Dataset):
     def __init__(
         self,
-        n_fibers,
         nifti_path,
+        mask=None,
+        lmax=6,
         subdivide=3,
-        kappa=100,
         healpix=False,
     ):
         """Load ODFs from a directory of .mat files from FISSILE outputs.
 
         Parameters
         ----------
-        n_fibers : int or str
-            Number of fibers in each ODF. If 'both', will use 2 and 3.
         nifti_path : str
-            Path to nifti file
+            Path to nifti file containing ODFs in spherical harmonic coefficients in tournier07/mrtrix convention
+        mask : str
+            Path to nifti file containing mask, by default None
+        lmax : int, optional
+            Maximum spherical harmonic order, by default 6
         subdivide : int, optional
             Number of times to subdivide the ico-hemisphere if healpix=False,
             otherwise corresponds to depth of Healpix sampling (where smaller is more vertices), by default 3
-        kappa : float, optional
-            Concentration parameter for von Mises-Fisher distribution, by default 100
         healpix : bool, optional
             If True, sample on healpix instead of icosphere, by default False
         """
-        self.n_fibers = n_fibers
         self.nifti_path = nifti_path
-        self.l_max = 6
+        self.l_max = lmax
 
         if healpix:
             n_side = 8
             depth = subdivide
             patch_size = 1
-            sh_degree = 6
+            sh_degree = lmax
             pooling_mode = "average"
             pooling_name = "mixed"
             use_hemisphere = True
@@ -666,13 +667,19 @@ def __init__(
             self.n_mesh = len(self.icosphere.vertices)
             self.sphere = self.icosphere
 
-        self.kappa = kappa
-
         # Load NIFTI
         nifti = nib.load(nifti_path)
+        self.shape = nifti.shape
+        self.affine = nifti.affine
 
         # Flatten all except last axis
         nifti_data = nifti.get_fdata().squeeze()
+
+        if mask is not None:
+            mask_nifti = nib.load(mask)
+            mask_data = mask_nifti.get_fdata().squeeze().astype(bool)
+            nifti_data = nifti_data[mask_data]
+
         nifti_data = nifti_data.reshape(-1, nifti_data.shape[-1])
 
         # Append each ODF in file to list
@@ -689,4 +696,4 @@ def __len__(self):
         return len(self.total_odf_meshes)
 
     def __getitem__(self, idx):
-        return self.total_odf_meshes[idx]
+        return self.total_odf_meshes[idx]

src/deep_fixel/utils.py

@@ -20,10 +20,13 @@
 from scipy.signal import argrelmax
 from scipy.spatial.transform import Rotation as R
 from scipy.io import loadmat
-import cmcrameri 
+import cmcrameri
 from line_profiler import profile
 
-def plot_odf(odf, ax=None, color="blue", basis="tournier", alpha=1, linewidth=0.1, sphere=None):
+
+def plot_odf(
+    odf, ax=None, color="blue", basis="tournier", alpha=1, linewidth=0.1, sphere=None
+):
     """Plot a spherical orientation distribution function represented by spherical harmonic coefficients.
 
     Parameters
@@ -413,7 +416,10 @@ def match_odfs(true_odfs, est_odfs):
 
     return matched_est_odfs, index_array
 
-def pdf2odfs(mesh, sphere, amp_threshold=0.5, use_dipy=False, **kwargs):
+
+def pdf2odfs(
+    mesh, sphere, amp_threshold=0.5, use_dipy=False, lmax=6, max_num=None, **kwargs
+):
     """
     Estimate ODFs from spherical PDF
 
@@ -426,7 +432,11 @@ def pdf2odfs(mesh, sphere, amp_threshold=0.5, use_dipy=False, **kwargs):
     amp_threshold : float, optional
         Amplitude threshold for maxima, by default 0.5
     use_dipy : bool, optional
-        Whether to use Dipy's peak finding, by default True.
+        Whether to use Dipy's peak finding, by default False.
+    lmax : int, optional
+        Maximum spherical harmonic degree, by default 6
+    max_num : int, optional
+        Maximum number of peaks to find, by default None
     **kwargs : dict, optional
         Additional keyword arguments for peak finding if using Dipy.
 
@@ -439,6 +449,8 @@ def pdf2odfs(mesh, sphere, amp_threshold=0.5, use_dipy=False, **kwargs):
     vol_fracs : NumPy array
         Estimated volume fractions
     """
+    m_list, l_list = sph_harm_ind_list(lmax)
+
     if use_dipy == "nl":
         points = np.array([sphere.theta, sphere.phi]).T
         pdf_mesh_interp = CloughTocher2DInterpolator(points, mesh, fill_value=0)
@@ -449,13 +461,19 @@ def pdf2odfs(mesh, sphere, amp_threshold=0.5, use_dipy=False, **kwargs):
             pdf_mesh_interp_sphere,
             relative_peak_threshold=amp_threshold,
             sphere=sphere,
-            **kwargs
+            **kwargs,
         )
+        if max_num is not None:
+            # Keep only top max_num peaks
+            if vals.shape[0] > max_num:
+                top_indices = np.argsort(vals)[-max_num:]
+                xyz = xyz[top_indices]
+                vals = vals[top_indices]
+
         vol_fracs = vals / np.sum(vals)  # Normalize volume fractions
         r, theta, phi = cart2sphere(xyz[:, 0], xyz[:, 1], xyz[:, 2])
         dirs = np.array([theta, phi]).T
         odfs = []
-        m_list, l_list = sph_harm_ind_list(6)
         for theta, phi, vol_frac in zip(dirs[:, 0], dirs[:, 1], vol_fracs):
             odfs.append(
                 convert_sh_descoteaux_tournier(
@@ -473,11 +491,21 @@ def pdf2odfs(mesh, sphere, amp_threshold=0.5, use_dipy=False, **kwargs):
             **kwargs,
         )
 
+        if len(vals) == 0:
+            return np.zeros((1, len(m_list))), np.zeros((1, 2)), np.zeros((1,))
+
+        if max_num is not None:
+            # Keep only top max_num peaks
+            if vals.shape[0] > max_num:
+                top_indices = np.argsort(vals)[-max_num:]
+                xyz = xyz[top_indices]
+                vals = vals[top_indices]
+
         vol_fracs = vals / np.sum(vals)  # Normalize volume fractions
         r, theta, phi = cart2sphere(xyz[:, 0], xyz[:, 1], xyz[:, 2])
         dirs = np.array([theta, phi]).T
         odfs = []
-        m_list, l_list = sph_harm_ind_list(6)
+        m_list, l_list = sph_harm_ind_list(lmax)
         for theta, phi, vol_frac in zip(dirs[:, 0], dirs[:, 1], vol_fracs):
             odfs.append(
                 convert_sh_descoteaux_tournier(
@@ -527,11 +555,18 @@ def pdf2odfs(mesh, sphere, amp_threshold=0.5, use_dipy=False, **kwargs):
         minima = minima[minima_vals > amp_threshold]
         minima_vals = minima_vals[minima_vals > amp_threshold]
 
+        if max_num is not None:
+            # Keep only top max_num peaks
+            if minima_vals.shape[0] > max_num:
+                top_indices = np.argsort(minima_vals)[-max_num:]
+                minima = minima[top_indices]
+                minima_vals = minima_vals[top_indices]
+
         # Estimate volume fraction using ratio of amplitude at minima
         minima_vals = minima_vals / np.sum(minima_vals)
 
         # Now estimate ODF at these points with these volume fractions
-        m_list, l_list = sph_harm_ind_list(6)
+        m_list, l_list = sph_harm_ind_list(lmax)
         odfs = []
         for min, min_val in zip(minima, minima_vals):
             odfs.append(
@@ -570,6 +605,7 @@ def angular_separation(angle1, angle2):
     delta = np.arccos(cos_delta)
     return delta
 
+
 def load_fissile_mat(path):
     """Load output from FISSILE and match estimated fibers to true fibers.
 
@@ -671,6 +707,7 @@ def load_fissile_mat(path):
 
     return data_dict
 
+
 def fiber_response(
     sphere,
     theta=0,
@@ -707,4 +744,4 @@ def fiber_response(
     response_amp = np.exp(-bval * lambda_perp) * np.exp(
         -3 * bval * (lambda_mean - lambda_perp) * (cos_theta**2)
     )
-    return response_amp
+    return response_amp