enh+ref: allow generation of multiple median files

Author: mgxd

nipype/algorithms/misc.py

@@ -1381,13 +1381,15 @@ def merge_rois(in_files, in_idxs, in_ref,
 
 
 class CalculateMedianInputSpec(BaseInterfaceInputSpec):
-    in_file = InputMultiPath(File(exists=True, mandatory=True,
+    in_files = InputMultiPath(File(exists=True, mandatory=True,
       desc="One or more realigned Nifti 4D timeseries"))
-    median_file = traits.Str('median.nii.gz', usedefault=True,
-      desc="Filename to store median image")
+    median_file = traits.Str(desc="Filename prefix to store median images")
+    median_per_file = traits.Bool(False, usedefault=True,
+      desc="Calculate a median file for each Nifti")
 
 class CalculateMedianOutputSpec(TraitedSpec):
-    median_file = File(exists=True)
+    median_files = OutputMultiPath(File(exists=True),
+      desc="One or more median images")
 
 class CalculateMedian(BaseInterface):
     """
@@ -1398,32 +1400,73 @@ class CalculateMedian(BaseInterface):
 
     >>> from nipype.algorithms.misc import CalculateMedian
     >>> mean = CalculateMedian()
-    >>> mean.inputs.in_file = 'functional.nii'
+    >>> mean.inputs.in_files = 'functional.nii'
     >>> mean.run() # doctest: +SKIP
 
     """
     input_spec = CalculateMedianInputSpec
     output_spec = CalculateMedianOutputSpec
 
+    def __init__(self, *args, **kwargs):
+        super(CalculateMedian, self).__init__(*args, **kwargs)
+        self._median_files = []
+
+    def _gen_fname(self, suffix, idx=None, ext=None):
+        if idx:
+            in_file = self.inputs.in_files[idx]
+        else:
+            if isinstance(self.inputs.in_files, list):
+                in_file = self.inputs.in_files[0]
+            else:
+                in_file = self.inputs.in_files
+        fname, in_ext = op.splitext(op.basename(in_file))
+        if in_ext == '.gz':
+            fname, in_ext2 = op.splitext(fname)
+            in_ext = in_ext2 + in_ext
+        if ext is None:
+            ext = in_ext
+        if ext.startswith('.'):
+            ext = ext[1:]
+        if self.inputs.median_file:
+            outname = self.inputs.median_file
+        else:
+            outname = '{}_{}'.format(fname, suffix)
+        if idx:
+            outname += str(idx)
+        return op.abspath('{}.{}'.format(outname, ext))
+
     def _run_interface(self, runtime):
         total = None
-        for idx, fname in enumerate(filename_to_list(self.inputs.in_file)):
+        self._median_files = []
+        for idx, fname in enumerate(filename_to_list(self.inputs.in_files)):
             img = nb.load(fname, mmap=NUMPY_MMAP)
             data = np.median(img.get_data(), axis=3)
-            if total is None:
-                total = data
+            if self.inputs.median_per_file:
+                self._median_files.append(self._write_nifti(img, data, idx))
             else:
-                total += data
-        median_img = nb.Nifti1Image(total/(idx + 1), img.affine, img.header)
-        filename = os.path.join(os.getcwd(), self.inputs.median_file)
-        median_img.to_filename(filename)
+                if total is None:
+                    total = data
+                else:
+                    total += data
+        if not self.inputs.median_per_file:
+            self._median_files.append(self._write_nifti(img, total, idx))
         return runtime
 
     def _list_outputs(self):
         outputs = self._outputs().get()
-        outputs['median_file'] = os.path.abspath(self.inputs.median_file)
+        outputs['median_files'] = self._median_files
         return outputs
 
+    def _write_nifti(self, img, data, idx, suffix='median'):
+        if self.inputs.median_per_file:
+            median_img = nb.Nifti1Image(data, img.affine, img.header)
+            filename = self._gen_fname(suffix, idx=idx)
+        else:
+            median_img = nb.Nifti1Image(data/(idx+1), img.affine, img.header)
+            filename = self._gen_fname(suffix)
+        median_img.to_filename(filename)
+        return filename
+
 
 # Deprecated interfaces ------------------------------------------------------
 

nipype/algorithms/tests/test_misc.py

@@ -10,6 +10,7 @@
 from nipype.utils.filemanip import fname_presuffix
 from nipype.testing.fixtures import create_analyze_pair_file_in_directory
 from nipype.utils import NUMPY_MMAP
+from nipype.testing import example_data
 
 
 def test_CreateNifti(create_analyze_pair_file_in_directory):
@@ -36,14 +37,12 @@ def test_CreateNifti(create_analyze_pair_file_in_directory):
 
 def test_CalculateMedian(create_analyze_pair_file_in_directory):
 
-    filelist, outdir = create_analyze_pair_file_in_directory
-
     mean = misc.CalculateMedian()
 
     with pytest.raises(TypeError): mean.run()
 
-    mean.inputs.in_file = filelist[0]
+    mean.inputs.in_files = example_data('ds003_sub-01_mc.nii.gz')
     eg = mean.run()
 
-    assert os.path.exists(eg.outputs.median_file)
-    assert nb.load(eg.outputs.median_file, mmap=NUMPY_MMAP)
+    assert os.path.exists(eg.outputs.median_files)
+    assert nb.load(eg.outputs.median_files, mmap=NUMPY_MMAP)