Use DIPY version of patch2self (#902)

* Use DIPY version of patch2self. * Update tests. * Update test_cli.py * Switch to Forrest Gump. * Fix stuff. * Calculate MSE. * Fix noise map calculation. * Fix filenames. * Don't remove unused outputs. * Update test_cli.py * Add unit tests for Patch2Self and DWIDenoise. * Update pyproject.toml * Drop the dwidenoise test. * Update dipy.py * Try this. * Try this. * Fix. * Update dipy.py * Let's give this a try. * Fix imports. * Clean up changes.
PennLINC · Dec 11, 2024 · 716aac4 · 716aac4
1 parent fa4dbe0
commit 716aac4
Show file tree

Hide file tree

Showing 10 changed files with 280 additions and 422 deletions.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -142,6 +142,26 @@ jobs:
           paths:
             - /tmp/src/qsiprep/.circleci/data/twoses
 
+  download_forrest_gump:
+    resource_class: small
+    docker:
+      - image: cimg/python:3.10.9
+    working_directory: /tmp/src/qsiprep
+    steps:
+      - checkout
+      - restore_cache:
+          key: forrestgump-01
+      - run: *runinstall
+      - run:
+          name: Download forrest_gump test data
+          command: |
+            cd /tmp/src/qsiprep/.circleci
+            python get_data.py $PWD/data forrest_gump
+      - save_cache:
+          key: forrestgump-01
+          paths:
+            - /tmp/src/qsiprep/.circleci/data/forrest_gump
+
   DSCSDSI:
     resource_class: large
     environment:
@@ -336,6 +356,8 @@ jobs:
     <<: *dockersetup
     steps:
       - checkout
+      - restore_cache:
+          key: forrestgump-01
       - run: *runinstall
       - run:
           name: Run QSIPrep on single-shell dataset with GRE field maps
@@ -354,6 +376,64 @@ jobs:
       - store_artifacts:
           path: /tmp/src/qsiprep/.circleci/out/forrest_gump/
 
+  forrest_gump_patch2self:
+    resource_class: xlarge
+    environment:
+      CIRCLE_CPUS: 4
+    <<: *dockersetup
+    steps:
+      - checkout
+      - restore_cache:
+          key: forrestgump-01
+      - run: *runinstall
+      - run:
+          name: Run QSIPrep on single-shell dataset with GRE field maps
+          no_output_timeout: 3h
+          command: |
+            pytest -rP -o log_cli=true -m "forrest_gump_patch2self" --cov-config=/tmp/src/qsiprep/pyproject.toml --cov-append --cov-report term-missing --cov=qsiprep --data_dir=/tmp/src/qsiprep/.circleci/data --output_dir=/tmp/src/qsiprep/.circleci/out --working_dir=/tmp/src/qsiprep/.circleci/work qsiprep
+            mkdir /tmp/src/coverage
+            mv /tmp/src/qsiprep/.coverage /tmp/src/coverage/.coverage.forrest_gump_patch2self
+            # remove nifti files before uploading artifacts
+            find /tmp/src/qsiprep/.circleci/out/ -name "*.nii.gz" -type f -delete
+            find /tmp/src/qsiprep/.circleci/out/ -name "*.fib.gz" -type f -delete
+      - persist_to_workspace:
+          root: /tmp/src/coverage/
+          paths:
+            - .coverage.forrest_gump_patch2self
+      - store_artifacts:
+          path: /tmp/src/qsiprep/.circleci/out/forrest_gump_patch2self/
+
+  unit_tests:
+    resource_class: large
+    environment:
+      CIRCLE_CPUS: 4
+    <<: *dockersetup
+    steps:
+      - checkout
+      - restore_cache:
+          key: forrestgump-01
+      - run: *runinstall
+      - run:
+          name: Run unit tests
+          no_output_timeout: 1h
+          command: |
+            pytest \
+              -n ${CIRCLE_CPUS} \
+              --cov-append \
+              --cov-branch \
+              --cov-report term-missing \
+              --cov=qsiprep \
+              --data_dir=/tmp/src/qsiprep/.circleci/data \
+              --output_dir=/tmp/src/qsiprep/.circleci/out \
+              --working_dir=/tmp/src/qsiprep/.circleci/work \
+              qsiprep
+            mkdir /tmp/src/coverage
+            mv /tmp/src/qsiprep/.coverage /tmp/src/coverage/.coverage.unit_tests
+      - persist_to_workspace:
+          root: /tmp/src/coverage/
+          paths:
+            - .coverage.unit_tests
+
   IntramodalTemplate:
     resource_class: large
     <<: *dockersetup
@@ -531,6 +611,13 @@ workflows:
             tags:
               only: /.*/
 
+      - download_forrest_gump:
+          requires:
+            - build
+          filters:
+            tags:
+              only: /.*/
+
       - DSCSDSI:
           requires:
             - download_DSCSDSI
@@ -582,7 +669,21 @@ workflows:
 
       - forrest_gump:
           requires:
-            - build
+            - download_forrest_gump
+          filters:
+            tags:
+              only: /.*/
+
+      - forrest_gump_patch2self:
+          requires:
+            - download_forrest_gump
+          filters:
+            tags:
+              only: /.*/
+
+      - unit_tests:
+          requires:
+            - download_forrest_gump
           filters:
             tags:
               only: /.*/
@@ -611,6 +712,8 @@ workflows:
             - DRBUDDI_TENSORLine_EPI
             - maternal_brain_project
             - forrest_gump
+            - forrest_gump_patch2self
+            - unit_tests
             - IntramodalTemplate
             - MultiT1w
           filters:
@@ -631,6 +734,8 @@ workflows:
             - DRBUDDI_TENSORLine_EPI
             - maternal_brain_project
             - forrest_gump
+            - forrest_gump_patch2self
+            - unit_tests
             - IntramodalTemplate
             - MultiT1w
           filters:
@@ -655,4 +760,4 @@ workflows:
             branches:
               ignore: /.*/
             tags:
-              only: /.*/
+              only: /.*/
diff --git a/pyproject.toml b/pyproject.toml
@@ -78,6 +78,7 @@ tests = [
     "pytest",
     "pytest-cov",
     "pytest-env",
+    "pytest-xdist",
 ]
 maint = [
     "fuzzywuzzy",
@@ -201,6 +202,7 @@ markers = [
     "multi_t1w: test 22",
     "maternal_brain_project: multi-shell with GRE field map",
     "forrest_gump: single-shell with GRE field map",
+    "forrest_gump_patch2self: single-shell with GRE field map and patch2self",
 ]
 env = [
     "RUNNING_PYTEST = 1",

diff --git a/qsiprep/interfaces/denoise.py b/qsiprep/interfaces/denoise.py
@@ -7,6 +7,8 @@
 
 """
 
+import os
+
 import nibabel as nb
 import numpy as np
 import pandas as pd
@@ -45,6 +47,9 @@ def _calculate_nmse(self, original_nii, corrected_nii):
         """Calculate NMSE from the applied preprocessing operation."""
         outputs = self._list_outputs()
         output_file = outputs.get('nmse_text')
+        if not output_file:
+            output_file = os.path.abspath('nmse_text.txt')
+
         pres = []
         posts = []
         differences = []
@@ -62,6 +67,7 @@ def _calculate_nmse(self, original_nii, corrected_nii):
         pd.DataFrame(
             {title + '_pre': pres, title + '_post': posts, title + '_change': differences}
         ).to_csv(output_file, index=False)
+        self._nmse_text = output_file
 
     def _generate_report(self):
         """Generate a reportlet."""

diff --git a/qsiprep/interfaces/dipy.py b/qsiprep/interfaces/dipy.py
@@ -22,7 +22,6 @@
     SeriesPreprocReportInputSpec,
     SeriesPreprocReportOutputSpec,
 )
-from .patch2self import patch2self
 
 LOGGER = logging.getLogger('nipype.interface')
 TAU_DEFAULT = 1.0 / (4 * np.pi**2)
@@ -37,7 +36,6 @@ def popen_run(arg_list):
 
 class Patch2SelfInputSpec(SeriesPreprocReportInputSpec):
     in_file = File(exists=True, mandatory=True, desc='4D diffusion MRI data file')
-    patch_radius = traits.Either(traits.Int(0), traits.Str('auto'), desc='patch radius in voxels.')
     bval_file = File(exists=True, mandatory=True, desc='bval file containing b-values')
     model = traits.Str('ols', usedefault=True, desc='Regression model for Patch2Self')
     alpha = traits.Float(1.0, usedefault=True, desc='Regularization parameter for Ridge and Lasso')
@@ -67,6 +65,8 @@ class Patch2Self(SeriesPreprocReport, SimpleInterface):
     output_spec = Patch2SelfOutputSpec
 
     def _run_interface(self, runtime):
+        from dipy.denoise.patch2self import patch2self
+
         in_file = self.inputs.in_file
         bval_file = self.inputs.bval_file
         denoised_file = fname_presuffix(
@@ -79,47 +79,19 @@ def _run_interface(self, runtime):
         noisy_arr = noisy_img.get_fdata()
         bvals = np.loadtxt(bval_file)
 
-        # Determine the patch radius
-        num_non_b0 = (bvals > self.inputs.b0_threshold).sum()
-        very_few_directions = num_non_b0 < 20
-        few_directions = num_non_b0 < 50
-        if self.inputs.patch_radius == 'auto':
-            if very_few_directions:
-                patch_radius = [3, 3, 3]
-            elif few_directions:
-                patch_radius = [1, 1, 1]
-            else:
-                patch_radius = [0, 0, 0]
-        else:
-            patch_radius = [self.inputs.patch_radius] * 3
-            if self.inputs.patch_radius > 3 and not very_few_directions:
-                LOGGER.info(
-                    'a very large patch radius is not necessary when more than '
-                    '20 gradient directions have been sampled.'
-                )
-            elif self.inputs.patch_radius > 1 and not few_directions:
-                LOGGER.info(
-                    'a large patch radius is not necessary when more than '
-                    '50 gradient directions have been sampled.'
-                )
-            elif self.inputs.patch_radius == 0 and few_directions:
-                LOGGER.warning(
-                    'When < 50 gradient directions are available, it is '
-                    'recommended to increase patch_radius to > 0.'
-                )
-
-        denoised_arr, noise_residuals = patch2self(
-            noisy_arr,
-            bvals,
+        denoised_arr = patch2self(
+            data=noisy_arr,
+            bvals=bvals,
             model=self.inputs.model,
             alpha=self.inputs.alpha,
-            patch_radius=patch_radius,
             b0_threshold=self.inputs.b0_threshold,
             verbose=True,
             b0_denoising=self.inputs.b0_denoising,
             clip_negative_vals=self.inputs.clip_negative_vals,
             shift_intensity=self.inputs.shift_intensity,
         )
+        # Calculate a "noise level" image
+        noise_residuals = np.sqrt(np.mean((noisy_arr - denoised_arr) ** 2, axis=3))
 
         # Back to nifti
         denoised_img = nb.Nifti1Image(denoised_arr, noisy_img.affine, noisy_img.header)
@@ -128,6 +100,7 @@ def _run_interface(self, runtime):
         p2s_residuals.to_filename(noise_file)
         self._results['out_file'] = denoised_file
         self._results['noise_image'] = noise_file
+        self._nmse_text = None
         return runtime
 
     def _get_plotting_images(self):
@@ -138,3 +111,7 @@ def _get_plotting_images(self):
         noise_name = outputs['noise_image']
         noisenii = load_img(noise_name)
         return input_dwi, denoised_nii, noisenii
+
+    def _list_outputs(self):
+        self._results['nmse_text'] = self._nmse_text
+        return super()._list_outputs()